# -*- coding: utf-8 -*-
# Copyright 2018 Google LLC. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Parameterizations for layer classes."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
# Dependency imports
import numpy as np
from scipy import fftpack
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
_matrix_cache = {}
def irdft_matrix(shape, dtype=dtypes.float32):
"""Matrix for implementing kernel reparameterization with `tf.matmul`.
This can be used to represent a kernel with the provided shape in the RDFT
domain.
Example code for kernel creation, assuming 2D kernels:
```
def create_kernel(init):
shape = init.shape.as_list()
matrix = irdft_matrix(shape[:2])
init = tf.reshape(init, (shape[0] * shape[1], shape[2] * shape[3]))
init = tf.matmul(tf.transpose(matrix), init)
kernel = tf.Variable(init)
kernel = tf.matmul(matrix, kernel)
kernel = tf.reshape(kernel, shape)
return kernel
```
Args:
shape: Iterable of integers. Shape of kernel to apply this matrix to.
dtype: `dtype` of returned matrix.
Returns:
`Tensor` of shape `(prod(shape), prod(shape))` and dtype `dtype`.
"""
shape = tuple(int(s) for s in shape)
dtype = dtypes.as_dtype(dtype)
key = (ops.get_default_graph(), "irdft", shape, dtype.as_datatype_enum)
matrix = _matrix_cache.get(key)
if matrix is None:
size = np.prod(shape)
rank = len(shape)
matrix = np.identity(size, dtype=np.float64).reshape((size,) + shape)
for axis in range(rank):
matrix = fftpack.rfft(matrix, axis=axis + 1)
slices = (rank + 1) * [slice(None)]
if shape[axis] % 2 == 1:
slices[axis + 1] = slice(1, None)
else:
slices[axis + 1] = slice(1, -1)
matrix[tuple(slices)] *= np.sqrt(2)
matrix /= np.sqrt(size)
matrix = np.reshape(matrix, (size, size))
matrix = array_ops.constant(
matrix, dtype=dtype, name="irdft_" + "x".join([str(s) for s in shape]))
_matrix_cache[key] = matrix
return matrix
