traveler.tables.polars_to_jax#

traveler.tables.polars_to_jax(df, cls=JaxTable, prng_keys=('_prng_key_',))[source]#

Convert a Polars data frame to a JAX-backed table.

Parameters:

  • df (polars.DataFrame) – Data frame to convert.

  • cls (type[JaxTable], optional) – Table class to instantiate.

  • prng_keys (collections.abc.Collection[str], optional) – Column names containing serialized JAX PRNG keys.

Returns:

Table containing JAX arrays and reconstructed categorical columns.

Return type:

JaxTable