def _trace(self):
# The diagonal of the [[nested] block] circulant operator is the mean of
# the spectrum.
# Proof: For the [0,...,0] element, this follows from the IDFT formula.
# Then the result follows since all diagonal elements are the same.
# Therefore, the trace is the sum of the spectrum.
# Get shape of diag along with the axis over which to reduce the spectrum.
# We will reduce the spectrum over all block indices.
if tensor_shape.TensorShape(self.spectrum.shape).is_fully_defined():
spec_rank = tensor_shape.TensorShape(self.spectrum.shape).ndims
axis = np.arange(spec_rank - self.block_depth, spec_rank, dtype=np.int32)
else:
spec_rank = array_ops.rank(self.spectrum)
axis = math_ops.range(spec_rank - self.block_depth, spec_rank)
# Real diag part "re_d".
# Suppose tensor_shape.TensorShape(spectrum.shape) = [B1,...,Bb, N1, N2]
# tensor_shape.TensorShape(self.shape) = [B1,...,Bb, N, N], with N1 * N2 = N.
# tensor_shape.TensorShape(re_d_value.shape) = [B1,...,Bb]
re_d_value = math_ops.reduce_sum(math_ops.real(self.spectrum), axis=axis)
if not np.issubdtype(self.dtype, np.complexfloating):
return _ops.cast(re_d_value, self.dtype)
# Imaginary part, "im_d".
if self.is_self_adjoint:
im_d_value = array_ops.zeros_like(re_d_value)
else:
im_d_value = math_ops.reduce_sum(math_ops.imag(self.spectrum), axis=axis)
return _ops.cast(math_ops.complex(re_d_value, im_d_value), self.dtype)
def _rotate_last_dim(x, rotate_right=False):
"""Rotate the last dimension either left or right."""
ndims = array_ops.rank(x)
if rotate_right:
transpose_perm = array_ops.concat(
[[ndims - 1], array_ops.range(0, ndims - 1)], axis=0)
else:
transpose_perm = array_ops.concat([array_ops.range(1, ndims), [0]],
axis=0)
return array_ops.transpose(x, transpose_perm)
