Merge

Merges a GraphDef and State back into a module (NNX Functional API).

Abstract Signature:

Merge(graphdef: GraphDef, state: State)

PyTorchFlax NNX

PyTorch

API: —

Strategy: Plugin (state_dict_load_shim)

Flax NNX

API: flax.nnx.merge

Strategy: Direct Mapping

Official Docs ↗