jaxsnn.base.dataset.dataloader

Classes

Any(*args, **kwargs)

Special type indicating an unconstrained type.

Functions

jaxsnn.base.dataset.dataloader.data_loader(dataset: Tuple[Any, Any], batch_size: int, num_batches: Optional[int, None] = None, rng: Optional[jax.Array, None] = None)

jaxsnn.base.dataset.dataloader.tree_leaves(tree: Any, is_leaf: Optional[Callable[[Any], bool], None] = None) → List[Any]

Gets the leaves of a pytree.

jaxsnn.base.dataset.dataloader.tree_map(f: Callable[[…], Any], tree: Any, *rest: Any, is_leaf: Optional[Callable[[Any], bool], None] = None) → Any

Maps a multi-input function over pytree args to produce a new pytree.

Args:

f: function that takes 1 + len(rest) arguments, to be applied at the

corresponding leaves of the pytrees.

tree: a pytree to be mapped over, with each leaf providing the first

positional argument to f.

rest: a tuple of pytrees, each of which has the same structure as tree

or has tree as a prefix.

is_leaf: an optionally specified function that will be called at each

flattening step. It should return a boolean, which indicates whether the flattening should traverse the current object, or if it should be stopped immediately, with the whole subtree being treated as a leaf.

Returns:

A new pytree with the same structure as tree but with the value at each leaf given by f(x, *xs) where x is the value at the corresponding leaf in tree and xs is the tuple of values at corresponding nodes in rest.

Examples:

>>> import jax.tree_util
>>> jax.tree_util.tree_map(lambda x: x + 1, {"x": 7, "y": 42})
{'x': 8, 'y': 43}

If multiple inputs are passed, the structure of the tree is taken from the first input; subsequent inputs need only have tree as a prefix:

>>> jax.tree_util.tree_map(lambda x, y: [x] + y, [5, 6], [[7, 9], [1, 2]])
[[5, 7, 9], [6, 1, 2]]