Function _pad_all_input

tensorflow/python/tpu/tpu.py:642–739

Pad all input tensors given padded_shapes. The real shape tensors will be concatenated with the padded original inputs. Args: inputs: The original inputs. padded_shapes: A list of padded shapes for each input. Returns: The padded inputs and a PaddingMap list which maps the padde

(inputs, padded_shapes)

def _pad_all_input(inputs, padded_shapes):
  """Pad all input tensors given padded_shapes.

  The real shape tensors will be concatenated with the padded original inputs.

  Args:
    inputs: The original inputs.
    padded_shapes: A list of padded shapes for each input.

  Returns:
    The padded inputs and a PaddingMap list which maps the padded input
    dimension to the real shape argument index.
  """
  # maximum_static_shapes[idx][i] indicates the maximum static size of ith
  # dimension of the idx input among all the replicas.
  maximum_static_shapes = []
  # need_padding[idx][i] indicates whether the ith dimension of the idx input
  # needs padding.
  need_padding = []
  input_shape_tensors = []
  for core_idx, inputs_per_core in enumerate(inputs):
    for idx, input_tensor in enumerate(inputs_per_core):
      input_shape = input_tensor.get_shape().as_list()
      if core_idx == 0:
        input_shape_tensors.append([])
        maximum_static_shapes.append(input_shape)
        need_padding.append(np.full_like(input_shape, False, dtype=bool))
      else:
        for i, s in enumerate(input_shape):
          if not s or s != maximum_static_shapes[idx][i]:
            need_padding[idx][i] = True
        maximum_static_shapes[idx] = max(input_shape,
                                         maximum_static_shapes[idx])
      input_shape_tensors[idx].append(array_ops.shape(input_tensor))

  maximum_shapes = []
  for shapes_per_input in input_shape_tensors:
    maximum_shapes.append(
        math_ops.reduce_max(array_ops.stack(shapes_per_input), axis=0))

  padded_inputs = []
  real_shapes = []
  padding_maps = []
  for core_idx, inputs_per_core in enumerate(inputs):
    padded_inputs.append([])
    real_shapes.append([])
    real_shape_idx = len(inputs_per_core) - 1
    for idx, input_tensor in enumerate(inputs_per_core):
      input_shape_tensor = input_shape_tensors[idx][core_idx]
      input_shape = input_tensor.get_shape().as_list()
      padded_shape = padded_shapes[idx]

      if any(need_padding[idx]):
        for i, s in enumerate(input_shape):
          if need_padding[idx][i]:
            if core_idx == 0:
              real_shape_idx += 1
              padding_map = dynamic_padding.PaddingMap()