Tf 2.11 Optimizers. Does anyone have any custom optimizer for the new version?
class Gravity(tf.keras.optimizers.Optimizer):
  def __init__(self,
               learning_rate=0.1,
               alpha=0.01,
               beta=0.9,
               name="Gravity",
               **kwargs):
    super(Gravity, self).__init__(name, **kwargs)
    self._set_hyper('learning_rate', kwargs.get('lr', learning_rate))
    self._set_hyper('alpha', alpha)
    self._set_hyper('beta', beta)
    self.epsilon = 1e-7

  def _create_slots(self, var_list):
    alpha = self._get_hyper("alpha")
    stddev = alpha / self._hyper['learning_rate']
    initializer = tf.keras.initializers.RandomNormal(mean=0.0,
                                                     stddev=stddev,
                                                     seed=None)
    for var in var_list:
      self.add_slot(var, "velocity", initializer=initializer)

  @tf.function
  def _resource_apply_dense(self, grad, var):
    # Get Data
    var_dtype = var.dtype.base_dtype
    lr_t = self._get_hyper('learning_rate', var_dtype)
    beta = self._get_hyper("beta", var_dtype)
    t = tf.cast(self.iterations, float)
    beta_hat = (beta * t + 1) / (t + 2)
    velocity = self.get_slot(var, "velocity")

    # Calculations
    max_step_grad = 1 / tf.math.reduce_max(tf.math.abs(grad))
    gradient_term = grad / (1 + (grad / max_step_grad)**2)

    # update variables
    updated_velocity = velocity.assign(beta_hat * velocity +
                                       (1 - beta_hat) * gradient_term)
    updated_var = var.assign(var - lr_t * updated_velocity)

  def get_config(self):
    config = super(Gravity, self).get_config()
    config.update({
        'learning_rate': self._get_hyper('learning_rate'),
        'alpha': self._get_hyper('alpha'),
        'beta': self._get_hyper('beta'),
        'epsilon': self.epsilon,
    })
    return config
 Thank you, sir, but still not working. Let me know if this one works:

```
from tensorflow.keras import optimizers

class Gravity(optimizers.Optimizer):
    def __init__(self, learning_rate=0.1, alpha=0.01, beta=0.9, name="Gravity", **kwargs):
        super(Gravity, self).__init__(name, **kwargs)
        self._set_hyper("learning_rate", kwargs.get("lr", learning_rate))
        self._set_hyper("decay", self._initial_decay)
        self._set_hyper("alpha", alpha)
        self._set_hyper("beta", beta)
        self.epsilon = 1e-7

    def _create_slots(self, var_list):
        alpha = self._get_hyper("alpha")
        stddev = alpha / self._get_hyper("learning_rate")
        initializer = tf.keras.initializers.RandomNormal(mean=0.0, stddev=stddev, seed=None)
        for var in var_list:
            self.add_slot(var, "velocity", initializer=initializer)

    @tf.function
    def _resource_apply_dense(self, grad, var):
        var_dtype = var.dtype.base_dtype
        lr_t = self._decayed_lr(var_dtype)
        beta = self._get_hyper("beta", var_dtype)
        t = tf.cast(self.iterations, var_dtype)
        beta_hat = (beta * t + 1) / (t + 2)
        velocity = self.get_slot(var, "velocity")

        max_step_grad = 1 / tf.math.reduce_max(tf.math.abs(grad))
        gradient_term = grad / (1 + (grad / max_step_grad)**2)

        updated_velocity = velocity.assign(beta_hat * velocity + (1 - beta_hat) * gradient_term)
        updated_var = var.assign_sub(lr_t * updated_velocity)

        return tf.group(updated_var.op, updated_velocity.op)

    def _resource_apply_sparse(self, grad, var):
        raise NotImplementedError

    def get_config(self):
        config = super(Gravity, self).get_config()
        config.update({
            "learning_rate": self._serialize_hyperparameter("learning_rate"),
            "decay": self._serialize_hyperparameter("decay"),
            "alpha": self._serialize_hyperparameter("alpha"),
            "beta": self._serialize_hyperparameter("beta"),
            "epsilon": self.epsilon,
        })
        return config

``` Still don't work "'Gravity' object has no attribute '_set_hyper'" @thewizardnet,
I tried to execute the given code on tensorflow v2.13.0-dev20230215 and was able to execute without any issue/error.
Kindly find the gist of it [here](https://colab.research.google.com/gist/tilakrayal/73916eea18f60811e7126614a21548bf/untitled965.ipynb).

The Optimizer issue comes from the fact that we call a **private _decayed_lr** method on the optimizer, which is not supported on the LossScaleOptimizer:

Please take a look at this https://github.com/tensorflow/models/issues/9262#issuecomment-696420425 and follow the steps as provided in the comment. Thank you! Hey, @tilakrayal thank you so much. I did not have a chance to look too much into the docs and know the issues, but I am glad that I got the solution. 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/59654">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/59654">No</a>
 > @thewizardnet, I tried to execute the given code on tensorflow v2.13.0-dev20230215 and was able to execute without any issue/error. Kindly find the gist of it [here](https://colab.research.google.com/gist/tilakrayal/73916eea18f60811e7126614a21548bf/untitled965.ipynb).
> 
> The Optimizer issue comes from the fact that we call a **private _decayed_lr** method on the optimizer, which is not supported on the LossScaleOptimizer:
> 
> Please take a look at this [tensorflow/models#9262 (comment)](https://github.com/tensorflow/models/issues/9262#issuecomment-696420425) and follow the steps as provided in the comment. Thank you!

Hi @tilakrayal @thewizardnet, I have run the snippet, replaced the version with `!pip install tf-nightly==2.14.0.dev20230701`, and instantiated the defined classed as

```console
g = Gravity()
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
[<ipython-input-4-7aa23255a1d6>](https://localhost:8080/#) in <cell line: 1>()
----> 1 g = Gravity()

[<ipython-input-3-65b6fe06e5d4>](https://localhost:8080/#) in __init__(self, learning_rate, alpha, beta, name, **kwargs)
     12                  **kwargs):
     13         super(Gravity, self).__init__(name, **kwargs)
---> 14         self._set_hyper('learning_rate', kwargs.get('lr', learning_rate))
     15         self._set_hyper('decay', self._initial_decay)
     16         self._set_hyper('alpha', alpha)

AttributeError: 'Gravity' object has no attribute '_set_hyper'
```

Meanwhile, I think the main question here is whether we have an internal `_set_hyper` in the new `tf.keras.optimizers.Optimizer` to offer similar functionality as in `tf.keras.optimizers.legacy.Optimizer`. Can you shed some light on that?
I know that we can use tf.keras.optimizers.legacy.Optimizer for making the older custom optimizers to work,but I'm wonder how I can update my code.This the original code that I want to make it function for tf 2.11

`class Gravity(tf.keras.optimizers.Optimizer):
    def __init__(self,
                 learning_rate=0.1,
                 alpha=0.01,
                 beta=0.9,
                 name="Gravity",
                 **kwargs):
        super(Gravity, self).__init__(name, **kwargs)
        self._set_hyper('learning_rate', kwargs.get('lr', learning_rate))
        self._set_hyper('decay', self._initial_decay)
        self._set_hyper('alpha', alpha)
        self._set_hyper('beta', beta)
        self.epsilon = 1e-7

    def _create_slots(self, var_list):
        alpha = self._get_hyper("alpha")
        stddev = alpha / self.learning_rate
        initializer = tf.keras.initializers.RandomNormal(mean=0.0,
                                                         stddev=stddev,
                                                         seed=None)
        for var in var_list:
            self.add_slot(var, "velocity", initializer=initializer)

    @tf.function
    def _resource_apply_dense(self, grad, var):
        # Get Data
        var_dtype = var.dtype.base_dtype
        lr_t = self._decayed_lr(var_dtype)
        beta = self._get_hyper("beta", var_dtype)
        t = tf.cast(self.iterations, float)
        beta_hat = (beta * t + 1) / (t + 2)
        velocity = self.get_slot(var, "velocity")

        # Calculations
        max_step_grad = 1 / tf.math.reduce_max(tf.math.abs(grad))
        gradient_term = grad / (1 + (grad / max_step_grad)**2)

        # update variables
        updated_velocity = velocity.assign(beta_hat * velocity +
                                           (1 - beta_hat) * gradient_term)
        updated_var = var.assign(var - lr_t * updated_velocity)

        # updates = [updated_var, updated_velocity]
        # return tf.group(*updates)
    def _resource_apply_sparse(self, grad, var):
        raise NotImplementedError

    def get_config(self):
        config = super(Gravity, self).get_config()
        config.update({
            'learning_rate':
            self._serialize_hyperparameter('learning_rate'),
            'decay':
            self._serialize_hyperparameter('decay'),
            'alpha':
            self._serialize_hyperparameter('alpha'),
            'beta':
            self._serialize_hyperparameter('beta'),
            'epsilon':
            self.epsilon,
        })
        return config`
