Functional

complextorch.nn.functional.apply_complex(real_module: Module, imag_module: Module, x: CVTensor) → CVTensor

Apply Complex

Naive complex computation between a complex-valued module defined by two real-valued modules and a complex-valued tensor (CVTensor).

Gauss’ trick is often faster and is implemented throughout this package.

Given a complex-valued tensor \(\mathbf{z} = \mathbf{x} + j \mathbf{y}\) and a linear function \(G(\cdot) = G_\mathbb{R}(\cdot) + j G_\mathbb{I}(\cdot)\), implements the following operation:

\[G(\mathbf{z}) = G_\mathbb{R}(\mathbf{x}) - G_\mathbb{I}(\mathbf{y}) + j(G_\mathbb{R}(\mathbf{y}) + G_\mathbb{I}(\mathbf{x}))\]

complextorch.nn.functional.apply_complex_polar(mag_fun, phase_fun, x: CVTensor) → CVTensor

Apply Complex Polar

Applies a polar function (\(G(\mathbf{z}) = G_{||}(|\mathbf{z}|) \odot \exp(j G_\angle(\angle\mathbf{z}))\)) to the magnitude and phase of the input tensor (\(\mathbf{z}\)) separately.

Implements the following operation:

\[G(\mathbf{z}) = G_{||}(|\mathbf{z}|) \odot \exp(j G_\angle(\angle\mathbf{z}))\]

param mag_fun:

function to be applied to the magnitude of the input tensor

param phase_fun:

function to be applied to the phase of the input tensor

param x:

input tensor

type x:

CVTensor

returns:

\(G_{||}(|\mathbf{z}|) \odot \exp(j G_\angle(\angle\mathbf{z}))\)

rtype:

CVTensor

complextorch.nn.functional.apply_complex_split(r_fun, i_fun, x: CVTensor) → CVTensor

Apply Complex Split

Applies a split function (\(G(\cdot) = G_\mathbb{R}(\cdot) + j G_\mathbb{I}(\cdot)\)) to the real and imaginary parts of the input tensor (\(\mathbf{z}\)) separately.

Implements the following operation:

\[G(\mathbf{z}) = G_\mathbb{R}(\mathbf{x}) + j G_\mathbb{I}(\mathbf{y})\]

param r_fun:

function to be applied to the real part of the input tensor

param i_fun:

function to be applied to the imaginary part of the input tensor

param x:

input tensor

type x:

CVTensor

returns:

\(G_\mathbb{R}(\mathbf{x}) + j G_\mathbb{I}(\mathbf{y})\)

rtype:

CVTensor

complextorch.nn.functional.cv_batch_norm(x: CVTensor, running_mean: Tensor | None = None, running_var: Tensor | None = None, weight: Tensor | None = None, bias: Tensor | None = None, training: bool = True, momentum: float = 0.1, eps: float = 1e-05) → CVTensor

Complex-Valued Batch Normalization

Applies complex-valued batch normalization as described in (Trabelsi et al., 2018) for each channel across a batch of data.

param x:

The input complex-valued data is expected to be at least 2d, with shape [B, F, …], where B is the batch dimension, F – the channels/features, … – the spatial dimensions (if present).

type x:

cvtorch.CVTensor

param running_mean:

The tensor with running mean statistics having shape [2, F]. Ignored if explicitly None.

type running_mean:

torch.tensor, or None

param running_var:

The tensor with running real-imaginary covariance statistics having shape [2, 2, F]. Ignored if explicitly None.

type running_var:

torch.tensor, or None

param weight:

The 2x2 weight matrix of the affine transformation of real and imaginary parts post normalization. Has shape [2, 2, F] . Ignored together with bias if explicitly None.

type weight:

torch.tensor, default=None

param bias:

The offest (bias) of the affine transformation of real and imaginary parts post normalization. Has shape [2, F] . Ignored together with weight if explicitly None.

type bias:

torch.tensor, or None

param training:

Determines whether to update running feature statistics, if they are provided, or use them instead of batch computed statistics. If False then running_mean and running_var MUST be provided.

type training:

bool, default=True

param momentum:

The weight in the exponential moving average used to keep track of the running feature statistics.

type momentum:

float, default=0.1

param eps:

The ridge coefficient to stabilize the estimate of the real-imaginary covariance.

type eps:

float, default=1e-5

complextorch.nn.functional.cv_layer_norm(x: CVTensor, normalized_shape: List[int], weight: Tensor | None = None, bias: Tensor | None = None, eps: float = 1e-05) → CVTensor

Complex-Valued Layer Normalization

Applies complex-valued layer normalization extending the work of (Trabelsi et al., 2018) for each channel across a batch of data.

param x:

The input complex-valued data is expected to be at least 2d, with shape [B, F, …], where B is the batch dimension, F – the

type x:

cvtorch.CVTensor

param weight:

The 2x2 weight matrix of the affine transformation of real and imaginary parts post normalization. Has shape [2, 2, F] . Ignored together with bias if explicitly None.

type weight:

torch.tensor, default=None

param bias:

The offest (bias) of the affine transformation of real and imaginary parts post normalization. Has shape [2, F] . Ignored together with weight if explicitly None.

type bias:

torch.tensor, or None

param eps:

The ridge coefficient to stabilize the estimate of the real-imaginary covariance.

type eps:

float, default=1e-5

complextorch.nn.functional.inv_sqrtm2x2(a: Tensor, b: Tensor, c: Tensor, d: Tensor, symmetric: bool = False)

Inverse Squareroot of 2x2 Matrix

Compute the inverse matrix square root of a 2x2 matrix: \(A^{-1/2}\) Improves computation speed of batch and layer normalization compared with PyTorch matrix inversion.

Following: https://en.wikipedia.org/wiki/Square_root_of_a_2_by_2_matrix

Given matrix \(\mathbf{A}\) as

\[\begin{split}\mathbf{A} = \begin{bmatrix} a & b \\ c & d \end{bmatrix}.\end{split}\]

Recall

\[\begin{split}\mathbf{A}^{-1} = \frac{1}{\text{det}(\mathbf{A})} \begin{bmatrix} d & -b \\ -c & a \end{bmatrix}.\end{split}\]

We define two parameters

\[ \begin{align}\begin{aligned}\delta &\triangleq \text{det}(\mathbf{A}) = ad - bc,\\\tau &\triangleq \text{trace}(\mathbf{A}) = a + d.\end{aligned}\end{align} \]

Using \(\delta\) and \(\tau\), we define two parameters to establish the relationship between \(\mathbf{A}\) and its matrix square root \(\mathbf{A}^{1/2}\) as

\[ \begin{align}\begin{aligned}s \triangleq \sqrt{\delta},\\t \triangleq \sqrt{\tau + 2s}.\end{aligned}\end{align} \]

The matrix square root can be expressed as

\[\begin{split}\mathbf{A}^{1/2} = \frac{1}{t} \begin{bmatrix} a+s & b \\ c & d+s \end{bmatrix}.\end{split}\]

Hence, the inverse of the matrix square root can be defined as

\[\begin{split}\mathbf{A}^{-1/2} = \frac{1}{st} \begin{bmatrix} d+s & -b \\ -c & a+s \end{bmatrix}.\end{split}\]

Finally, defining

\[\begin{split}\mathbf{B} \triangleq \begin{bmatrix} w & x \\ y & z \end{bmatrix} \triangleq \mathbf{A}^{-1/2}.\end{split}\]

Hence,

\[ \begin{align}\begin{aligned}w &= \frac{d + s}{st},\\x &= \frac{-b}{st},\\y &= \frac{-c}{st},\\z &= \frac{a + s}{st}.\end{aligned}\end{align} \]

param a:

a11 element of matrix A

type a:

torch.Tensor

param b:

a12 element of matrix A

type b:

torch.Tensor

param c:

a21 element of matrix A

type c:

torch.Tensor

param d:

a22 element of matrix A

type d:

torch.Tensor

param symmetric:

Boolean whether or not matrix A is symmetric. Defaults to False.

type symmetric:

bool, optional

returns:

\(w, x, y, z\) elemets of matrix B, the matrix inverse square root of matrix A

rtype:

Tuple[torch.Tensor]