code: def test_adam(self):
      lin = haiku.transform(snt.Linear(1))
      x = tf.ones([1, 1])
      params = lin.init(x)
  
      optimizer = adam(learning_rate=0.01)
      opt_state = optimizer.init(params)
      # Step + (m, v) per parameter.
      self.assertLen(tree.flatten(opt_state), 5)
  
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 2 for calling the OptimizersTest obj's underlying member method code with example usage: obj.test_adam() without return types