code: def testTheoreticalFldj(self):
      width = 4
      bijector = highway_flow.build_trainable_highway_flow(
          width,
          activation_fn=tf.nn.softplus,
          gate_first_n=2,
          seed=test_util.test_seed(),
      )
      self.evaluate(
          [v.initializer for v in bijector.trainable_variables]
      )
  
      x = self.evaluate(
          samplers.uniform(
              [width],
              minval=-1.0,
              maxval=1.0,
              seed=test_util.test_seed(sampler_type="stateless"),
          )
      )
      y = self.evaluate(bijector.forward(x))
      bijector_test_util.assert_bijective_and_finite(
          bijector,
          x,
          y,
          eval_func=self.evaluate,
          event_ndims=1,
          inverse_event_ndims=1,
          rtol=1e-5,
      )
      fldj = bijector.forward_log_det_jacobian(x, event_ndims=1)
      # The jacobian is not yet broadcast, since it is constant.
      fldj_theoretical = bijector_test_util.get_fldj_theoretical(
          bijector, x, event_ndims=1
      )
      self.assertAllClose(
          self.evaluate(fldj_theoretical),
          self.evaluate(fldj),
          atol=1e-5,
          rtol=1e-5,
      )
  
file_path: probability/tensorflow_probability/python/experimental/bijectors/highway_flow_test.py
imports: from absl.testing import parameterized
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import bijector_test_util
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.experimental.bijectors import highway_flow
  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 5 for calling the HighwayFlowTests obj's underlying member method code with example usage: obj.testTheoreticalFldj() without return types