code: def testStatefulDistFn(self):
      class StatefulNormal(distribution.Distribution):
          def __init__(self, loc):
              self._loc = tf.convert_to_tensor(loc)
              super(StatefulNormal, self).__init__(
                  dtype=tf.float32,
                  reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,
                  validate_args=False,
                  allow_nan_stats=False,
              )
  
          def _batch_shape(self):
              return self._loc.shape
  
          def _event_shape(self):
              return []
  
          def _sample_n(self, n, seed=None):
              return self._loc + tf.random.normal(
                  tf.concat([[n], tf.shape(self._loc)], axis=0),
                  seed=seed,
              )
  
      def dist_fn(s):
          return StatefulNormal(loc=s)
  
      ar = autoregressive.Autoregressive(
          dist_fn,
          sample0=normal.Normal(0.0, 1.0).sample(
              7, seed=test_util.test_seed()
          ),
          num_steps=7,
      )
  
      with warnings.catch_warnings(record=True) as triggered:
          self.evaluate(ar.sample(seed=test_util.test_seed()))
      self.assertTrue(
          any(
              "Falling back to stateful sampling for `distribution_fn(sample0)`"
              in str(warning.message)
              for warning in triggered
          )
      )
  
      num_steps = tf.Variable(9)
      self.evaluate(num_steps.initializer)
      with warnings.catch_warnings(record=True) as triggered:
          self.evaluate(
              ar.copy(num_steps=num_steps).sample(
                  seed=test_util.test_seed()
              )
          )
      self.assertTrue(
          any(
              "Falling back to stateful sampling for `distribution_fn(sample0)`"
              in str(warning.message)
              for warning in triggered
          )
      )
  
imports: import warnings
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import masked_autoregressive
  from tensorflow_probability.python.bijectors import scale_matvec_tril
  from tensorflow_probability.python.distributions import autoregressive
  from tensorflow_probability.python.distributions import categorical
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.internal import reparameterization
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import linalg
label_desc: Use image node_id 1 for calling the SamplerBackwardCompatibilityTest obj's underlying member method code with example usage: obj.testStatefulDistFn() and returns: StatefulNormal, self, list, unknown