Custom Keras Optimizer over TPU strategy error
Ah, just a little bit of an embarrassing mistake from my side. The optimizer needs to be called inside strategy.scope() block for it to work.
```

with strategy.scope():
  model = return_model()
  opt = GradientDescent(learning_rate = 0.001)
  model.compile(optimizer=opt, loss='mse')

  model(x_train[:2])
  model.summary()
  model.fit(x_train, y_train, epochs = 20, batch_size = 1024)
  opt.temp
```

Now, I'm just left with my second question only.

> Apart from that I wish the addition of a new feature - batch size of the data being currently used by the optimizer to compute gradients, as we know is pretty straightforward in a single CPU/GPU, but becomes difficult while executing the same code in distributed systems. If possible, kindly amend my existing standalone code below to add that small tiny feature into my Optimizer.

I'll close the issue after this. Hi @abhaskumarsinha ,

The Optimizers by default takes batch_size from model.fit only or from the tf.data.Dataset if set batch explicitly. I am not quite follow through your requirement of batch_size passing to Optimizer here. Normally for multi workers or nodes we can increase the batch_size by the no of GPUs we have.

Suppose we want `batch_size` of `64` on each of the GPU and total no of GPUs enabled under distribution strategy can be got by `strategy.num_replicas_in_sync`. Please refer the below code.

```
BATCH_SIZE_PER_REPLICA = 64
BATCH_SIZE = BATCH_SIZE_PER_REPLICA * strategy.num_replicas_in_sync
```
Then BATCH_SIZE shall be passed to model.fit.

Please refer to some official tutorials on how to configure distribution training on multiple workers here [link1](https://www.tensorflow.org/tutorials/distribute/keras) and [link2](https://www.tensorflow.org/tutorials/distribute/multi_worker_with_ctl).

If you still have any queries please come back and let us know.

Thanks!
 Reference: https://github.com/tensorflow/tensorflow/issues/60592#issuecomment-1563789462

Hello @SuryanarayanaY 

Thank you so much for your response and for bearing me with regarding the issue.
The thing is that we are working on some sort of special optimizer that works on variable _batch_sizes_ to make it work on some very specific non-convex functions on very large distributed systems.

So, I'm left with two requirements for now.
1. Have variable batch size for the optimizer (the optimizer decides the batch size of _$t_n$_ given the optimizer parameters at _$t_n-1$_.
2. We manually specify the upper bound of the batch size, i.e. all the _batch_size at all t_i_ for i in 0, 1, 2, ... should be less than _upper_bound_ i.e. _batch_size at t_i < upper_bound._ In that case, we mask the last entries of our gradients with 0 and calculate the gradient of only the required first _batch_size_ at _t_i_ entries and train the network

From my guess, 2 should be faster and should outperform 1 because the graph wouldn't need to be reconstructed using CUDA API for training the model, unlike case 1.

Since we need this optimizer to run on distributed training settings on a very large system, I wanted to make sure if any of these features are implementable on Keras or not. (Masking the gradient of certain batch indices with 0).

Thank you in Advance.
 HI @abhaskumarsinha ,

Thanks for your explanation on your requirement. If you want to customize the training step you need to build custom model class subclassing the `keras.Model` and need to override `train_step`. If you can able to accommodate your requirements in the `train_step` method then may be it can possible to implement your requirement. 

Please have a look into sample [tutorial](https://www.tensorflow.org/guide/keras/customizing_what_happens_in_fit) for customizing the train_step of keras model training.

Please have a look and let us know if it is helpful. Thanks! Hello @SuryanarayanaY 

Sorry for the late response. Yes, I've also checked to extend `tf.keras.Model` class and custom training step function. My only concern was whether this whole thing would be applicable to distribution strategies.

The second thing is to mask certain batches to 0 while keeping the gradients from other batches. If this gives control over batch-level editing and supports distribution, then this would definitely be helpful for me. Thank you so much. 


@abhaskumarsinha 

> Sorry for the late response. Yes, I've also checked to extend `tf.keras.Model` class and custom training step function. My only concern was whether this whole thing would be applicable to distribution strategies.

Ofcourse custom models also can be used with distribution strategies.Please refer a sample [tutorial](https://www.tensorflow.org/tutorials/distribute/custom_training).



> The second thing is to mask certain batches to 0 while keeping the gradients from other batches. If this gives control over batch-level editing and supports distribution, then this would definitely be helpful for me. Thank you so much.

I am not quiet follow through you here. AFAIK if you want some batches to be masked Zero then there is no standard way in Tensorflow.This may not possible during training but you can try with [tf.data.Dataset](https://www.tensorflow.org/api_docs/python/tf/data/Dataset) generators and implement your own logic for dataset generation.

Thanks!
 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/60711">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/60711">No</a>

<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Bug

### Have you reproduced the bug with TF nightly?

No

### Source

binary

### Tensorflow Version

v2.12.0-rc1-12-g0db597d0d75 2.12.0

### Custom Code

Yes

### OS Platform and Distribution

_No response_

### Mobile device

_No response_

### Python version

3.10.11

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current Behaviour?

Hello,

I'd wish to draw your attention to a bug that affects custom Optimizers extended from the class `tf.keras.optimizers.experimental.Optimizer`  while running on TPU clusters on native Google Colab.

```
_________________________________________________________________
Epoch 1/2
---------------------------------------------------------------------------
UnavailableError                          Traceback (most recent call last)
[<ipython-input-7-7647619ded22>](https://localhost:8080/#) in <cell line: 1>()
      5   model(x_train[:2])
      6   model.summary()
----> 7   model.fit(x_train, y_train, epochs = 2, batch_size = 1024)
      8   opt.temp

1 frames
[/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py](https://localhost:8080/#) in _numpy(self)
   1126       return self._numpy_internal()
   1127     except core._NotOkStatusException as e:  # pylint: disable=protected-access
-> 1128       raise core._status_to_exception(e) from None  # pylint: disable=protected-access
   1129 
   1130   @property

UnavailableError: failed to connect to all addresses; last error: UNKNOWN: ipv4:127.0.0.1:40812: Failed to connect to remote host: Connection refused
Additional GRPC error information from remote target /job:localhost/replica:0/task:0/device:CPU:0:
:UNKNOWN:failed to connect to all addresses; last error: UNKNOWN: ipv4:127.0.0.1:40812: Failed to connect to remote host: Connection refused {created_time:"2023-05-25T19:42:27.514045935+00:00", grpc_status:14}
```

While the code refuses to execute in order to compute gradients while using a custom optimizer - opt, it happens to run very well while using 'adam' on default settings.

The standalone code to reproduce the problem is provided and doesn't requires any extra tool or dependency and can be easily pasted into Colab to run and inspect the execution for bug investigations.

_______________________________________________________________________________________________________________________

Apart from that I wish addition of a new feature - *batch size* of the data being currently used by the optimizer to compute gradients, as we know is pretty straightforward in a single CPU/GPU, but becomes difficult while executing the same code in distributed systems. If possible, kindly amend my existing standalone code below to add that small tiny feature into my Optimizer.
Thank You.

Best Regards


### Standalone code to reproduce the issue

```shell
import numpy as np
import tensorflow as tf

class GradientDescent(tf.keras.optimizers.experimental.Optimizer):
  def __init__(self, learning_rate = 0.01, name='GDST'):
    super().__init__(name=name)
    #self.learning_rate = learning_rate
    self._learning_rate = self._build_learning_rate(learning_rate)
    self.temp = None
  
  def build(self, var_list):
    super().build(var_list)
  
  def update_step(self, gradient, variable):
    lr = tf.cast(self._learning_rate, gradient.dtype)
    output = tf.clip_by_value(self._learning_rate*gradient, clip_value_max = gradient.dtype.max, clip_value_min = gradient.dtype.min)
    variable.assign_sub(output)
    self.temp = output


  def get_config(self):
    return super().get_config()

opt = GradientDescent(learning_rate = 0.0001)

input_shape = 10000
output_shape = 100

def return_model(input_shape = 10000, output_shape = 500):
  model = tf.keras.Sequential([
    tf.keras.layers.Dense(input_shape, activation='relu'),
    tf.keras.layers.Dense(10000, activation='relu'),
    #tf.keras.layers.Dense(10000, activation='tanh'),
    #tf.keras.layers.Dense(10000, activation='linear'),
    #tf.keras.layers.Dense(10000, activation='tanh'),
    tf.keras.layers.Dense(output_shape, activation='linear')
])
  return model


# Compile the model
#model.compile(optimizer=opt, loss='mse')

# Create a random dataset
x_train = tf.random.uniform(shape=[10000, 10000])
y_train = tf.random.uniform(shape=[10000, 500])

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(resolver)
# This is the TPU initialization code that has to be at the beginning.
tf.tpu.experimental.initialize_tpu_system(resolver)
print("All devices: ", tf.config.list_logical_devices('TPU'))

strategy = tf.distribute.TPUStrategy(resolver)

with strategy.scope():
  model = return_model()
  model.compile(optimizer=opt, loss='mse')

  model(x_train[:2])
  model.summary()
  model.fit(x_train, y_train, epochs = 2, batch_size = 1024)
  opt.temp

```
```


### Relevant log output

_No response_</details>