RegisterDataType

Registers a type as pytree data type recognized by Object.

Abstract Signature:

RegisterDataType(type_: type)

JAX (Core)

API: jax.tree_util.register_pytree_node
Strategy: Direct Mapping
Official Docs ↗

Flax NNX

API: flax.nnx.register_data_type
Strategy: Direct Mapping
Official Docs ↗