A gradient function for RFFT with the provided `rank` and `irfft_fn`.
(op, grad)
| 221 | assert rank in (1, 2), "Gradient for RFFT3D is not implemented." |
| 222 | |
| 223 | def _grad(op, grad): |
| 224 | """A gradient function for RFFT with the provided `rank` and `irfft_fn`.""" |
| 225 | fft_length = op.inputs[1] |
| 226 | input_shape = _array_ops.shape(op.inputs[0]) |
| 227 | is_even = _math_ops.cast(1 - (fft_length[-1] % 2), _dtypes.complex64) |
| 228 | |
| 229 | def _tile_for_broadcasting(matrix, t): |
| 230 | expanded = _array_ops.reshape( |
| 231 | matrix, |
| 232 | _array_ops.concat([ |
| 233 | _array_ops.ones([_array_ops.rank(t) - 2], _dtypes.int32), |
| 234 | _array_ops.shape(matrix) |
| 235 | ], 0)) |
| 236 | return _array_ops.tile( |
| 237 | expanded, _array_ops.concat([_array_ops.shape(t)[:-2], [1, 1]], 0)) |
| 238 | |
| 239 | def _mask_matrix(length): |
| 240 | """Computes t_n = exp(sqrt(-1) * pi * n^2 / line_len).""" |
| 241 | # TODO(rjryan): Speed up computation of twiddle factors using the |
| 242 | # following recurrence relation and cache them across invocations of RFFT. |
| 243 | # |
| 244 | # t_n = exp(sqrt(-1) * pi * n^2 / line_len) |
| 245 | # for n = 0, 1,..., line_len-1. |
| 246 | # For n > 2, use t_n = t_{n-1}^2 / t_{n-2} * t_1^2 |
| 247 | a = _array_ops.tile( |
| 248 | _array_ops.expand_dims(_math_ops.range(length), 0), (length, 1)) |
| 249 | b = _array_ops.transpose(a, [1, 0]) |
| 250 | return _math_ops.exp( |
| 251 | -2j * np.pi * _math_ops.cast(a * b, _dtypes.complex64) / |
| 252 | _math_ops.cast(length, _dtypes.complex64)) |
| 253 | |
| 254 | def _ymask(length): |
| 255 | """A sequence of [1+0j, -1+0j, 1+0j, -1+0j, ...] with length `length`.""" |
| 256 | return _math_ops.cast(1 - 2 * (_math_ops.range(length) % 2), |
| 257 | _dtypes.complex64) |
| 258 | |
| 259 | y0 = grad[..., 0:1] |
| 260 | if rank == 1: |
| 261 | ym = grad[..., -1:] |
| 262 | extra_terms = y0 + is_even * ym * _ymask(input_shape[-1]) |
| 263 | elif rank == 2: |
| 264 | # Create a mask matrix for y0 and ym. |
| 265 | base_mask = _mask_matrix(input_shape[-2]) |
| 266 | |
| 267 | # Tile base_mask to match y0 in shape so that we can batch-matmul the |
| 268 | # inner 2 dimensions. |
| 269 | tiled_mask = _tile_for_broadcasting(base_mask, y0) |
| 270 | |
| 271 | y0_term = _math_ops.matmul(tiled_mask, _math_ops.conj(y0)) |
| 272 | extra_terms = y0_term |
| 273 | |
| 274 | ym = grad[..., -1:] |
| 275 | ym_term = _math_ops.matmul(tiled_mask, _math_ops.conj(ym)) |
| 276 | |
| 277 | inner_dim = input_shape[-1] |
| 278 | ym_term = _array_ops.tile( |
| 279 | ym_term, |
| 280 | _array_ops.concat([ |
nothing calls this directly
no test coverage detected