Multi-Dim Datasets

Traveler uses JaxDataset for named collections of JAX arrays that do not naturally form a table. Each variable has a name, but the variables may have different numbers of dimensions and different shapes. This makes a JaxDataset useful for multidimensional model inputs such as skims, other related level-of-service (LOS) matrices, and accessibility arrays.

Like a JaxTable, a JaxDataset:

• stores its numerical data as JAX arrays;

• provides variables through attributes and string keys;

• carries a separate metadata dictionary;

• is a JAX pytree, so its arrays can participate in JAX transformations; and

• supports functional updates with assign().

The relationship is one of specialization: JaxTable is a subclass of JaxDataset that adds table rules.

	JaxDataset	JaxTable
Variable shapes	May differ	Must be scalar or one-dimensional
Common leading length	Not required	Required for non-scalar columns
Typical meaning	Named model arrays	Rows and columns
Typical use	Skims and other multidimensional data	Households, persons, tours, or land use

A dataset is therefore not an xarray-style labeled array and does not perform alignment by coordinate name. Its variables are ordinary JAX arrays indexed by position. The variable names, class annotations, categorical coordinates, and metadata provide the model-specific meaning around those arrays.

Multidimensional LOS data

A skim describes the impedance between an origin and a destination. Some LOS measures have two dimensions:

distance[origin, destination]

Others vary by time period and need a third dimension:

travel_time[origin, destination, time_period]

These variables cannot be columns of one JaxTable: they are not one-dimensional and do not represent one record per row. They can coexist in one JaxDataset because the dataset does not impose a shared rank or shape.

import jax.numpy as jnp

from traveler.trees import JaxDataset


distance = jnp.array(
    [
        [0, 4, 7],
        [4, 0, 3],
        [7, 3, 0],
    ],
    dtype=jnp.float32,
)

am_time = jnp.array(
    [
        [0, 8, 13],
        [8, 0, 7],
        [13, 7, 0],
    ],
    dtype=jnp.float32,
)

pm_time = jnp.array(
    [
        [0, 11, 18],
        [11, 0, 9],
        [18, 9, 0],
    ],
    dtype=jnp.float32,
)

skims = JaxDataset(
    {
        "distance": distance,
        "travel_time": jnp.stack([am_time, pm_time], axis=-1),
    },
    name="skims",
    metadata={"time_periods": ["AM", "PM"]},
)

The two variables retain their independent shapes:

skims.distance.shape  # (3, 3)
skims.travel_time.shape  # (3, 3, 2)

skims["distance"] is skims.distance  # True
skims.travel_time[:, :, 0]  # the complete AM matrix

For model records with coded origin, destination, and period indices, JAX advanced indexing selects one LOS value per record:

origins = jnp.array([0, 2], dtype=jnp.int32)
destinations = jnp.array([2, 0], dtype=jnp.int32)
periods = jnp.array([1, 0], dtype=jnp.int32)

skims.travel_time[origins, destinations, periods]
# Array([18., 13.], dtype=float32)

The same dataset can be placed in a Store and accessed as store["skims"]. Model code can then use the appropriate two- or three-dimensional variable without copying it into a row-oriented table.

Describing the data model with TypeMaker

A plain JaxDataset accepts any mapping of names to array-like values. For an application-specific dataset, a subclass can declare which variables are expected and describe their shapes and dtypes with TypeMaker.

TypeMaker binds a symbolic shape expression to JAX array types. Accessing a dtype attribute produces a jaxtyping annotation with that shape and dtype; it does not allocate an array.

from traveler import TypeMaker


taz = TypeMaker("TAZ")
od = TypeMaker("TAZ TAZ")
odt = TypeMaker("TAZ TAZ TIMEPERIOD")

These three TypeMaker instances describe the shapes of the skims arrays, and can be used to annotate the variables of a JaxDataset subclass. Each annotation also should add a dtype to specify the expected data type of the array, using attributes on the TypeMaker instance. For example:

• taz.Int32 means a one-dimensional int32 JAX array indexed by TAZ;

• od.Float32 means a two-dimensional float32 JAX array whose two TAZ axes have the same size; and

• odt.Float32 means a three-dimensional float32 JAX array with origin TAZ, destination TAZ, and time-period axes.

Reusing a dimension name communicates that its size should be consistent wherever it appears. In particular, "TAZ TAZ" describes a square origin-destination matrix. Different names such as "ORIGIN DESTINATION" would describe two axes without asserting that they have the same size.

The annotations form the data model for a skim-specific subclass:

from traveler.trees import JaxDataset


class Skims(JaxDataset):
    zone_id: taz.Int32
    distance: od.Float32
    travel_time: odt.Float32

We can now construct the typed dataset with the same arrays:

skims = Skims(
    {
        "zone_id": jnp.array([101, 102, 103], dtype=jnp.int32),
        "distance": distance,
        "travel_time": jnp.stack([am_time, pm_time], axis=-1),
    },
    name="skims",
    metadata={"time_periods": ["AM", "PM"]},
)

skims.distance.shape  # (3, 3)
skims.travel_time.dtype  # dtype('float32')

This pattern puts domain vocabulary in one place. A loader is responsible for reading source data and constructing the arrays, while the Skims class says which variables downstream model code expects and how their axes and dtypes are intended to be organized. TypeMaker also provides attributes for other common JAX dtypes, including Bool, Int8, Int16, Int64, and Float64.

skims.validate() after loading data to check required variables, dtypes, shapes, and the consistency of named dimensions. Validation safely coerces a dtype by default only when no values would be lost; pass coerce=False to require exact dtypes without conversion. Use skims.schema() to inspect the shapes and dtypes of the arrays actually present in an instance. :::

When zone or time-period labels need to travel with the arrays, they can be stored as categorical variables or metadata. Numerical model calculations still index the LOS arrays with integer codes, keeping the data directly usable by JAX.