MergeState

Merges multiple State objects.

Abstract Signature:

MergeState(state: Mapping, states)

JAX (Core)Flax NNX

JAX (Core)

API: jax.tree_util.tree_map

Strategy: Direct Mapping

Official Docs ↗

Flax NNX

API: flax.nnx.merge_state

Strategy: Direct Mapping

Official Docs ↗