penzai.core

Penzai core types and functions.

Submodules

penzai.core.context	Helpers for scoped contextual value providers.
penzai.core.dataclass_util	Utilities for working with dataclasses.
penzai.core.formatting_util	Utilities for formatting and rendering.
penzai.core.layer	Penzai's base layer class and associated utilities.
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 simple random stream abstraction.
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.