Tensorflow Partitioned Variable checkpoint Inconsistency
I think logic that allows legacy.optimizer variables to handle sharding with checkpointing lives [here](https://github.com/keras-team/keras/blob/68f9af408a1734704746f7e6fa9cfede0d6879d8/keras/optimizers/legacy/optimizer_v2.py#L936), _create_slots_for_sharded_variables while new experimental.optimizer is missing equivalent [logic](https://github.com/keras-team/keras/blob/v2.14.0/keras/optimizers/optimizer.py). So maybe fits more as keras repo issue although you can trigger this bug without explicitly import keras or even referencing tf.keras once as default optimizer when model.compile(optimizer="Adam") will lead to this.

I see also a bit more sharded specific logic inside `get_slot` [here](https://github.com/keras-team/keras/blob/68f9af408a1734704746f7e6fa9cfede0d6879d8/keras/optimizers/legacy/optimizer_v2.py#L1090)

edit: Looks like this is key [pr](https://github.com/keras-team/keras/commit/75d70a610dffe927d89ceb400d79bb7f9027b26e) where legacy.optimizer got support for handling sharded variables well. Hi Mehdi,

Thanks for reporting this. I can reproduce the shape map you provided, however when I load the model with different sharding or no sharding, it loads correctly. Could you share the exact code you're using to load the checkpoint and the error message you are getting? Thanks! > however when I load the model with different sharding or no sharding, it loads correctly

My guess is you reloaded model, but didn't reload optimizer too.

```python
from typing import Any

import os
import tempfile

import tensorflow as tf
from tensorflow.python.distribute.sharded_variable import ShardedVariable


# More realistic example this would be done by Parameter Server Strategy.
def shard_variables_creator(partitioner) -> Any:
    def _creator(next_creator, **kwargs):
        var = next_creator(**kwargs)
        if var.shape.rank == 0:
            return var

        num_shards = partitioner(var.shape, var.dtype, axis=0)
        if num_shards[0] == 1:
            return var

        shard_count = num_shards[0]
        shards = []
        start = 0
        for index in range(shard_count):
            shard_name = var.name.removesuffix(":0") + f"/part_{index}"
            size = var.shape[0] // shard_count + (1 if var.shape[0] % shard_count > index else 0)
            shards.append(tf.Variable(var[start : start + size], name=shard_name))
            start += size

        return ShardedVariable(shards)

    return _creator


partitioner = tf.distribute.experimental.partitioners.MaxSizePartitioner(max_shard_bytes=100 * 16 * 4)

with tf.variable_creator_scope(shard_variables_creator(partitioner)):
    toy_model = tf.keras.Sequential(
        [tf.keras.layers.Embedding(100, 32), tf.keras.layers.Dense(1, activation="sigmoid")]
    )
    toy_model.compile(loss="binary_crossentropy", optimizer=tf.optimizers.experimental.Adam())
    toy_model.build(input_shape=(None, 1))
    toy_model.optimizer.build(toy_model.trainable_variables)  # type: ignore

temp_dir = tempfile.gettempdir()
weights_path = os.path.join(temp_dir, "model_weights")
toy_model.save_weights(weights_path)

toy_model2 = tf.keras.Sequential([tf.keras.layers.Embedding(100, 32), tf.keras.layers.Dense(1, activation="sigmoid")])
toy_model2.compile(loss="binary_crossentropy", optimizer=tf.optimizers.experimental.Adam())
toy_model2.build(input_shape=(None, 1))
toy_model2.optimizer.build(toy_model2.trainable_variables)  # type: ignore
toy_model2.load_weights(weights_path)
```

Crashes with an error message of,

```
    raise e.with_traceback(filtered_tb) from None
  File "/Users/pa-loaner/Snapchat/Dev/.venvs/bento/lib/python3.9/site-packages/tensorflow/python/training/saving/saveable_object_util.py", line 133, in restore
    raise ValueError(
ValueError: Received incompatible tensor with shape (50, 32) when attempting to restore variable with shape (100, 32) and name optimizer/_momentums/0/.ATTRIBUTES/VARIABLE_VALUE.
```

The issue is optimizer variables are not saved properly when partitioning is used. You can load model without optimizer, but if you try to reload model and optimizer (to continue training) it crashes. This bug is specific to new optimizer api and older optimizer api does not crash and handles optimizer variable partitioning as needed.
 Thanks! I've reproduced this now -- indeed the optimizer variables are saved in sharded form and probably shouldn't be. I'll look into this.

FYI, loading does work when the last lines are placed under a ParameterServerStrategy scope, i.e.:
```python
with strategy.scope():
  toy_model2 = tf.keras.Sequential([tf.keras.layers.Embedding(100, 32), tf.keras.layers.Dense(1, activation="sigmoid")])
  toy_model2.compile(loss="binary_crossentropy", optimizer=tf.optimizers.experimental.Adam())
  toy_model2.build(input_shape=(None, 1))
  toy_model2.optimizer.build(toy_model2.trainable_variables)  # type: ignore
  toy_model2.load_weights(weights_path)
```
regardless of the sharding used by the strategy. So if you are planning on restoring to continue training with PSStrategy, it should work as long as you load within scope like this. Else, report back here.  This script has 1 dependency (pip install portpicker) besides tf, but is self contained ps example.

```python
from typing import Any, Mapping

import json
import os
import tempfile
from multiprocessing import Process

import tensorflow as tf

from portpicker import pick_unused_port

__spec__ = None


def create_tf_configs(worker_count: int, ps_count: int):
    """Create TF_CONFIGs for a cluster."""
    cluster_dict: dict[str, list[str]] = {}
    if worker_count:
        cluster_dict["worker"] = [f"localhost:{pick_unused_port()}" for _ in range(worker_count)]
    if ps_count:
        cluster_dict["ps"] = [f"localhost:{pick_unused_port()}" for _ in range(ps_count)]

    cluster_dict["chief"] = [f"localhost:{pick_unused_port()}"]

    tf_configs = []
    for i in range(worker_count):
        tf_configs.append({"cluster": cluster_dict, "task": {"type": "worker", "index": i}})

    for i in range(ps_count):
        tf_configs.append({"cluster": cluster_dict, "task": {"type": "ps", "index": i}})

    tf_configs.append({"cluster": cluster_dict, "task": {"type": "chief", "index": 0}})

    return tf_configs


def _create_process(tf_config: Mapping[str, Any]):
    name = tf_config["task"]["type"] + "_" + str(tf_config["task"]["index"])

    print(f"Starting {name} process...")
    os.environ["TF_CONFIG"] = json.dumps(tf_config)
    p = Process(target=run)
    p.start()


def run():
    resolver = tf.distribute.cluster_resolver.TFConfigClusterResolver()

    task_type = resolver.task_type
    if task_type in ("worker", "ps"):
        print("Starting server...")
        server = tf.distribute.Server(
            resolver.cluster_spec(),
            job_name=resolver.task_type,
            task_index=resolver.task_id,
            protocol=resolver.rpc_layer,
            start=True,
        )
        server.join()

    partitioner = tf.distribute.experimental.partitioners.MaxSizePartitioner(max_shard_bytes=100 * 16 * 4)
    strategy = tf.distribute.experimental.ParameterServerStrategy(
        cluster_resolver=resolver, variable_partitioner=partitioner
    )

    print("Building model...")
    with strategy.scope():
        toy_model = tf.keras.Sequential(
            [tf.keras.layers.Embedding(100, 32), tf.keras.layers.Dense(1, activation="sigmoid")]
        )
        toy_model.compile(loss="binary_crossentropy", optimizer=tf.optimizers.experimental.Adam())
        toy_model.build(input_shape=(None, 1))
        toy_model.optimizer.build(toy_model.trainable_variables)  # type: ignore

    print("Saving weights...")
    temp_dir = tempfile.gettempdir()
    weights_path = os.path.join(temp_dir, "model_weights")
    toy_model.save_weights(weights_path)

    # No partitioner used for second ps strategy.
    strategy2 = tf.distribute.experimental.ParameterServerStrategy(cluster_resolver=resolver)
    with strategy2.scope():
        toy_model2 = tf.keras.Sequential(
            [tf.keras.layers.Embedding(100, 32), tf.keras.layers.Dense(1, activation="sigmoid")]
        )
        toy_model2.compile(loss="binary_crossentropy", optimizer=tf.optimizers.experimental.Adam())
        toy_model2.build(input_shape=(None, 1))
        toy_model2.optimizer.build(toy_model2.trainable_variables)  # type: ignore
        print("Loading weights...")
        toy_model2.load_weights(weights_path).assert_consumed()

    print("Done!")


def main():
    tf_configs = create_tf_configs(2, 1)
    chief_config = tf_configs[-1]
    for tf_config in tf_configs[:-1]:
        _create_process(tf_config)

    os.environ["TF_CONFIG"] = json.dumps(chief_config)
    run()


if __name__ == "__main__":
    main()
```

This script also fails and uses ps strategy for both original model and reloading. It does change partitioner which would be nice to do. Typical usage we have is train model on cluster for base model and then incrementally train it later (daily) on different clusters. It was common for number of partitions to change which also causes reloading to fail.

I think right now for model to be reloaded even in PS it requires partitioning to stay same. If you change the number of parameter servers/partitioning then it mismatches.

I do have workarounds for this, but they're complex and would be nice to save optimizer variables unpartitioned and support reloading them. My current workaround is I read checkpoint, rewrite optimizer variables in checkpoint with new target partitioning, and save it back. This relies on couple heuristics/examining checkpoint metadata to even find optimizer variables (including heuristic of if name looks like legacy optimizer skip it) Thanks for this script -- I'll have to dig in, there must be some other discrepancy with the test I'm using. I reload optimizer variables as well and the checkpoint reading works and re-partitions appropriately. The root cause is likely whatever is causing the checkpoint to contain variables in sharded form, and if we fix that the rest of the workflow will work as intended, so I'll start there. I just submitted https://github.com/keras-team/tf-keras/pull/703 which should resolve this, please reopen if not. You would need to wait for the next nightly release, or build from source, to test. Thanks for the report! Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/62215">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/62215">No</a>
 Thank you very much. I've confirmed the pr does fix my issue and tested it on some more cases where it behaved well.
### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

No

### Source

source

### TensorFlow version

2.14

### Custom code

No

### OS platform and distribution

Mac/Linux

### Mobile device

_No response_

### Python version

3.9

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

For tf.optimizers.experimental vs tf.optimizers.legacy, the former saves variables in checkpoint partitioned while latter does not. Model variables are saved not partitioned. This leads to save_weights/load_weights incompatibility when changing partition count. For example if you train model with 10 parameter servers using partitioner, save weights, and then load it outside PS with no partitioning the model fails to load.

There is even a test case focused on supporting saving -> loading model with different partitioning [here](https://github.com/tensorflow/tensorflow/blob/4dacf3f368eb7965e9b5c3bbdd5193986081c3b2/tensorflow/python/training/saver_test.py#L1070). I would not expect tf.optimizers.experimental.Optimizer variables to require partitioning to match while model variables/legacy optimizer do not.

I tried to use tf-nightly but it uses keras-nightly which is already keras 3 and has various breaking changes. model.save_weights does not supports non h5 format currently with nightly nor did tf.train.Checkpoint(model=model).write() work with the example code.

When you print checkpoint weights you can see model variable shapes are [100, 32], but optimizer variable shapes are [50, 32], [50, 32]. I've simplified code to make partitioned variables directly, but in realistic example they would be made by PS strategy.

### Standalone code to reproduce the issue

```shell
import tensorflow as tf
from tensorflow.python.distribute.sharded_variable import ShardedVariable
from typing import Any
import tempfile
import os

# More realistic example this would be done by Parameter Server Strategy.
def shard_variables_creator(partitioner) -> Any:
    def _creator(next_creator, **kwargs):
        var = next_creator(**kwargs)
        if var.shape.rank == 0:
            return var

        num_shards = partitioner(var.shape, var.dtype, axis=0)
        if num_shards[0] == 1:
            return var
        
        shard_count = num_shards[0]
        shards = []
        start = 0
        for index in range(shard_count):
            shard_name = var.name.removesuffix(":0") + f"/part_{index}"
            size = var.shape[0] // shard_count + (1 if var.shape[0] % shard_count > index else 0)
            shards.append(tf.Variable(var[start:start+size], name=shard_name))
            start += size
        
        return ShardedVariable(shards)
    
    return _creator

partitioner = tf.distribute.experimental.partitioners.MaxSizePartitioner(max_shard_bytes=100*16*4)

with tf.variable_creator_scope(shard_variables_creator(partitioner)):
    toy_model = tf.keras.Sequential([tf.keras.layers.Embedding(100, 32), tf.keras.layers.Dense(1, activation="sigmoid")])
    toy_model.compile(loss="binary_crossentropy", optimizer=tf.optimizers.experimental.Adam())
    toy_model.build(input_shape=(None, 1))
    toy_model.optimizer.build(toy_model.trainable_variables) # type: ignore

temp_dir = tempfile.gettempdir()
weights_path = os.path.join(temp_dir, "model_weights")
toy_model.save_weights(weights_path)

reader = tf.train.load_checkpoint(weights_path)
print(reader.get_variable_to_shape_map())
```


### Relevant log output

```shell
{'optimizer/_velocities/3/.ATTRIBUTES/VARIABLE_VALUE': [1],
 'optimizer/_velocities/2/.ATTRIBUTES/VARIABLE_VALUE': [32, 1],
 'optimizer/_velocities/1/.ATTRIBUTES/VARIABLE_VALUE': [50, 32],
 'optimizer/_momentums/3/.ATTRIBUTES/VARIABLE_VALUE': [1],
 'optimizer/_learning_rate/.ATTRIBUTES/VARIABLE_VALUE': [],
 'layer_with_weights-1/kernel/.ATTRIBUTES/VARIABLE_VALUE': [32, 1],
 'layer_with_weights-1/bias/.ATTRIBUTES/VARIABLE_VALUE': [1],
 'optimizer/_momentums/1/.ATTRIBUTES/VARIABLE_VALUE': [50, 32],
 'layer_with_weights-0/embeddings/.ATTRIBUTES/VARIABLE_VALUE': [100, 32],
 'optimizer/_velocities/0/.ATTRIBUTES/VARIABLE_VALUE': [50, 32],
 'optimizer/_momentums/2/.ATTRIBUTES/VARIABLE_VALUE': [32, 1],
 '_CHECKPOINTABLE_OBJECT_GRAPH': [],
 'optimizer/_momentums/0/.ATTRIBUTES/VARIABLE_VALUE': [50, 32],
 'optimizer/_iterations/.ATTRIBUTES/VARIABLE_VALUE': []}
```
