common#

Backend infrastructure: numpy/jax dispatch, array predicates, scans/maps, shared mixins.

get_backend(is_uniform, use_jax) -> (xnp, xmap, xscan) selects the array module and loop machinery; is_ndarray/is_jax_ndarray/tree_contains_jax are the type-inference predicates behind the no-use_jax-parameter convention (dispatch is inferred from the input arrays at the lowest level); ValueHashedMasks is the value-based hash/eq mixin that keeps a rebuilt-but-identical mask holder on the same jit cache key.

Attributes#

`jax_available`
`NDArray`
`is_ndarray`
`is_jax_ndarray`
`is_numpy_ndarray`
`to_jax`
`to_numpy`
`jax_scan`
`jax_map`

Classes#

ValueHashedMasks

Mixin giving uniform-layer mask holders VALUE-based __hash__/__eq__ over mask content.

Functions#

`is_boolean_ndarray`(x)
`ragged_scan`(f, init, xs)	Similar to jax.lax.scan, except for ragged-sized arrays
`numpy_scan`(f, init, xs)	Similar to jax.lax.scan, except returns numpy arrays instead of jax arrays.
`ragged_map`(f, xs)
`numpy_map`(f, xs)
`get_backend`(is_uniform, use_jax)
`xwhile`(cond, body, init_state[, use_jit])	Data-dependent `while` with the numpy / eager-jax / jit dispatch -- the `xscan` precedent for a
`xcat`(x, y)	Concatenate arrays or sequences.
`xappend`(S, x)	Append slice to array or element to sequence
`xprepend`(x, S)	Prepend slice to array or element to sequence
`randn`(*args, use_jax)
`tree_contains_jax`(T)
`tree_to_jax`(T)	Move every array leaf of a pytree (nested tuples/lists of arrays) onto jax, preserving the tree
`items_are_uniform`(xx)	Checks if an object can be treated as uniform for the purposes of jax.scan and jax.map.
`save_core_families`(file, families)	Save a sequence of core-families (each a sequence of arrays) to a `.npz` file.
`load_core_families`(file)	Inverse of `save_core_families()`: load a `.npz` file into a tuple of core-families.
`require_concrete_masks`(*masks)	Guard the uniform-mask contract: masks are concrete host (numpy) arrays, never jax tracers.
`prefix_mask`(ranks, pad)	Boolean prefix indicator: slot `j` is real iff `j < rank` -- the canonical (prefix) form.

Module Contents#

jax_available = False#

NDArray#

is_ndarray#

is_jax_ndarray#

is_numpy_ndarray#

to_jax#

to_numpy#

jax_scan#

jax_map#