GetAbstractModel

Get Abstract model definition.

Abstract Signature:

GetAbstractModel(init_fn, mesh)

JAX (Core)Flax NNX

JAX (Core)

API: —

Strategy: Custom / Partial

Flax NNX

API: flax.nnx.get_abstract_model

Strategy: Direct Mapping

Official Docs ↗