Can't get optimizer to apply gradients with Keras and DTensor based model
@pentney,
In the given code snippet you have defined the class and its methods but are not calling them anywhere. Could you please provide the complete code to debug the issue. Kindly find the gist of it [here](https://colab.research.google.com/gist/tilakrayal/d4baeba4cdc407efec93501bbb2ce525/untitled1374.ipynb). Thank you! This issue is stale because it has been open for 7 days with no activity. It will be closed if no further activity occurs. Thank you. This issue was closed because it has been inactive for 7 days since being marked as stale. Please reopen if you'd like to work on this further. 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/61932">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/61932">No</a>

### Issue type

Support

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

No

### Source

source

### TensorFlow version

2.13.0 (also tried with 2.9.1)

### Custom code

Yes

### OS platform and distribution

Linux Ubuntu 16.04

### Mobile device

_No response_

### Python version

3.8

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

Trying to do distributed training, for a rather wide model, with DTensors and Keras. I'm creating an Adam optimizer, trying to update it manually with the following code for a custom DTensor training step (this is within my own model object, with self.model being a Keras model):

 ```
  @tf.function
    def train_step(self, x, y, w, optimizer, train_metrics): 
        if not self.init:   # this runs once to initialize the model variables
            self.model(x)
            self.init = True 
        with tf.GradientTape() as tape:
            logits = self.model(x, training=True)
            logits = tf.reshape(logits, (logits.shape[1], logits.shape[0]))
            loss = tf.reduce_sum(tf.math.multiply(
                tf.keras.losses.binary_crossentropy(
                    y, logits, from_logits=True), w))
        gradients = tape.gradient(loss, self.model.trainable_variables)

        optimizer.apply_gradients(zip(gradients, self.model.trainable_variables))

        loss_per_sample = loss / len(x)
        results = {'loss': loss_per_sample}
        for metric in train_metrics.values():
            metric.update_state(y_true=y, y_pred=logits)
        return results 
        
    def dtensor_fit(        
        self,
        x_train,
        y_train,
        x_val,
        y_val,
        w_train=None):
        num_epochs = 5
        train_metrics = {
            "accuracy" : metrics.Accuracy(),
            "tp": metrics.TruePositives(),
            "fp": metrics.FalsePositives(),
            "fn": metrics.FalseNegatives(),
            "tn": metrics.TrueNegatives(),
            "auc": metrics.AUC(curve="PR"),
        }
        optimizer = tf.keras.dtensor.experimental.optimizers.Adam(mesh=self.mesh)

        eval_metrics = dict(train_metrics)   
        for epoch in range(num_epochs):
            print("============================") 
            print("Epoch: ", epoch)
            for metric in train_metrics.values():
                metric.reset_state()
            step = 0
            results = {}
            pbar = tf.keras.utils.Progbar(target=None, stateful_metrics=[])
            def batch(x, y, w, n):
                num_samples = x.shape[0]
                l = 0
                while l < num_samples:
                    yield x[l:l+n], y[l:l+n], w[l:l+n]
                    l += n
            self.init = False
            for inputs, labels, weights in batch(x_train, y_train, w_train, NNModel.BATCH_SIZE):
                indices = np.transpose(inputs.nonzero())
                inputs.eliminate_zeros()
                values = inputs.data
                inputs, labels, weights = self.pack_dtensor_inputs(
                    tf.SparseTensor(indices=indices, values=values, dense_shape=(inputs.shape[0], self.input_size)),
                    tf.convert_to_tensor([labels], dtype =tf.bfloat16),
                    tf.convert_to_tensor([weights], dtype=tf.bfloat16), self.input_layout, self.label_layout, self.weight_layout)
                optimizer.build(self.model.trainable_variables)

                results.update(self.train_step(inputs, labels, weights, optimizer, train_metrics))
                for metric_name, metric in train_metrics.items():
                    results[metric_name] = metric.result()

                pbar.update(step, values=results.items(), finalize=False)
                step += 1 
``` 

I get the following error when I reach the apply_gradients step in the train_step function:

ValueError: in user code:

```
File "/usr/local/home/bill/.cache/bazel/_bazel_bill/254c50c69c3701cca4e904bef759573b/execroot/_main/bazel-out/k8-fastbuild/bin/keystone/training/training.runfiles/_main/keystone/training/model.py", line 263, in train_step  *
    optimizer.apply_gradients(zip(gradients, self.model.trainable_variables))
File "/usr/local/home/bill/.cache/bazel/_bazel_bill/254c50c69c3701cca4e904bef759573b/execroot/_main/bazel-out/k8-fastbuild/bin/keystone/training/training.runfiles/rules_python~0.21.0~pip~pip_keras/site-packages/keras/dtensor/optimizers.py", line 141, in apply_gradients  **
    optimizer_lib._BaseOptimizer.apply_gradients(self, grads_and_vars)
File "/usr/local/home/bill/.cache/bazel/_bazel_bill/254c50c69c3701cca4e904bef759573b/execroot/_main/bazel-out/k8-fastbuild/bin/keystone/training/training.runfiles/rules_python~0.21.0~pip~pip_keras/site-packages/keras/optimizers/optimizer.py", line 650, in apply_gradients
    iteration = self._internal_apply_gradients(grads_and_vars)
File "/usr/local/home/bill/.cache/bazel/_bazel_bill/254c50c69c3701cca4e904bef759573b/execroot/_main/bazel-out/k8-fastbuild/bin/keystone/training/training.runfiles/rules_python~0.21.0~pip~pip_keras/site-packages/keras/dtensor/optimizers.py", line 153, in _internal_apply_gradients
    optimizer_lib._BaseOptimizer._internal_apply_gradients(
File "/usr/local/home/bill/.cache/bazel/_bazel_bill/254c50c69c3701cca4e904bef759573b/execroot/_main/bazel-out/k8-fastbuild/bin/keystone/training/training.runfiles/rules_python~0.21.0~pip~pip_keras/site-packages/keras/optimizers/optimizer.py", line 680, in _internal_apply_gradients
    self._update_step(grad, var)
File "/usr/local/home/bill/.cache/bazel/_bazel_bill/254c50c69c3701cca4e904bef759573b/execroot/_main/bazel-out/k8-fastbuild/bin/keystone/training/training.runfiles/rules_python~0.21.0~pip~pip_keras/site-packages/keras/optimizers/optimizer.py", line 240, in _update_step
    self.update_step(gradient, variable)
File "/usr/local/home/bill/.cache/bazel/_bazel_bill/254c50c69c3701cca4e904bef759573b/execroot/_main/bazel-out/k8-fastbuild/bin/keystone/training/training.runfiles/rules_python~0.21.0~pip~pip_keras/site-packages/keras/optimizers/adam.py", line 194, in update_step
    m.assign_add((gradient - m) * (1 - self.beta_1))

ValueError: Dimensions must be equal, but are 5 and 0 for '{{node sub_2}} = Sub[T=DT_FLOAT](gradient_tape/keystone/feature/MatMul_1/Cast/Cast, sub_2/ReadVariableOp)' with input shapes: [8299614,5], [0].

```
It looks like the optimizer momentums are not of the correct shape. I tried explicitly adding `optimizer.build(self.model.trainable_variables) `after the `self.model(x)` in `train_step` to force correct population, but the variables are still wrong.

I have seen the same results with TF versions 2.9.1 and 2.13.0.

Is this a bug? How do I get the optimizer to be correctly set up for training?



### Standalone code to reproduce the issue

```shell
see above
```


### Relevant log output

_No response_