penzai.experimental.v2.core#
Submodules
Helper class to define automatic arbitrary orderings across different types. |
|
A stateful random stream helper. |
|
Support for named pockets of mutable/shared state within a JAX pytree. |
Submodules
Helper class to define automatic arbitrary orderings across different types. |
|
A stateful random stream helper. |
|
Support for named pockets of mutable/shared state within a JAX pytree. |