traveler.tables.JaxDataset#

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

Bases: AttrPyTree

A collection of JAX arrays of with metadata.

The arrays are not required to have the same shape or dimensions.

Allocate a pytree with empty data and metadata mappings.

Parameters:

  • *args (Any) – Positional initialization arguments accepted by the subclass.

  • **kwargs (Any) – Keyword initialization arguments accepted by the subclass.

  • data (dict[str, Any] | Dataset | DataFrame)

  • id_col (str | None)

  • name (str | None)

  • metadata (dict[str, Any] | None)

Returns:

Newly allocated, empty instance of cls.

Return type:

AttrPyTree

Methods#

assign([vars])

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

copy()

Create a copy of the JaxTable, without any validation.

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.

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

Validate dataset fields against the class annotations.