traveler.JaxTable

class traveler.JaxTable(*args, **kwargs)

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.
assign_in_place([vars])	Assign new variables to the JaxDataset, in place.
copy()	Create a copy of the JaxTable, without any validation.
from_pandas(df[, id_col])	Create a JaxTable from a pandas DataFrame.
from_parquet(filename[, id_col])	Create a dataset from a Parquet file.
from_polars(df[, id_col])	Create a JaxTable from a polars DataFrame.
info()	Print a compact summary of fields, shapes, and data types.
read_csv(filename)	Read a CSV file into a JaxTable.
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.

Attributes

global_random_seed
shape	Return the one-dimensional row shape of the table.