class_docstr: Direct tests for `_make_gradient_fn_of_y`.
class_docstr_tok: ["Direct","tests","for","`","_make_gradient_fn_of_y","`","."]
code: def test_grad_fn_of_y_grad_with_argument(self):
      if not tf1.control_flow_v2_enabled():
          self.skipTest(
              "TF1 appears to find incorrect gradients. This is likely "
              "not worth debugging."
          )
      p = tf.convert_to_tensor([2.0, 3.0, 4.0])
      y = tf.convert_to_tensor([2.0, 4.0, 6.0])
      x = tf.math.pow(y, 1.0 / p)
      grad_fn_of_y = sfwii._make_gradient_fn_of_y(
          lambda x, p: x**p, x
      )
      true_dy_dx_fn = lambda y, p: p * tf.math.pow(y, (p - 1) / p)
      true_dy_dpower_fn = lambda y, p: y * tf.math.log(y) / p
  
      (
          dy_dx,
          (grad_dy_dx_wrt_y, grad_dy_dx_wrt_power),
      ) = gradient.value_and_gradient(
          lambda y, p: grad_fn_of_y(y, p)[0], y, p
      )
      (
          true_dy_dx,
          (true_grad_dy_dx_wrt_y, true_grad_dy_dx_wrt_power),
      ) = gradient.value_and_gradient(true_dy_dx_fn, y, p)
      self.assertAllClose(dy_dx, true_dy_dx)
      self.assertAllClose(grad_dy_dx_wrt_y, true_grad_dy_dx_wrt_y)
      self.assertAllClose(
          grad_dy_dx_wrt_power, true_grad_dy_dx_wrt_power
      )
  
      (
          dy_dpower,
          (grad_dy_dpower_wrt_y, grad_dy_dpower_wrt_power),
      ) = gradient.value_and_gradient(
          lambda y, p: grad_fn_of_y(y, p)[1], y, p
      )
      (
          true_dy_dpower,
          (true_grad_dy_dpower_wrt_y, true_grad_dy_dpower_wrt_power),
      ) = gradient.value_and_gradient(true_dy_dpower_fn, y, p)
      self.assertAllClose(dy_dpower, true_dy_dpower)
      self.assertAllClose(
          grad_dy_dpower_wrt_y, true_grad_dy_dpower_wrt_y
      )
      self.assertAllClose(
          grad_dy_dpower_wrt_power, true_grad_dy_dpower_wrt_power
      )
  
imports: from absl.testing import parameterized
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import chain
  from tensorflow_probability.python.bijectors import power
  from tensorflow_probability.python.bijectors import scale
  from tensorflow_probability.python.bijectors import shift
  from tensorflow_probability.python.distributions import beta
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import student_t
  from tensorflow_probability.python.distributions import uniform
  from tensorflow_probability.python.experimental.bijectors import scalar_function_with_inferred_inverse
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import gradient
label_desc: Use image node_id 2 for calling the TestMakeGradientFunctionOfY obj's underlying member method code with example usage: obj.test_grad_fn_of_y_grad_with_argument() without return types