Version bump 3.11.2 and nnx fix #21565 (#21570)

* Fix nnx object state (#21565)

* Update operation.py

* Update actions.yml

* Update operation.py

* Update actions.yml

* Update operation.py

* Update operation.py

* Update operation.py

* fix test

* code reformat

* Version bump to 3.11.2

---------

Co-authored-by: Divyashree Sreepathihalli <divyashreepathihalli@gmail.com>

Version bump 3.11.1 (#21535)

* Version bump to 3.11.1

* Bug fix with serialization when walking through operations. (#21532)

Also update error message which was inaccurate.

---------

Co-authored-by: hertschuh <1091026+hertschuh@users.noreply.github.com>

Update flax (#21527)

* Update actions.yml

* Update requirements.txt

* Update requirements-jax-cuda.txt

* Update requirements-jax-cuda.txt

* Update requirements.txt

Update layer.py add_metric error message (#21290)

* Update layer.py add_metric error message

* Update layer.py

---------

Co-authored-by: François Chollet <francois.chollet@gmail.com>

Fix Remat error when called with a model (#21094)

* add print

* fix remat issue

* simplify code

* enable traceback filtering and update the function sig

* add a wrapper for activations

* change to except

* add layer call decorator

* fix remat call

Fix incorrect argument in JAX flash attention. (#21014)

The mask is named `array` in `NumpyMask`.

Fix `Discretization` serialization when `num_bins` is used. (#20971)

Previously, serialization / deserialization would fail if:
- the layer was saved / restored before `adapt` was called
- the layer was saved / restored after `adapt` was called, but the dataset was such that the number of bins learned was fewer than `num_bins`

The fix consists in adding a `from_config` to handle `bin_boundaries` separately. This is because at initial creation, `bin_boundaries` and `num_bins` cannot be both set, but when restoring the layer after `adapt`, they are both set.

Tightened the error checking:
- never allow `num_bins` and `bin_boundaries` to be specified at the same time, even if they match (same as `tf_keras`)
- don't allow `num_bins` and `bin_boundaries` to be `None` at the same time
- verify that `adapt` has been called in `call`

Also removed `init_bin_boundaries` as the value was never used and its presence can be inferred.

Fixing batch_dim_name attribute (#20674)

* fixing wrong trainer assumption that batch dim is always the first one in the mesh

* need functools partial

* lint

* fix test failure when distribution=None

* lint2

* fix for test failure

* added data sharding for 3D+ meshes

* lint3

* added @Property for batch_dim_name + refactoring

* fix typo

Better input validation for InputLayer with input_tensor provided

Fix optimizer docstring