UseEagerSharding

Sets whether Variables should use eager sharding by default.

Abstract Signature:

UseEagerSharding(value: bool)

JAX (Core)Flax NNX

JAX (Core)

API: —

Strategy: Custom / Partial

Flax NNX

API: flax.nnx.use_eager_sharding

Strategy: Direct Mapping

Official Docs ↗