code: def test_functional_api(self):
      # Create a trainable distribution using the functional API.
      dummy_input = tf_keras.Input(shape=())
      x = variable_input.VariableLayer(
          shape=[2, 3, 4],
          dtype=tf.float64,
          trainable=False,  # You'd probably never want this in IRL.
      )(dummy_input)
      # The Dense serves no real purpose; it will change the event_shape.
      x = tf_keras.layers.Dense(5, use_bias=False, dtype=tf.float64)(x)
      x = distribution_layer.DistributionLambda(
          lambda t: independent.Independent(
              normal.Normal(
                  loc=t[0], scale=t[1]
              ),  # pylint: disable=g-long-lambda
              reinterpreted_batch_ndims=1,
          ),
          dtype=tf.float64,
      )(x)
      model = tf_keras.Model(dummy_input, x)
  
      # Instantiate the model (as a TFP distribution).
      dist = model(tf.zeros([]))
  
      # Check the weights.
      self.assertEqual(2, len(model.weights))
  
      # Check the VariableLayer layer.
      self.assertIs(tf.float64, tf.as_dtype(model.weights[0].dtype))
      self.assertEqual((2, 3, 4), model.layers[1].weights[0].shape)
      self.assertFalse(model.layers[1].trainable)
      self.assertFalse(model.layers[1].weights[0].trainable)
  
      # Check the Dense layer.
      self.assertIs(tf.float64, tf.as_dtype(model.weights[1].dtype))
      self.assertEqual((4, 5), model.layers[2].weights[0].shape)
      self.assertTrue(model.layers[2].trainable)
      self.assertTrue(model.layers[2].weights[0].trainable)
  
      # Check the distribution.
      self.assertIsInstance(
          dist.tensor_distribution, independent.Independent
      )
      self.assertIs(tf.float64, dist.dtype)
      self.assertEqual((3,), dist.batch_shape)
      self.assertEqual((5,), dist.event_shape)
  
file_path: probability/tensorflow_probability/python/layers/variable_input_test.py
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.internal import tf_keras
  from tensorflow_probability.python.layers import distribution_layer
  from tensorflow_probability.python.layers import variable_input
label_desc: Use image node_id 2 for calling the VariableInputLayerTest obj's underlying member method code with example usage: obj.test_functional_api() without return types