penzai.core

Penzai core types and functions.

Submodules

penzai.core.auto_order_types	Helper class to define automatic arbitrary orderings across different types.
penzai.core.named_axes	A lightweight and minimal implementation of named axes.
penzai.core.partitioning	Tags and utility functions for working with partitioned models.
penzai.core.random_stream	A stateful random stream helper.
penzai.core.selectors	A toolkit for pytree manipulation.
penzai.core.shapecheck	Utilities for runtime shape-checking.
penzai.core.struct	Core struct abstraction.
penzai.core.syntactic_sugar	Syntactic sugar utilities for common operations.
penzai.core.tree_util	Additional tree functionality related to jax.tree_util.
penzai.core.variables	Support for named pockets of mutable/shared state within a JAX pytree.