def curfft(inp, norm=None):
r"""
Performs the fast Fourier transform of a real-valued input on the GPU.
The input must be a real-valued float32 variable of dimensions (m, ..., n).
It performs FFTs of size (..., n) on m batches.
The output is a GpuArray of dimensions (m, ..., n//2+1, 2). The second to
last dimension of the output contains the n//2+1 non-trivial elements of
the real-valued FFTs. The real and imaginary parts are stored as a pair of
float32 arrays.
Parameters
----------
inp
Array of real-valued float32 of size (m, ..., n), containing m inputs of
size (..., n).
norm : {None, 'ortho', 'no_norm'}
Normalization of transform. Following numpy, default *None* normalizes
only the inverse transform by n, 'ortho' yields the unitary transform
(:math:`1/\sqrt n` forward and inverse). In addition, 'no_norm' leaves
the transform unnormalized.
"""
s = inp.shape[1:]
cond_norm = _unitary(norm)
scaling = 1
if cond_norm == "ortho":
scaling = sqrt(s.prod().astype("float32"))
return curfft_op(inp, s) / scaling
def local_abstract_batch_norm_inference(fgraph, node):
if not isinstance(node.op, AbstractBatchNormInference):
return None
x, scale, bias, estimated_mean, estimated_variance, epsilon = node.inputs
if (
not isinstance(x.type, TensorType)
or not isinstance(scale.type, TensorType)
or not isinstance(bias.type, TensorType)
or not isinstance(estimated_mean.type, TensorType)
or not isinstance(estimated_variance.type, TensorType)
or not isinstance(epsilon.type, TensorType)
):
return None
# The epsilon should not upcast the dtype.
if estimated_variance.dtype == "float32" and epsilon.dtype == "float64":
epsilon = epsilon.astype("float32")
result = (x - estimated_mean) * (scale / sqrt(estimated_variance + epsilon)) + bias
result = aet.patternbroadcast(result, node.outputs[0].broadcastable)
for var in aesara.graph.basic.vars_between(node.inputs, [result]):
if var not in node.inputs:
copy_stack_trace(node.outputs[0], var)
return [result]
def grad(self, inputs, grads):
x, scale, bias, est_mean, est_var, epsilon = inputs
dy = grads[0]
axes = self.axes
if min(axes) < 0 or max(axes) >= x.ndim:
raise ValueError(
f"axes should be less than ndim (<{x.ndim}), but {axes} given"
)
scale, bias, est_mean, est_var = (
aet.addbroadcast(t, *axes) for t in (scale, bias, est_mean, est_var)
)
# define helper expressions
est_var_eps = est_var + epsilon
est_std = sqrt(est_var_eps)
two = aet.constant(2.0)
# define and return gradients
dx = dy * (scale / est_std)
dscale = (dy * (x - est_mean)).sum(axes, keepdims=True) / est_std
dbias = dy.sum(axes, keepdims=True)
dmean = -dy.sum(axes, keepdims=True) * (scale / est_std)
dvar = -(dy * (x - est_mean)).sum(axes, keepdims=True) * (
scale / (two * est_var_eps * est_std)
)
return [dx, dscale, dbias, dmean, dvar, aesara.gradient.DisconnectedType()()]
def local_abstract_batch_norm_train(fgraph, node):
if not isinstance(node.op, AbstractBatchNormTrain):
return None
x, scale, bias, epsilon, running_average_factor = node.inputs[:5]
axes = node.op.axes
if min(axes) < 0 or max(axes) > x.ndim:
return None
if (
not isinstance(x.type, TensorType)
or not isinstance(scale.type, TensorType)
or not isinstance(bias.type, TensorType)
or not isinstance(epsilon.type, TensorType)
or not isinstance(running_average_factor.type, TensorType)
):
return None
# optional running_mean and running_var
if len(node.inputs) > 5 and not isinstance(node.inputs[5].type, TensorType):
return None
if len(node.inputs) > 6 and not isinstance(node.inputs[6].type, TensorType):
return None
mean = x.mean(axes, keepdims=True)
var = x.var(axes, keepdims=True)
# The epsilon should not upcast the dtype.
if var.dtype == "float32" and epsilon.dtype == "float64":
epsilon = epsilon.astype("float32")
invstd = inv(sqrt(var + epsilon))
out = (x - mean) * (scale * invstd) + bias
results = [out, mean, invstd]
if len(node.inputs) > 5:
running_mean = node.inputs[5]
running_mean = (
running_mean * (1.0 - running_average_factor)
+ mean * running_average_factor
)
results.append(running_mean)
if len(node.inputs) > 6:
m = aet.cast(prod(x.shape) / prod(scale.shape), config.floatX)
running_var = node.inputs[6]
running_var = (
running_var * (1.0 - running_average_factor)
+ (m / (m - 1)) * var * running_average_factor
)
results.append(running_var)
results = [
aet.patternbroadcast(r, r_orig.broadcastable)
for (r, r_orig) in zip(results, node.outputs)
]
for var in aesara.graph.basic.vars_between(node.inputs, results):
if var not in node.inputs:
copy_stack_trace(node.outputs[0], var)
return results
def irfft(inp, norm=None, is_odd=False):
r"""
Performs the inverse fast Fourier Transform with real-valued output.
The input is a variable of dimensions (m, ..., n//2+1, 2)
representing the non-trivial elements of m real-valued Fourier transforms
of initial size (..., n). The real and imaginary parts are stored as a
pair of float arrays.
The output is a real-valued variable of dimensions (m, ..., n)
giving the m inverse FFTs.
Parameters
----------
inp
Array of size (m, ..., n//2+1, 2), containing m inputs
with n//2+1 non-trivial elements on the last dimension and real
and imaginary parts stored as separate real arrays.
norm : {None, 'ortho', 'no_norm'}
Normalization of transform. Following numpy, default *None* normalizes
only the inverse transform by n, 'ortho' yields the unitary transform
(:math:`1/\sqrt n` forward and inverse). In addition, 'no_norm' leaves
the transform unnormalized.
is_odd : {True, False}
Set to True to get a real inverse transform output with an odd last dimension
of length (N-1)*2 + 1 for an input last dimension of length N.
"""
if is_odd not in (True, False):
raise ValueError(
f"Invalid value {is_odd} for id_odd, must be True or False")
s = inp.shape[1:-1]
if is_odd:
s = set_subtensor(s[-1], (s[-1] - 1) * 2 + 1)
else:
s = set_subtensor(s[-1], (s[-1] - 1) * 2)
cond_norm = _unitary(norm)
scaling = 1
# Numpy's default normalization is 1/N on the inverse transform.
if cond_norm is None:
scaling = s.prod().astype(inp.dtype)
elif cond_norm == "ortho":
scaling = sqrt(s.prod().astype(inp.dtype))
return irfft_op(inp, s) / scaling
def normal(
self,
size,
avg=0.0,
std=1.0,
ndim=None,
dtype=None,
nstreams=None,
truncate=False,
**kwargs,
):
"""
Sample a tensor of values from a normal distribution.
Parameters
----------
size : int_vector_like
Array dimensions for the output tensor.
avg : float_like, optional
The mean value for the truncated normal to sample from (defaults to 0.0).
std : float_like, optional
The standard deviation for the truncated normal to sample from (defaults to 1.0).
truncate : bool, optional
Truncates the normal distribution at 2 standard deviations if True (defaults to False).
When this flag is set, the standard deviation of the result will be less than the one specified.
ndim : int, optional
The number of dimensions for the output tensor (defaults to None).
This argument is necessary if the size argument is ambiguous on the number of dimensions.
dtype : str, optional
The data-type for the output tensor. If not specified,
the dtype is inferred from avg and std, but it is at least as precise as floatX.
kwargs
Other keyword arguments for random number generation (see uniform).
Returns
-------
samples : TensorVariable
A Aesara tensor of samples randomly drawn from a normal distribution.
"""
size = _check_size(size)
avg = undefined_grad(as_tensor_variable(avg))
std = undefined_grad(as_tensor_variable(std))
if dtype is None:
dtype = aes.upcast(config.floatX, avg.dtype, std.dtype)
avg = at.cast(avg, dtype=dtype)
std = at.cast(std, dtype=dtype)
# generate even number of uniform samples
# Do manual constant folding to lower optiimizer work.
if isinstance(size, Constant):
n_odd_samples = size.prod(dtype="int64")
else:
n_odd_samples = prod(size, dtype="int64")
n_even_samples = n_odd_samples + n_odd_samples % 2
uniform = self.uniform(
(n_even_samples, ),
low=0.0,
high=1.0,
ndim=1,
dtype=dtype,
nstreams=nstreams,
**kwargs,
)
# box-muller transform
u1 = uniform[:n_even_samples // 2]
u2 = uniform[n_even_samples // 2:]
r = sqrt(-2.0 * log(u1))
theta = np.array(2.0 * np.pi, dtype=dtype) * u2
cos_theta, sin_theta = cos(theta), sin(theta)
z0 = r * cos_theta
z1 = r * sin_theta
if truncate:
# use valid samples
to_fix0 = (z0 < -2.0) | (z0 > 2.0)
to_fix1 = (z1 < -2.0) | (z1 > 2.0)
z0_valid = z0[at.nonzero(~to_fix0)]
z1_valid = z1[at.nonzero(~to_fix1)]
# re-sample invalid samples
to_fix0 = at.nonzero(to_fix0)[0]
to_fix1 = at.nonzero(to_fix1)[0]
n_fix_samples = to_fix0.size + to_fix1.size
lower = at.constant(1.0 / np.e**2, dtype=dtype)
u_fix = self.uniform(
(n_fix_samples, ),
low=lower,
high=1.0,
ndim=1,
dtype=dtype,
nstreams=nstreams,
**kwargs,
)
r_fix = sqrt(-2.0 * log(u_fix))
z0_fixed = r_fix[:to_fix0.size] * cos_theta[to_fix0]
z1_fixed = r_fix[to_fix0.size:] * sin_theta[to_fix1]
# pack everything together to a useful result
norm_samples = at.join(0, z0_valid, z0_fixed, z1_valid, z1_fixed)
else:
norm_samples = at.join(0, z0, z1)
if isinstance(n_odd_samples, Variable):
samples = norm_samples[:n_odd_samples]
elif n_odd_samples % 2 == 1:
samples = norm_samples[:-1]
else:
samples = norm_samples
samples = reshape(samples, newshape=size, ndim=ndim)
samples *= std
samples += avg
return samples
