Repository navigation
Tags
Tags
Release Flax 0.12.10 and use attribute lookup for HiPrim fallback Updates Flax version to 0.12.10 so the recent `jax.experimental.hijax` migration (`HiPrimitive` -> `HiPrim`) is included in a public release compatible with `jax>=0.11.2`. Also replaces the version-tuple check for `HiPrim` with an attribute lookup (`getattr(hjx, 'HiPrim', None) or hjx.VJPHiPrimitive`) and removes a shadowed duplicate definition of `_bind_new_variable` in `flax.nnx.variablelib`. PiperOrigin-RevId: 987564781
simplify set_view API * `nnx.view` now uses `inspect.signature` to automatically infer which parameters are required by `set_view` and only passes those. * updates the `view` guide to refect the above * simplifies the `set_view` definition for `Dropout`, `BatchNorm` and `MultiHeadAttention` PiperOrigin-RevId: 878229069
Rename sharding_names to out_sharding in NNX Variable metadata This CL renames the sharding_names attribute to out_sharding for better consistency with the sharding API. The new name more clearly indicates the purpose of this metadata field. ## Changes - Bump Flax version to 0.12.4 - Core changes in variablelib.py: - Add sharding_names to out_sharding metadata remapping for backward compatibility - Add deprecated sharding_names property that returns out_sharding with a warning - Update nnx/spmd.py, core/spmd.py, core/meta.py, linen/spmd.py to use out_sharding - Update all NNX tests to use the new attribute name - Update qwix flax_util.py to check for out_sharding first, with fallback to sharding_names - Update maxtext initializers.py to check for out_sharding first - Update documentation and examples to use out_sharding ## Backward Compatibility Existing code using sharding_names will continue to work via: - Metadata remapping during Variable creation - Deprecated Variable.sharding_names property PiperOrigin-RevId: 869269899
PreviousNext