MapState

Maps a function over a State object.

Abstract Signature:

MapState(f: Callable, state: State)

JAX (Core)Flax NNX

JAX (Core)

API: jax.tree_util.tree_map

Strategy: Direct Mapping

Official Docs ↗

Flax NNX

API: flax.nnx.map_state

Strategy: Direct Mapping

Official Docs ↗