code: def fold_in_axis_index(seed, axis_name=None):
      """Folds the active axis index into a seed according to its axis name."""
      if axis_name is None:
          return seed
      nest.assert_shallow_structure(seed, axis_name)
      axis_names = nest.map_structure_up_to(
          seed, canonicalize_named_axis, axis_name
      )
  
      def fold_in(seed, axes):
          for name in axes:
              axis_index = get_axis_index(name)
              seed = samplers.fold_in(
                  seed, tf.cast(axis_index, tf.int32)
              )
          return seed
  
      return nest.map_structure_up_to(seed, fold_in, seed, axis_names)
  
imports: import collections
  import functools
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import custom_gradient
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow.python.util import nest
label_desc: Use image node_id 13 for calling a global function with example usage: fold_in_axis_index(seed, axis_name) and returns: nest, seed, seed