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

Restore version string resolution in pip_build.

One time fix to _tf_keras API.

Move _tf_keras directory to the root of the pip package.

Fix JAX vectorize.