class_docstr: Tests correctness of the Y = g(X) = FFJORD(X) transformation.
code: def testJacobianDiagonalScaling(self, dtype):
      tf_dtype = tf.as_dtype(dtype)
      num_dims = 10
      matrix_diagonal = np.random.uniform(size=[num_dims]).astype(dtype)
      scaling_matrix = np.diag(matrix_diagonal)
      one_time_scale_matrix = np.diag(np.exp(matrix_diagonal))
      scale_ode_fn = lambda t, z: tf.linalg.matvec(scaling_matrix, z)
      trace_augmentation_fn = ffjord.trace_jacobian_exact
      bijector = ffjord.FFJORD(
          trace_augmentation_fn=trace_augmentation_fn,
          state_time_derivative_fn=scale_ode_fn,
          dtype=tf_dtype,
      )
      x = np.random.uniform(size=[1, num_dims]).astype(dtype)
      y = np.matmul(x, one_time_scale_matrix)
  
      expected_forward_log_det_jacobian_value = np.log(
          np.prod(np.exp(matrix_diagonal))
      )
      expected_fldj = np.array(
          [expected_forward_log_det_jacobian_value]
      )
      expected_ildj = np.array(
          [-expected_forward_log_det_jacobian_value]
      )
  
      self.assertAllClose(
          y, self.evaluate(bijector.forward(x)), rtol=0.0, atol=1e-3
      )
      self.assertAllClose(
          x, self.evaluate(bijector.inverse(y)), rtol=0.0, atol=1e-3
      )
      self.assertAllClose(
          expected_ildj,
          self.evaluate(
              bijector.inverse_log_det_jacobian(y, event_ndims=1)
          ),
      )
      self.assertAllClose(
          expected_fldj,
          self.evaluate(
              bijector.forward_log_det_jacobian(x, event_ndims=1)
          ),
      )
  
file_path: probability/tensorflow_probability/python/bijectors/ffjord_test.py
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import ffjord
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import gradient
label_desc: Use image node_id 4 for calling the FFJORDBijectorTest obj's underlying member method code with example usage: obj.testJacobianDiagonalScaling(dtype) without return types