Pads `input_tensor` to `fft_length` on its inner-most `fft_rank` dims.
(input_tensor, fft_rank, fft_length, is_reverse=False)
| 61 | |
| 62 | |
| 63 | def _maybe_pad_for_rfft(input_tensor, fft_rank, fft_length, is_reverse=False): |
| 64 | """Pads `input_tensor` to `fft_length` on its inner-most `fft_rank` dims.""" |
| 65 | fft_shape = _tensor_util.constant_value_as_shape(fft_length) |
| 66 | |
| 67 | # Edge case: skip padding empty tensors. |
| 68 | if (input_tensor.shape.ndims is not None and |
| 69 | any(dim.value == 0 for dim in input_tensor.shape.dims)): |
| 70 | return input_tensor |
| 71 | |
| 72 | # If we know the shapes ahead of time, we can either skip or pre-compute the |
| 73 | # appropriate paddings. Otherwise, fall back to computing paddings in |
| 74 | # TensorFlow. |
| 75 | if fft_shape.is_fully_defined() and input_tensor.shape.ndims is not None: |
| 76 | # Slice the last FFT-rank dimensions from input_tensor's shape. |
| 77 | input_fft_shape = input_tensor.shape[-fft_shape.ndims:] |
| 78 | |
| 79 | if input_fft_shape.is_fully_defined(): |
| 80 | # In reverse, we only pad the inner-most dimension to fft_length / 2 + 1. |
| 81 | if is_reverse: |
| 82 | fft_shape = fft_shape[:-1].concatenate( |
| 83 | fft_shape.dims[-1].value // 2 + 1) |
| 84 | |
| 85 | paddings = [[0, max(fft_dim.value - input_dim.value, 0)] |
| 86 | for fft_dim, input_dim in zip( |
| 87 | fft_shape.dims, input_fft_shape.dims)] |
| 88 | if any(pad > 0 for _, pad in paddings): |
| 89 | outer_paddings = [[0, 0]] * max((input_tensor.shape.ndims - |
| 90 | fft_shape.ndims), 0) |
| 91 | return _array_ops.pad(input_tensor, outer_paddings + paddings) |
| 92 | return input_tensor |
| 93 | |
| 94 | # If we can't determine the paddings ahead of time, then we have to pad. If |
| 95 | # the paddings end up as zero, tf.pad has a special-case that does no work. |
| 96 | input_rank = _array_ops.rank(input_tensor) |
| 97 | input_fft_shape = _array_ops.shape(input_tensor)[-fft_rank:] |
| 98 | outer_dims = _math_ops.maximum(0, input_rank - fft_rank) |
| 99 | outer_paddings = _array_ops.zeros([outer_dims], fft_length.dtype) |
| 100 | # In reverse, we only pad the inner-most dimension to fft_length / 2 + 1. |
| 101 | if is_reverse: |
| 102 | fft_length = _array_ops.concat([fft_length[:-1], |
| 103 | fft_length[-1:] // 2 + 1], 0) |
| 104 | fft_paddings = _math_ops.maximum(0, fft_length - input_fft_shape) |
| 105 | paddings = _array_ops.concat([outer_paddings, fft_paddings], 0) |
| 106 | paddings = _array_ops.stack([_array_ops.zeros_like(paddings), paddings], |
| 107 | axis=1) |
| 108 | return _array_ops.pad(input_tensor, paddings) |
| 109 | |
| 110 | |
| 111 | def _rfft_wrapper(fft_fn, fft_rank, default_name): |