docstr: Tests that the diagonal of the Jacobian matrix computes correctly.
docstr_tok: ["Tests","that","the","diagonal","of","the","Jacobian","matrix","computes","correctly","."]
code: def testJacobianDiagonal3DListInput(self):
      """Tests that the diagonal of the Jacobian matrix computes correctly."""
  
      dtype = np.float32
      true_mean = dtype([0, 0, 0])
      true_cov = dtype(
          [[1, 0.25, 0.25], [0.25, 2, 0.25], [0.25, 0.25, 3]]
      )
      chol = tf.linalg.cholesky(true_cov)
      target = mvn_tril.MultivariateNormalTriL(
          loc=true_mean, scale_tril=chol
      )
  
      # Assume that the state is passed as a list of tensors `x` and `y`.
      # Then the target function is defined as follows:
      def target_fn(x, y):
          # Stack the input tensors together
          z = tf.concat([x, y], axis=-1) - true_mean
          return target.log_prob(z)
  
      sample_shape = [3, 5]
      state = [
          tf.ones(sample_shape + [2], dtype=dtype),
          tf.ones(sample_shape + [1], dtype=dtype),
      ]
      fn_val, grads = gradient.value_and_gradient(target_fn, state)
      grad_fn = lambda *args: gradient.value_and_gradient(
          target_fn, args
      )[1]
  
      _, diag_jacobian_shape_passed = diag_jacobian(
          xs=state, ys=grads, fn=grad_fn, sample_shape=ps.shape(fn_val)
      )
      _, diag_jacobian_shape_none = diag_jacobian(
          xs=state, ys=grads, fn=grad_fn
      )
  
      true_diag_jacobian_1 = np.zeros(sample_shape + [2])
      true_diag_jacobian_1[..., 0] = -1.05
      true_diag_jacobian_1[..., 1] = -0.52
  
      true_diag_jacobian_2 = -0.34 * np.ones(sample_shape + [1])
  
      self.assertAllClose(
          self.evaluate(diag_jacobian_shape_passed[0]),
          true_diag_jacobian_1,
          atol=0.01,
          rtol=0.01,
      )
      self.assertAllClose(
          self.evaluate(diag_jacobian_shape_none[0]),
          true_diag_jacobian_1,
          atol=0.01,
          rtol=0.01,
      )
  
      self.assertAllClose(
          self.evaluate(diag_jacobian_shape_passed[1]),
          true_diag_jacobian_2,
          atol=0.01,
          rtol=0.01,
      )
      self.assertAllClose(
          self.evaluate(diag_jacobian_shape_none[1]),
          true_diag_jacobian_2,
          atol=0.01,
          rtol=0.01,
      )
  
file_path: probability/tensorflow_probability/python/math/diag_jacobian_test.py
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import mvn_tril
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import gradient
  from tensorflow_probability.python.math.diag_jacobian import diag_jacobian
label_desc: Use image node_id 1 for calling the JacobianTest obj's underlying member method code with example usage: obj.testJacobianDiagonal3DListInput() and returns: target