code: def testRematLambdaFunction(self):
      f = lambda a, b: tf_np.sum(tf_np.sqrt(tf_np.exp(a)) + b)
      f_remat = extensions.remat(f)
  
      shape = [10]
      a = tf_np.random.randn(*shape)
      b = tf_np.random.randn(*shape)
  
      actual = extensions.grad(f_remat)(a, b)
      expected = extensions.grad(f)(a, b)
      self.assertAllClose(actual, expected)
  
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 45 for calling the ExtensionsTest obj's underlying member method code with example usage: obj.testRematLambdaFunction() without return types