code: def test_adam_with_variable_lr(self, trainable_lr):
      lin = haiku.transform(snt.Linear(1))
      x = tf.ones([1, 1])
      initial_params = lin.init(x)
  
      with haiku.variables():
          lr = tf.Variable(0.01, trainable=trainable_lr, name="lr")
  
      optimizer = adam(learning_rate=lr)
      initial_opt_state = optimizer.init(initial_params)
      # Learning rate, step + (m, v) per parameter.
      self.assertLen(tree.flatten(initial_opt_state), 6)
  
      grads = tree.map_structure(tf.ones_like, initial_params)
      params, opt_state = optimizer.apply(
          initial_opt_state, grads, initial_params
      )
  
      tree.assert_same_structure(initial_opt_state, opt_state)
      tree.assert_same_structure(initial_params, params)
  
file_path: sonnet/sonnet/src/functional/optimizers_test.py
imports: from absl.testing import parameterized
  import sonnet
  from sonnet.src import test_utils
  from sonnet.src.functional import haiku
  from sonnet.src.functional import optimizers
  import tensorflow
  import tree
label_desc: Use image node_id 3 for calling the OptimizersTest obj's underlying member method code with example usage: obj.test_adam_with_variable_lr(trainable_lr) without return types