code: def testBijectorIntegration(self):
      @jdc.JointDistributionCoroutine
      def bijector_model():
          bijectors = []
          for i in range(3):
              bijector_net = yield Root(
                  jdlayers.Sequential(
                      jdlayers.Affine(5, 5),
                      tf.tanh,
                      jdlayers.Affine(10, 5),
                      jdlayers.Lambda(
                          lambda x: chain.Chain(  # pylint: disable=g-long-lambda
                              [
                                  shift.Shift(x[..., :5]),
                                  scale.Scale(log_scale=x[..., 5:]),
                              ]
                          )
                      ),
                      name=f"net{i}",
                  )
              )
  
              bijectors.append(
                  real_nvp.RealNVP(
                      fraction_masked=0.5 * (-1) ** i,
                      bijector_fn=lambda x, _, bn=bijector_net: bn(x),
                  )
              )
  
          yield jdlayers.Lambda(lambda: chain.Chain(bijectors))
  
      *params, _ = bijector_model.sample(
          2, seed=test_util.test_seed(sampler_type="stateless")
      )
  
      def ldj_fn(params):
          bijector_fn = bijector_model.sample(
              value=params,
              seed=test_util.test_seed(sampler_type="stateless"),
          )[-1]
          return bijector_fn().forward_log_det_jacobian(
              tf.ones([2, 10]), 1
          )
  
      ldj, (params_grad,) = gradient.value_and_gradient(
          ldj_fn, (params,)
      )
      self.assertEqual([2], ldj.shape)
      self.assertAllAssertsNested(
          lambda x, g: self.assertTrue(
              x.shape[-1] == 0
              or custom_gradient.is_valid_gradient(  # pylint: disable=g-long-lambda
                  g
              )
          ),
          params,
          params_grad,
      )
  
imports: from absl.testing import parameterized
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import chain
  from tensorflow_probability.python.bijectors import real_nvp
  from tensorflow_probability.python.bijectors import scale
  from tensorflow_probability.python.bijectors import shift
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.distributions import joint_distribution_coroutine
  from tensorflow_probability.python.distributions import lognormal
  from tensorflow_probability.python.experimental.joint_distribution_layers import layers
  from tensorflow_probability.python.internal import custom_gradient
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import gradient
label_desc: Use image node_id 1 for calling the _JDLayersTestBase obj's underlying member method code with example usage: obj.testBijectorIntegration() and returns: bijector_fn