traveler.tables.JaxTable#

class traveler.tables.JaxTable(*args, **kwargs)[source]#

Bases: JaxDataset

Analog of a pyarrow Table, but with JAX arrays.

Every column is a JAX array, and each should be a 1D array of common length, or a 0D array (scalar). The table can be indexed or sliced by row. Column ordering is loosely preserved, but not guaranteed.

Allocate a table with an empty row-count cache.

Parameters:

  • *args (Any) – Positional initialization arguments.

  • **kwargs (Any) – Keyword initialization arguments.

Returns:

Newly allocated instance of cls.

Return type:

JaxTable

Methods#

assign([vars])

Assign new variables to the JaxDataset, returning a new instance.

copy()

Create a copy of the JaxTable, without any validation.

from_pandas(df[, id_col])

Create a JaxTable from a pandas DataFrame.

from_polars(df[, id_col])

Create a JaxTable from a polars DataFrame.

info()

Print a compact summary of fields, shapes, and data types.

schema()

Return a schema of the JaxDataset, describing its fields and types.

to_dict()

Convert the JaxDataset to a dictionary of JAX arrays.

to_pandas([include_prng_key])

Convert the JaxTable to a pandas DataFrame.

to_polars([include_prng_key])

Convert the table to a Polars data frame.

validate([sizes, verbose, sources, ...])

Validate dataset fields against the class annotations.

with_prng_keys(id_col, *, seed, namespace[, ...])

Return a table with deterministic per-row PRNG keys materialized.

Attributes#

shape

Return the one-dimensional row shape of the table.