code: def batchify_op(op, op_min_input_ndims, x, *other_op_args):
      """Reshape `op` input `x` to be a vec of `op_min_input_ndims`-rank tensors."""
      if x.shape.rank == op_min_input_ndims + 1:
          # Input is already a vector of `op_min_input_ndims`-rank tensors.
          return op(x, *other_op_args)
      batch_shape, op_shape = ps.split(
          ps.shape(x), num_or_size_splits=[-1, op_min_input_ndims]
      )
      flat_shape = ps.pad(
          op_shape, paddings=[[1, 0]], constant_values=-1
      )
      y = tf.reshape(x, flat_shape)
      y = op(y, *other_op_args)
      unflat_shape = ps.concat(
          [
              batch_shape,
              ps.shape(y)[1:],
          ],
          axis=0,
      )
      y = tf.reshape(y, unflat_shape)
      return y
  
imports: import collections
  import contextlib
  import functools
  import sys
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
label_desc: Use image node_id 14 for calling a global function with example usage: batchify_op(op, op_min_input_ndims, x) and returns: y, op