def _to_dense(self): row = ops.convert_to_tensor(self.row) col = ops.convert_to_tensor(self.col) total_shape = array_ops.broadcast_dynamic_shape( prefer_static.shape(row), prefer_static.shape(col)) n = prefer_static.shape(row)[-1] row = _ops.broadcast_to(row, total_shape) col = _ops.broadcast_to(col, total_shape) # We concatenate the column in reverse order to the row. # This gives us 2*n + 1 elements. elements = prefer_static.concat( [array_ops.reverse(col, axis=[-1]), row[..., 1:]], axis=-1) # Given the above vector, the i-th row of the Toeplitz matrix # is the last n elements of the above vector shifted i right # (hence the first row is just the row vector provided, and # the first element of each row will belong to the column vector). # We construct these set of indices below. indices = math_ops.mod( # How much to shift right. This corresponds to `i`. array_ops.range(0, n) + # Specifies the last `n` indices. array_ops.range(n - 1, -1, -1)[..., _ops.newaxis], # Mod out by the total number of elements to ensure the index is # non-negative (for tf.gather) and < 2 * n - 1. 2 * n - 1) return array_ops.gather(elements, indices, axis=-1)
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)
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 = array_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)