code: def _CheckFuns(self, optimizer, loss, x0, *args):
      init_fun, update_fun, get_params = optimizer(*args)
      opt_state = init_fun(x0)
      self.assertAllClose(x0, get_params(opt_state))
      opt_state2 = update_fun(
          0, grad(loss)(x0), opt_state
      )  # doesn't crash
      self.assertEqual(
          tree_util.tree_structure(opt_state),
          tree_util.tree_structure(opt_state2),
      )
  
imports: import functools
  from absl.testing import absltest
  import numpy
  import jax.numpy
  import jax._src.test_util
  from jax import jit, grad, jacfwd, jacrev
  from jax import tree_util
  from jax import lax
  from jax.example_libraries import optimizers
  from jax import config
label_desc: Use image node_id 2 for calling the OptimizerTests obj's underlying member method code with example usage: obj._CheckFuns(optimizer, loss, x0) without return types