def pad_along_dimension()

in tensorflow_text/python/ops/pad_along_dimension_op.py [0:0]

• 27 lines of code
• 15 McCabe index (conditional complexity)

def pad_along_dimension(data, axis=-1, left_pad=None, right_pad=None,
                        name=None):
  """Add padding to the beginning and end of data in a specific dimension.

  Returns a tensor constructed from `data`, where each row in dimension `axis`
  is replaced by the concatenation of the left padding followed by the row
  followed by the right padding.  I.e., if `L=left_pad.shape[0]` and
  `R=right_pad.shape[0]`, then:

  ```python
  result[i1...iaxis, 0:L] = left_pad
  result[i1...iaxis, L:-R] = data[i0...iaxis]
  result[i1...iaxis, -R:] = right_pad
  ```

  Args:
    data: `<dtype>[O1...ON, A, I1...IM]` A potentially ragged `K` dimensional
      tensor with outer dimensions of size `O1...ON`; axis dimension of size
      `A`; and inner dimensions of size `I1...IM`.  I.e. `K = N + 1 + M`, where
      `N>=0` and `M>=0`.
    axis: An integer constant specifying the axis along which padding is added.
      Negative axis values from `-K` to `-1` are supported.
    left_pad: `<dtype>[L, I1...IM]` An `M+1` dimensional tensor that should be
      prepended to each row along dimension `axis`; or `None` if no padding
      should be added to the left side.
    right_pad: `<dtype>[R, I1...IM]` An `M+1` dimensional tensor that should be
      appended to each row along dimension `axis`; or `None` if no padding
      should be added to the right side.
    name: The name of this op (optional).

  Returns:
    `<dtype>[O1...ON, L + A + R, I1...IM]`
      A potentially ragged `K` dimensional tensor with outer dimensions of size
      `O1...ON`; padded axis dimension size `L+A+R`; and inner dimensions of
      size `I1...IM`.  If `data` is a `RaggedTensor`, then the returned tensor
      is a `RaggedTensor` with the same `ragged_rank`.
  """
  data = ragged_tensor.convert_to_tensor_or_ragged_tensor(data, name="data")

  if not isinstance(axis, int):
    raise TypeError("axis must be an int; got %s" % type(axis).__name__)

  if left_pad is None and right_pad is None:
    return data

  with ops.name_scope(name, "PadAlongDimension", [data]):
    if data.shape.ndims is not None and (axis < -data.shape.ndims or
                                         axis >= data.shape.ndims):
      raise errors.InvalidArgumentError(
          None, None, "axis must be between -k <= axis <= -1 OR 0 <= axis < k")
    if isinstance(data, ragged_tensor.RaggedTensor):
      axis = _get_positive_axis(axis, data.shape.ndims)

    if left_pad is not None:
      left_pad = ragged_tensor.convert_to_tensor_or_ragged_tensor(
          left_pad, dtype=data.dtype, name="left_pad")
    if right_pad is not None:
      right_pad = ragged_tensor.convert_to_tensor_or_ragged_tensor(
          right_pad, dtype=data.dtype, name="left_pad")

    left_padding = _padding_for_dimension(data, axis, left_pad)
    right_padding = _padding_for_dimension(data, axis, right_pad)

    pieces = [left_padding, data, right_padding]
    if isinstance(data, ragged_tensor.RaggedTensor):
      return array_ops.concat([p for p in pieces if p is not None], axis)
    else:
      return array_ops.concat([p for p in pieces if p is not None], axis)