(tensor_or_dict, signals)
return nest.map_structure(slice_fn, tensor_or_dict)
Slice the real Tensors according to padding mask in signals.
Slice the real Tensors according to padding mask in signals.
Slice the real Tensors according to padding mask in signals .
def slice_tensor_or_dict(tensor_or_dict, signals):
    """Slice the real Tensors according to padding mask in signals."""

    padding_mask = signals['padding_mask']
    batch_size = array_ops.shape(padding_mask)[0]

    def verify_batch_size(tensor):
      check_batch_size = math_ops.equal(batch_size, tensor.shape[0])
      with ops.control_dependencies([check_batch_size]):
        return array_ops.identity(tensor)

    def slice_single_tensor(tensor):
      rank = len(tensor.shape)
      assert rank > 0
      real_batch_size = batch_size - math_ops.reduce_sum(padding_mask)
      return verify_batch_size(tensor)[0:real_batch_size]

    # As we split the Tensors to all TPU cores and concat them back, it is
    # important to ensure the real data is placed before padded ones, i.e.,
    # order is preserved. By that, the sliced padding mask should have all 0's.
    # If this assertion failed, # the slice logic here would not hold.
    sliced_padding_mask = slice_single_tensor(padding_mask)
    assert_padding_mask = math_ops.equal(
        math_ops.reduce_sum(sliced_padding_mask), 0)

    with ops.control_dependencies([assert_padding_mask]):
      should_stop = _StopSignals.should_stop(
          _StopSignals.as_scalar_stopping_signal(signals))

    is_full_batch = math_ops.equal(math_ops.reduce_sum(padding_mask), 0)

    def slice_fn(tensor):
      # If the current batch is full batch or part of stopping signals, we do
      # not need to slice to save performance.
      return control_flow_ops.cond(
          math_ops.logical_or(should_stop, is_full_batch),
          (lambda: verify_batch_size(tensor)),
          (lambda: slice_single_tensor(tensor)))

    return nest.map_structure(slice_fn, tensor_or_dict)
def slice_tensor_or_dict ( tensor_or_dict , signals ) : padding_mask = signals [ 'padding_mask' ] batch_size = array_ops . shape ( padding_mask ) [ 0 ] def verify_batch_size ( tensor ) : check_batch_size = math_ops . equal ( batch_size , tensor . shape [ 0 ] ) with ops . control_dependencies ( [ check_batch_size ] ) : return array_ops . identity ( tensor ) def slice_single_tensor ( tensor ) : rank = len ( tensor . shape ) assert rank > 0 real_batch_size = batch_size - math_ops . reduce_sum ( padding_mask ) return verify_batch_size ( tensor ) [ 0 : real_batch_size ] # As we split the Tensors to all TPU cores and concat them back, it is # important to ensure the real data is placed before padded ones, i.e., # order is preserved. By that, the sliced padding mask should have all 0's. # If this assertion failed, # the slice logic here would not hold. sliced_padding_mask = slice_single_tensor ( padding_mask ) assert_padding_mask = math_ops . equal ( math_ops . reduce_sum ( sliced_padding_mask ) , 0 ) with ops . control_dependencies ( [ assert_padding_mask ] ) : should_stop = _StopSignals . should_stop ( _StopSignals . as_scalar_stopping_signal ( signals ) ) is_full_batch = math_ops . equal ( math_ops . reduce_sum ( padding_mask ) , 0 ) def slice_fn ( tensor ) : # If the current batch is full batch or part of stopping signals, we do # not need to slice to save performance. return control_flow_ops . cond ( math_ops . logical_or ( should_stop , is_full_batch ) , ( lambda : verify_batch_size ( tensor ) ) , ( lambda : slice_single_tensor ( tensor ) ) ) return nest . map_structure ( slice_fn , tensor_or_dict )