flax

5e60ced0 - Fix Flax GitHub Actions CI (`pytest` `DeprecationWarning` and `pytype` in `examples/gemma/train.py`).

Fix Flax GitHub Actions CI (`pytest` `DeprecationWarning` and `pytype` in `examples/gemma/train.py`). 1. Add `DeprecationWarning` for using `Module.sow()` outside `nnx.capture()` to `filterwarnings` in `oss/pyproject.toml` so `test_sow_capture_noop_config_flag` passes under `filterwarnings = ["error", ...]`. 2. Use `jnp.astype(input_mask, jnp.float32)` in `examples/gemma/train.py` so `pytype` does not fail with `attribute-error` (`No attribute 'astype' on bool`). PiperOrigin-RevId: 995384363