code: def make_rank_polymorphic(fn, core_ndims, name=None):
      """Lift a function to one that vectorizes across arbitrary-rank inputs.
  
      Args:
        fn: Python `callable` `result = fn(*args)` where all arguments
          and the returned result(s) are (structures of) Tensors. Non-`Tensor`
          arguments may also be passed through by specifying a value of `None`
          in `core_ndims`.
        core_ndims: structure of `int` Tensors and/or `None` values, of the same
          structure as `args`. Each `int` contains the
          expected rank of the corresponding `Tensor` argument in an unvectorized
          call to `fn`; `None` values denote arguments that should not be vectorized
          (e.g., non-`Tensor` arguments). May alternately be a single scalar
          Tensor `int` applicable to all `args`.
        name: Python `str` name prefixed to ops created by `vectorized_fn`.
      Returns:
        vectorized_fn: a new function, equivalent to `fn`, but which automatically
          accepts arguments of any (combination of) ranks above the `core_ndims`,
          and returns a value with the broadcast batch shape of its arguments.
  
      #### Example
  
      ```python
      def add(a, b):
        return a + b
      add(tf.constant([1., 2.]), tf.constant(3.))
        # ==> Returns [4., 5.]
  
      # Naively passing a batch of three values for `b` raises an error.
      add(tf.constant([1., 2.]), tf.constant([3., 4., 5.]))
        # ==> Raises InvalidArgumentError: Incompatible shapes.
  
      # By annotating that `b` is scalar, we clarify how to batch the results.
      add_vector_to_scalar = make_rank_polymorphic(add, core_ndims=(1, 0))
      add_vector_to_scalar(tf.constant([1., 2.]), tf.constant([3., 4., 5.]))
        # ==> Returns [[4., 5.], [5., 6.], [6., 7.]]
      ```
  
      Lifted functions may accept non-`Tensor` args, denoted by `None` values in
      `core_ndims`. These values will be passed unmodified to the underlying
      function. For example, we could generalize `add` above to take an
      arbitrary binary operation, specified as a Python callable.
  
      ```python
      import operator
      def apply_binop(fn, a, b):
        return fn(a, b)
      apply_binop(operator.mul, tf.constant([1., 2.]), tf.constant(3.))
        # ==> Returns [3., 6.]
  
      # Batching, we pass the `fn` arg by specifying `core_ndims` of `None`.
      apply_binop_to_vector_and_scalar = make_rank_polymorphic(
        apply_binop, core_ndims=(None, 1, 0))
      apply_binop_to_vector_and_scalar(
        operator.mul, tf.constant([1., 2.]), tf.constant([3., 4., 5.]))
        # ==> Returns [[3., 6.], [4., 8.], [5., 10.]]
      ```
  
      #### Limitations
  
      The current form of this function has several limitations due to
      `vectorize_map`'s requirement that all inputs and outputs share a common batch
      dimension. When an automatically vectorized function is called, its inputs are
      padded up to the common batch shape, and all outputs are returned with the
      full broadcast batch shape. This can waste memory when a function has multiple
      outputs or multiple, differently-sized inputs. For example, if we define a
      simple function of two arguments:
  
      ```python
      def silly_increment(a, b):
        return a + 1., b + 1.
      vectorized_increment = make_rank_polymorphic(silly_increment, core_ndims=0)
      ```
  
      and call it with one very large argument (in this case, a billion entries):
  
      ```python
      a1, b1 = vectorized_increment(0., tf.ones([1000, 1000, 1000]))
      print(a1.shape, b1.shape)
        # ==> [1000, 1000, 1000], [1000, 1000, 1000]
      ```
  
      then both of the returned results will have a billion entries, even though
      the first result `a1` could have been computed as a scalar. In addition,
      another unnecessary billion-entry tensor will be created internally
      representing the vectorized input `a`. In this case, the vectorization was
      inefficient because not all outputs depended on all inputs; it would have
      been more efficient to exploit this by incrementing `a` and `b` separately,
      in different calls.
  
      """
  
      def vectorized_fn(*args):
          """Vectorized version of `fn` that accepts arguments of any rank."""
          with tf.name_scope(name or "make_rank_polymorphic"):
              # If we got a single value for core_ndims, tile it across all args.
              core_ndims_structure = (
                  core_ndims
                  if tf.nest.is_nested(core_ndims)
                  else tf.nest.map_structure(lambda _: core_ndims, args)
              )
  
              # Build flat lists of all argument parts and their corresponding core
              # ndims.
              flat_core_ndims = tf.nest.flatten(core_ndims_structure)
              flat_args = nest.flatten_up_to(
                  core_ndims_structure, args, check_types=False
              )
  
              # Filter to only the `Tensor`-valued args (taken to be those with `None`
              # values for `core_ndims`). Other args will be passed through to `fn`
              # unmodified.
              (
                  vectorized_arg_core_ndims,
                  vectorized_args,
                  fn_of_vectorized_args,
              ) = _lock_in_non_vectorized_args(
                  fn,
                  arg_structure=core_ndims_structure,
                  flat_core_ndims=flat_core_ndims,
                  flat_args=flat_args,
              )
  
              # `vectorized_map` requires all inputs to have a single, common batch
              # dimension `[n]`. So we broadcast all input parts to a common
              # batch shape, then flatten it down to a single dimension.
              vectorized_arg_shapes = [
                  ps.shape(arg) for arg in vectorized_args
              ]
  
              vectorized_arg_actual_core_ndims = []
              batch_shapes, core_shapes = [], []
              for arg_shape, core_nd in zip(
                  vectorized_arg_shapes, vectorized_arg_core_ndims
              ):
                  arg_nd = ps.rank_from_shape(arg_shape)
                  # Shrink 'core' ndims of rank-deficient args. This guarantees that
                  # `batch_ndims` is always nonnegative.
                  actual_core_nd = ps.minimum(arg_nd, core_nd)
                  vectorized_arg_actual_core_ndims.append(
                      actual_core_nd
                  )
  
                  batch_ndims = arg_nd - actual_core_nd
                  batch_ndims_ = tf.get_static_value(batch_ndims)
                  if batch_ndims_ is not None:
                      batch_shapes.append(arg_shape[:batch_ndims_])
                      core_shapes.append(arg_shape[batch_ndims_:])
                  else:
                      arg_shape = tf.convert_to_tensor(arg_shape)
                      batch_shapes.append(arg_shape[:batch_ndims])
                      core_shapes.append(arg_shape[batch_ndims:])
  
              # Flatten all of the batch dimensions into one.
              broadcast_batch_shape = functools.reduce(
                  ps.broadcast_shape, batch_shapes, []
              )
              n = ps.cast(
                  ps.reduce_prod(broadcast_batch_shape), tf.int32
              )
  
              static_n = tf.get_static_value(n)
              if static_n == 1:
                  # We can bypass `vectorized_map` if the batch shape is `[]`, `[1]`,
                  # `[1, 1]`, etc., just by flattening to batch shape `[]`.
                  result_batch_dims = 0
                  batched_result = fn_of_vectorized_args(
                      tf.nest.map_structure(
                          lambda x, nd: tf.reshape(
                              x, ps.shape(x)[ps.rank(x) - nd :]
                          ),
                          vectorized_args,
                          vectorized_arg_actual_core_ndims,
                          check_types=False,
                      )
                  )
              else:
                  # Pad all input parts to the common shape, then flatten
                  # into the single leading dimension `[n]`.
                  # TODO(b/145227909): If/when vmap supports broadcasting, use nested vmap
                  # when batch rank is static so that we can exploit broadcasting.
                  broadcast_vectorized_args = [
                      tf.broadcast_to(
                          part,
                          ps.concat(
                              [broadcast_batch_shape, core_shape],
                              axis=0,
                          ),
                      )
                      for (part, core_shape) in zip(
                          vectorized_args, core_shapes
                      )
                  ]
                  vectorized_args_with_flattened_batch_dim = [
                      tf.reshape(
                          part, ps.concat([[n], core_shape], axis=0)
                      )
                      for (part, core_shape) in zip(
                          broadcast_vectorized_args, core_shapes
                      )
                  ]
                  result_batch_dims = 1
                  batched_result = tf.vectorized_map(
                      fn_of_vectorized_args,
                      vectorized_args_with_flattened_batch_dim,
                  )
  
              # Unflatten any `Tensor`s in the result.
              unflatten = lambda x: tf.reshape(
                  x,
                  ps.concat(
                      [  # pylint: disable=g-long-lambda
                          broadcast_batch_shape,
                          ps.shape(x)[result_batch_dims:],
                      ],
                      axis=0,
                  ),
              )
              result = tf.nest.map_structure(
                  lambda x: unflatten(x) if tf.is_tensor(x) else x,
                  batched_result,
                  expand_composites=True,
              )
          return result
  
      return vectorized_fn
  
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 4 for calling a global function with example usage: make_rank_polymorphic(fn, core_ndims, name) and returns: vectorized_fn, result