code: def test_updates_optimizer_state(self):
      final_x, losses = self.evaluate(
          minimize_stateless(
              lambda x: x**2,
              init=10.0,
              num_steps=3,
              optimizer=StatelessSGD(
                  learning_rate=1.0, decay_learning_rate=True
              ),
          )
      )
      # The loss has gradient `2 * x`, so with decaying learning rate the
      # optimization path should be:
      # step   x    loss   decayed_learning_rate   grad    new_x
      # 0     10.   100.   1.                       20.     -10.
      # 1    -10.   100.   1. / 2.                 -20.       0.
      # 2      0.   0.     1. / 3.                   0.       0.
      # ...    0.   0.     ...
      # (with no decay, `x` will just bounce between -10. and 10. forever).
      self.assertAllClose(final_x, 0.0)
      self.assertAllClose(losses, [100.0, 100.0, 0.0])
  
imports: import collections
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python import optimizer
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.internal import tf_keras
  from tensorflow_probability.python.math.minimize import minimize
  from tensorflow_probability.python.math.minimize import minimize_stateless
label_desc: Use image node_id 2 for calling the MinimizeStatelessTests obj's underlying member method code with example usage: obj.test_updates_optimizer_state() without return types