UpdateState

Updates a graph node with a new state dictionary in-place.

Abstract Signature:

UpdateState(node, state)

PyTorchJAX (Core)KerasFlax NNX

PyTorch

API: —

Strategy: Macro '{node}.load_state_dict({state})'

JAX (Core)

API: —

Strategy: Custom / Partial

Keras

API: keras.Model.set_weights

Strategy: Direct Mapping

Official Docs ↗

Flax NNX

API: flax.nnx.update

Strategy: Direct Mapping

Official Docs ↗