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):
  # 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 ) : # 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.
Like batch_gather but broadcasts to the left of axis .