class_docstr: Tests correctness of the Y = g(X) = FFJORD(X) transformation.
code: def testJacobianScaling(self, dtype):
      tf_dtype = tf.as_dtype(dtype)
      scaling_by_two_exp = np.log(2.0).astype(dtype)
      scale_ode_fn = lambda t, z: scaling_by_two_exp * 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.array([[0.0], [1.0]], dtype=dtype)
      y = np.array([[0.0], [2.0]], dtype=dtype)
      expected_forward_log_det_jacobian = np.array(
          [np.log(2.0)] * 2, dtype=dtype
      )
      expected_inverse_log_det_jacobian = np.array(
          [np.log(0.5)] * 2, dtype=dtype
      )
      self.assertAllClose(
          y, self.evaluate(bijector.forward(x)), rtol=0.0, atol=1e-4
      )
      self.assertAllClose(
          x, self.evaluate(bijector.inverse(y)), rtol=0.0, atol=1e-4
      )
      self.assertAllClose(
          expected_inverse_log_det_jacobian,
          self.evaluate(
              bijector.inverse_log_det_jacobian(y, event_ndims=1)
          ),
      )
      self.assertAllClose(
          expected_forward_log_det_jacobian,
          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 3 for calling the FFJORDBijectorTest obj's underlying member method code with example usage: obj.testJacobianScaling(dtype) without return types