code: def _lock_in_non_vectorized_args(
      fn, arg_structure, flat_core_ndims, flat_args
  ):
      """Wraps `fn` to take only those args with non-`None` core ndims."""
  
      # Extract the indices and values of args where core_ndims is not `None`.
      (
          vectorized_arg_indices,
          vectorized_arg_core_ndims,
          vectorized_args,
      ) = ([], [], [])
      if any(nd is not None for nd in flat_core_ndims):
          (
              vectorized_arg_indices,
              vectorized_arg_core_ndims,
              vectorized_args,
          ) = zip(
              *[
                  (i, nd, tf.convert_to_tensor(t))
                  for i, (nd, t) in enumerate(
                      zip(flat_core_ndims, flat_args)
                  )
                  if nd is not None
              ]
          )
  
      vectorized_arg_index_set = set(vectorized_arg_indices)
  
      def fn_of_vectorized_args(vectorized_args):
          with tf.name_scope("fn_of_vectorized_args"):
              vectorized_args_by_index = dict(
                  zip(vectorized_arg_indices, vectorized_args)
              )
              # Substitute the vectorized args into the original argument
              # structure.
              new_args_with_original_structure = (
                  tf.nest.pack_sequence_as(
                      arg_structure,
                      [
                          vectorized_args_by_index[i]
                          if i in vectorized_arg_index_set
                          else v
                          for (i, v) in enumerate(flat_args)
                      ],
                  )
              )
              return tf.nest.map_structure(
                  # If `fn` returns any Distribution instances, ensure that their
                  # parameter batch shapes are padded to align after vectorization.
                  _maybe_rectify_parameter_shapes,
                  fn(*new_args_with_original_structure),
              )
  
      return (
          vectorized_arg_core_ndims,
          vectorized_args,
          fn_of_vectorized_args,
      )
  
file_path: probability/tensorflow_probability/python/internal/vectorization_util.py
imports: import functools
  import warnings
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.util import SeedStream
  from tensorflow.python.ops.parallel_for import control_flow_ops
  from tensorflow.python.util import nest
label_desc: Use image node_id 2 for calling a global function with example usage: _lock_in_non_vectorized_args(fn, arg_structure, flat_core_ndims, flat_args) and returns: vectorized_arg_core_ndims,  vectorized_args,  fn_of_vectorized_args, tf