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

Update Keras version.

Fix `save_model` and `load_model` (#19924)

Update version number