tensorflow_probability/python/math/interpolation.py
def _batch_gather_with_broadcast(params, indices, axis):
  """Like batch_gather, but broadcasts to the left of axis."""
  # batch_gather assumes...
  #   params.shape =  [A1,...,AN, B1,...,BM]
  #   indices.shape = [A1,...,AN, C]
  # which gives output of shape
  #                   [A1,...,AN, C, B1,...,BM]
  # Here we broadcast dims of each to the left of `axis` in params, and left of
  # the rightmost dim in indices, e.g. we can
  # have
  #   params.shape =  [A1,...,AN, B1,...,BM]
  #   indices.shape = [a1,...,aN, C],
  # where ai broadcasts with Ai.

  # leading_bcast_shape is the broadcast of [A1,...,AN] and [a1,...,aN].
  leading_bcast_shape = tf.broadcast_dynamic_shape(
      tf.shape(input=params)[:axis],
      tf.shape(input=indices)[:-1])
  params += tf.zeros(
      tf.concat((leading_bcast_shape, tf.shape(input=params)[axis:]), axis=0),
      dtype=params.dtype)
  indices += tf.zeros(
      tf.concat((leading_bcast_shape, tf.shape(input=indices)[-1:]), axis=0),
      dtype=indices.dtype)
  return tf.compat.v1.batch_gather(params, indices)
def _batch_gather_with_broadcast(params, indices, axis):
  """Like batch_gather, but broadcasts to the left of axis."""
  # batch_gather assumes...
  #   params.shape =  [A1,...,AN, B1,...,BM]
  #   indices.shape = [A1,...,AN, C]
  # which gives output of shape
  #                   [A1,...,AN, C, B1,...,BM]
  # Here we broadcast dims of each to the left of `axis` in params, and left of
  # the rightmost dim in indices, e.g. we can
  # have
  #   params.shape =  [A1,...,AN, B1,...,BM]
  #   indices.shape = [a1,...,aN, C],
  # where ai broadcasts with Ai.

  # leading_bcast_shape is the broadcast of [A1,...,AN] and [a1,...,aN].
  leading_bcast_shape = tf.broadcast_dynamic_shape(
      tf.shape(input=params)[:axis],
      tf.shape(input=indices)[:-1])
  params += tf.zeros(
      tf.concat((leading_bcast_shape, tf.shape(input=params)[axis:]), axis=0),
      dtype=params.dtype)
  indices += tf.zeros(
      tf.concat((leading_bcast_shape, tf.shape(input=indices)[-1:]), axis=0),
      dtype=indices.dtype)
  return tf.compat.v1.batch_gather(params, indices)
Like batch_gather, but broadcasts to the left of axis.