tensorflow_probability/python/internal/distribution_util.py
def pad(x, axis, front=False, back=False, value=0, count=1, name=None):
  """Pads `value` to the front and/or back of a `Tensor` dim, `count` times.

  Args:
    x: `Tensor` input.
    axis: Scalar `int`-like `Tensor` representing the single dimension to pad.
      (Negative indexing is supported.)
    front: Python `bool`; if `True` the beginning of the `axis` dimension is
      padded with `value`, `count` times. If `False` no front padding is made.
    back: Python `bool`; if `True` the end of the `axis` dimension is padded
      with `value`, `count` times. If `False` no end padding is made.
    value: Scalar `int`-like `Tensor` representing the actual value added to the
      front and/or back of the `axis` dimension of `x`.
    count: Scalar `int`-like `Tensor` representing number of elements added to
      the front and/or back of the `axis` dimension of `x`. E.g., if `front =
      back = True` then `2 * count` elements are added.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    pad: The padded version of input `x`.

  Raises:
    ValueError: if both `front` and `back` are `False`.
    TypeError: if `count` is not `int`-like.
  """
  with tf.name_scope(name or "pad"):
    x = tf.convert_to_tensor(value=x, name="x")
    value = tf.convert_to_tensor(value=value, dtype=x.dtype, name="value")
    count = tf.convert_to_tensor(value=count, name="count")
    if not dtype_util.is_integer(count.dtype):
      raise TypeError("`count.dtype` (`{}`) must be `int`-like.".format(
          dtype_util.name(count.dtype)))
    if not front and not back:
      raise ValueError("At least one of `front`, `back` must be `True`.")
    ndims = (
        tensorshape_util.rank(x.shape)
        if tensorshape_util.rank(x.shape) is not None else tf.rank(
            x, name="ndims"))
    axis = tf.convert_to_tensor(value=axis, name="axis")
    axis_ = tf.get_static_value(axis)
    if axis_ is not None:
      axis = axis_
      if axis < 0:
        axis = ndims + axis
      count_ = tf.get_static_value(count)
      if axis_ >= 0 or tensorshape_util.rank(x.shape) is not None:
        head = x.shape[:axis]
        mid_dim_value = tf.compat.dimension_value(x.shape[axis])
        if count_ is None or mid_dim_value is None:
          middle = tf.TensorShape(None)
        else:
          middle = tf.TensorShape(mid_dim_value + count_ * (front + back))
        tail = x.shape[axis + 1:]
        final_shape = head.concatenate(middle.concatenate(tail))
      else:
        final_shape = None
    else:
      axis = tf.where(axis < 0, ndims + axis, axis)
      final_shape = None
    x = tf.pad(
        tensor=x,
        paddings=tf.one_hot(
            indices=tf.stack([axis if front else -1, axis if back else -1]),
            depth=ndims,
            axis=0,
            on_value=count,
            dtype=tf.int32),
        constant_values=value)
    if final_shape is not None:
      tensorshape_util.set_shape(x, final_shape)
    return x
def pad(x, axis, front=False, back=False, value=0, count=1, name=None):
  """Pads `value` to the front and/or back of a `Tensor` dim, `count` times.

  Args:
    x: `Tensor` input.
    axis: Scalar `int`-like `Tensor` representing the single dimension to pad.
      (Negative indexing is supported.)
    front: Python `bool`; if `True` the beginning of the `axis` dimension is
      padded with `value`, `count` times. If `False` no front padding is made.
    back: Python `bool`; if `True` the end of the `axis` dimension is padded
      with `value`, `count` times. If `False` no end padding is made.
    value: Scalar `int`-like `Tensor` representing the actual value added to the
      front and/or back of the `axis` dimension of `x`.
    count: Scalar `int`-like `Tensor` representing number of elements added to
      the front and/or back of the `axis` dimension of `x`. E.g., if `front =
      back = True` then `2 * count` elements are added.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    pad: The padded version of input `x`.

  Raises:
    ValueError: if both `front` and `back` are `False`.
    TypeError: if `count` is not `int`-like.
  """
  with tf.name_scope(name or "pad"):
    x = tf.convert_to_tensor(value=x, name="x")
    value = tf.convert_to_tensor(value=value, dtype=x.dtype, name="value")
    count = tf.convert_to_tensor(value=count, name="count")
    if not dtype_util.is_integer(count.dtype):
      raise TypeError("`count.dtype` (`{}`) must be `int`-like.".format(
          dtype_util.name(count.dtype)))
    if not front and not back:
      raise ValueError("At least one of `front`, `back` must be `True`.")
    ndims = (
        tensorshape_util.rank(x.shape)
        if tensorshape_util.rank(x.shape) is not None else tf.rank(
            x, name="ndims"))
    axis = tf.convert_to_tensor(value=axis, name="axis")
    axis_ = tf.get_static_value(axis)
    if axis_ is not None:
      axis = axis_
      if axis < 0:
        axis = ndims + axis
      count_ = tf.get_static_value(count)
      if axis_ >= 0 or tensorshape_util.rank(x.shape) is not None:
        head = x.shape[:axis]
        mid_dim_value = tf.compat.dimension_value(x.shape[axis])
        if count_ is None or mid_dim_value is None:
          middle = tf.TensorShape(None)
        else:
          middle = tf.TensorShape(mid_dim_value + count_ * (front + back))
        tail = x.shape[axis + 1:]
        final_shape = head.concatenate(middle.concatenate(tail))
      else:
        final_shape = None
    else:
      axis = tf.where(axis < 0, ndims + axis, axis)
      final_shape = None
    x = tf.pad(
        tensor=x,
        paddings=tf.one_hot(
            indices=tf.stack([axis if front else -1, axis if back else -1]),
            depth=ndims,
            axis=0,
            on_value=count,
            dtype=tf.int32),
        constant_values=value)
    if final_shape is not None:
      tensorshape_util.set_shape(x, final_shape)
    return x
def pad ( x , axis , front = False , back = False , value = 0 , count = 1 , name = None ) : with tf . name_scope ( name or "pad" ) : x = tf . convert_to_tensor ( value = x , name = "x" ) value = tf . convert_to_tensor ( value = value , dtype = x . dtype , name = "value" ) count = tf . convert_to_tensor ( value = count , name = "count" ) if not dtype_util . is_integer ( count . dtype ) : raise TypeError ( "`count.dtype` (`{}`) must be `int`-like." . format ( dtype_util . name ( count . dtype ) ) ) if not front and not back : raise ValueError ( "At least one of `front`, `back` must be `True`." ) ndims = ( tensorshape_util . rank ( x . shape ) if tensorshape_util . rank ( x . shape ) is not None else tf . rank ( x , name = "ndims" ) ) axis = tf . convert_to_tensor ( value = axis , name = "axis" ) axis_ = tf . get_static_value ( axis ) if axis_ is not None : axis = axis_ if axis < 0 : axis = ndims + axis count_ = tf . get_static_value ( count ) if axis_ >= 0 or tensorshape_util . rank ( x . shape ) is not None : head = x . shape [ : axis ] mid_dim_value = tf . compat . dimension_value ( x . shape [ axis ] ) if count_ is None or mid_dim_value is None : middle = tf . TensorShape ( None ) else : middle = tf . TensorShape ( mid_dim_value + count_ * ( front + back ) ) tail = x . shape [ axis + 1 : ] final_shape = head . concatenate ( middle . concatenate ( tail ) ) else : final_shape = None else : axis = tf . where ( axis < 0 , ndims + axis , axis ) final_shape = None x = tf . pad ( tensor = x , paddings = tf . one_hot ( indices = tf . stack ( [ axis if front else - 1 , axis if back else - 1 ] ) , depth = ndims , axis = 0 , on_value = count , dtype = tf . int32 ) , constant_values = value ) if final_shape is not None : tensorshape_util . set_shape ( x , final_shape ) return x
Pads `value` to the front and/or back of a `Tensor` dim, `count` times.

  Args:
    x: `Tensor` input.
    axis: Scalar `int`-like `Tensor` representing the single dimension to pad.
      (Negative indexing is supported.)
    front: Python `bool`; if `True` the beginning of the `axis` dimension is
      padded with `value`, `count` times. If `False` no front padding is made.
    back: Python `bool`; if `True` the end of the `axis` dimension is padded
      with `value`, `count` times. If `False` no end padding is made.
    value: Scalar `int`-like `Tensor` representing the actual value added to the
      front and/or back of the `axis` dimension of `x`.
    count: Scalar `int`-like `Tensor` representing number of elements added to
      the front and/or back of the `axis` dimension of `x`. E.g., if `front =
      back = True` then `2 * count` elements are added.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    pad: The padded version of input `x`.

  Raises:
    ValueError: if both `front` and `back` are `False`.
    TypeError: if `count` is not `int`-like.
Pads value to the front and / or back of a Tensor dim count times .