ReplaceByPureDict

Updates the state of a module/node using a dictionary of values.

Abstract Signature:

ReplaceByPureDict(state, pure_dict: dict)

PyTorchKerasFlax NNX

PyTorch

API: load_state_dict

Strategy: Plugin (load_state_dict_wrapper)

Keras

API: set_weights

Strategy: Plugin (keras_weights_wrapper)

Flax NNX

API: flax.nnx.replace_by_pure_dict

Strategy: Direct Mapping

Official Docs ↗