def fourier_poisson(tensor, times=1): """ Inverse operation to `fourier_laplace`. """ frequencies = math.fft(math.to_complex(tensor)) k = fftfreq(math.staticshape(tensor)[1:-1], mode='square') fft_laplace = -(2 * np.pi)**2 * k fft_laplace[(0, ) * math.ndims(k)] = np.inf inv_fft_laplace = 1 / fft_laplace inv_fft_laplace[(0, ) * math.ndims(k)] = 0 return math.real(math.ifft(frequencies * inv_fft_laplace**times))
def normalize_to(target, source=1, epsilon=1e-5, batch_dims=1): """ Multiplies the target so that its total content matches the source. :param target: a tensor :param source: a tensor or number :param epsilon: small number to prevent division by zero or None. :return: normalized tensor of the same shape as target """ target_total = math.sum(target, axis=tuple(range(batch_dims, math.ndims(target))), keepdims=True) denominator = math.maximum(target_total, epsilon) if epsilon is not None else target_total source_total = math.sum(source, axis=tuple(range(batch_dims, math.ndims(source))), keepdims=True) return target * (source_total / denominator)
def batch_align_scalar(tensor, innate_spatial_dims, target): if rank(tensor) == 0: assert innate_spatial_dims == 0 return math.expand_dims(tensor, 0, len(math.staticshape(target))) if math.staticshape(tensor)[-1] != 1 or math.ndims(tensor) <= 1: tensor = math.expand_dims(tensor, -1) result = batch_align(tensor, innate_spatial_dims + 1, target) return result
def spatial_rank(tensor): """ The spatial rank of a tensor is ndims - 2. """ return math.ndims(tensor) - 2
def is_scalar(tensor): return math.ndims(tensor) == 0