code: def testGradSimpleModel(self):
      (
          params,
          params_true,
          inputs,
          targets,
      ) = generate_params_inputs_targets()
  
      for _ in range(50):
          params = train_step(params, inputs, targets)
  
      # This is not trained super well, but it usually gets "close".
      self.assertAllClose(params[0], params_true[0], atol=1e-1)
      self.assertAllClose(params[1], params_true[1], atol=1e-1)
  
imports: import functools
  import itertools
  from absl import flags
  from absl.testing import parameterized
  from jax import lax
  import numpy
  import tensorflow.compat.v2
  from trax.tf_numpy import extensions
  import trax.tf_numpy.numpy
label_desc: Use image node_id 35 for calling the ExtensionsTest obj's underlying member method code with example usage: obj.testGradSimpleModel() without return types