code: def test_init_supports_arg_or_kwarg_seed(self):
      seed = test_util.test_seed(sampler_type="stateless")
      init_fn, _ = trainable_state_util.as_stateless_builder(
          seed_generator
      )()
      self.assertLen(init_fn(seed=seed), 5)
      # Check that we can invoke init_fn with an arg or kwarg seed,
      # regardless of how the inner functions are parameterized.
      self.assertAllCloseNested(init_fn(seed), init_fn(seed=seed))
  
      if not JAX_MODE:
          # Check that we can initialize with no seed.
          self.assertLen(init_fn(), 5)
  
imports: import functools
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import identity
  from tensorflow_probability.python.bijectors import joint_map
  from tensorflow_probability.python.bijectors import softplus
  from tensorflow_probability.python.distributions import deterministic
  from tensorflow_probability.python.distributions import joint_distribution_auto_batched
  from tensorflow_probability.python.distributions import joint_distribution_sequential
  from tensorflow_probability.python.distributions import lognormal
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.internal import tf_keras
  from tensorflow_probability.python.internal import trainable_state_util
  from tensorflow_probability.python.math import gradient
  from tensorflow_probability.python.math.minimize import minimize
  from tensorflow_probability.python.math.minimize import minimize_stateless
label_desc: Use image node_id 1 for calling the TestWrapGeneratorAsStateless obj's underlying member method code with example usage: obj.test_init_supports_arg_or_kwarg_seed() without return types