class_docstr: Tests correctness of the Y = g(X) = FFJORD(X) transformation.
code: def testBijectorConditionKwargs(self, dtype):
      if not tf1.control_flow_v2_enabled():
          self.skipTest("b/152464477")
  
      tf_dtype = tf.as_dtype(dtype)
  
      def conditional_ode_fn(t, z, c):
          del t  # unused.
          return tf.ones_like(z) * c**2
  
      trace_augmentation_fn = ffjord.trace_jacobian_exact
      bijector = ffjord.FFJORD(
          trace_augmentation_fn=trace_augmentation_fn,
          state_time_derivative_fn=conditional_ode_fn,
          dtype=tf_dtype,
      )
      x = tf.zeros((2, 5), dtype=tf_dtype)
      y = tf.ones((2, 5), dtype=tf_dtype) * 4
      c = tf.ones((2, 5), dtype=tf_dtype) * 2
      expected_log_det_jacobian = np.zeros(2, dtype=dtype)
      expected_dy_dc = np.ones((2, 5), dtype=dtype) * 4
  
      def grad_fn(c):
          y = bijector.forward(x, c=c)
          return y
  
      dy_dc = self.evaluate(
          tfp_gradient.value_and_gradient(grad_fn, c)[1]
      )
  
      self.assertStartsWith(bijector.name, "ffjord")
      self.assertAllClose(
          self.evaluate(y),
          self.evaluate(bijector.forward(x, c=c)),
          atol=1e-5,
      )
      self.assertAllClose(
          self.evaluate(x),
          self.evaluate(bijector.inverse(y, c=c)),
          atol=1e-5,
      )
      self.assertAllClose(
          expected_log_det_jacobian,
          self.evaluate(
              bijector.inverse_log_det_jacobian(y, event_ndims=1, c=c)
          ),
      )
      self.assertAllClose(
          expected_log_det_jacobian,
          self.evaluate(
              bijector.forward_log_det_jacobian(x, event_ndims=1, c=c)
          ),
      )
      self.assertAllClose(expected_dy_dc, dy_dc)
  
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 2 for calling the FFJORDBijectorTest obj's underlying member method code with example usage: obj.testBijectorConditionKwargs(dtype) and returns: unknown, y