flax

a1cbcab4 - Fix pre-existing Pyrefly type-check errors in Flax and remove migration suppressions

Fix pre-existing Pyrefly type-check errors in Flax and remove migration suppressions - Remove the temporary per-file `[[tool.pyrefly.sub-config]]` migration error suppressions from `pyproject.toml` now that `pyrefly` targets Python 3.14. - Initialize conditional local variables before branching in `flax/core/nn/attention.py`, `flax/linen/linear.py`, `flax/nnx/nn/linear.py`, and `flax/nnx/transforms/autodiff.py`, and narrow `promote_dtype` outputs in `DenseGeneral`. - Annotate `Scope._submodule_adoption_cache` and `Scope.rng_counters` in `flax/core/scope.py`, and align `fp8_meta_dtype_wrapper._rules` in `flax/linen/fp8_ops.py` with `jax.dtypes.ExtendedDType`. - Fix `ModuleMeta.__call__` typing and `functools.partial` `__self__` attachment in `flax/nnx/bridge`, `classes: set[type[Any]]` and `sow`/`capture` typing in `flax/nnx/module.py`, `split` / `split_state` overload signatures in `flax/nnx/statelib.py`, `Stage.args_info` property overrides and `Updates` indexing in `flax/nnx/transforms/compilation.py`, and `StateAxes` mapping narrowing in `flax/nnx/transforms/iteration.py`. - Align `HiPrim` method signatures (`expand`, `vjp_fwd`, `vjp_bwd_retval`, `transpose`) and expose `TYPE_CHECKING` binary operator stubs on `Variable` in `flax/nnx/variablelib.py`, and type `_split_mp_arrays` (`PyTree`) and `os.fspath` `str` narrowing in `flax/training/checkpoints.py`. PiperOrigin-RevId: 1013412680