lgatr.primitives.nonlinearities

Gated nonlinearities on multivectors.

Functions

gated_gelu(x, gates)

Pin-equivariant gated GeLU nonlinearity without division.

gated_relu(x, gates)

Pin-equivariant gated ReLU nonlinearity.

gated_sigmoid(x, gates)

Pin-equivariant gated sigmoid nonlinearity.

gated_silu(x, gates)

Pin-equivariant gated SiLU or Swish nonlinearity without division.
lgatr.primitives.nonlinearities.gated_gelu(x, gates)[source]

Pin-equivariant gated GeLU nonlinearity without division.

Given multivector input x and scalar input gates (with matching batch dimensions), computes GeLU(gates) * x.

References

Dan Hendrycks, Kevin Gimpel, “Gaussian Error Linear Units (GELUs)”, arXiv:1606.08415

Parameters:

  • x (torch.Tensor) – Multivector input with shape (…, 16)

  • gates (torch.Tensor) – Pin-invariant gates with shape (…, 1).

Returns:

outputs – Computes GeLU(gates) * x, with broadcasting along the last dimension. Result has shape (…, 16)

Return type:

torch.Tensor

lgatr.primitives.nonlinearities.gated_relu(x, gates)[source]

Pin-equivariant gated ReLU nonlinearity.

Given multivector input x and scalar input gates (with matching batch dimensions), computes ReLU(gates) * x.

Parameters:

  • x (torch.Tensor) – Multivector input with shape (…, 16)

  • gates (torch.Tensor) – Pin-invariant gates with shape (…, 1).

Returns:

outputs – Computes ReLU(gates) * x, with broadcasting along the last dimension. Result has shape (…, 16)

Return type:

torch.Tensor

lgatr.primitives.nonlinearities.gated_sigmoid(x, gates)[source]

Pin-equivariant gated sigmoid nonlinearity.

Given multivector input x and scalar input gates (with matching batch dimensions), computes sigmoid(gates) * x.

Parameters:

  • x (torch.Tensor) – Multivector input with shape (…, 16)

  • gates (torch.Tensor) – Pin-invariant gates with shape (…, 1).

Returns:

outputs – Computes Sigmoid(gates) * x, with broadcasting along the last dimension. Result has shape (…, 16)

Return type:

torch.Tensor

lgatr.primitives.nonlinearities.gated_silu(x, gates)[source]

Pin-equivariant gated SiLU or Swish nonlinearity without division.

Given multivector input x and scalar input gates (with matching batch dimensions), computes Swish(gates) * x.

References

Stefan Elfwing, Eiji Uchibe, Kenji Doya, “Sigmoid-Weighted Linear Units for Neural Network Function Approximation in Reinforcement Learning”, arXiv:1702.03118

Parameters:

  • x (torch.Tensor) – Multivector input with shape (…, 16)

  • gates (torch.Tensor) – Pin-invariant gates with shape (…, 1).

Returns:

outputs – Computes Swish(gates) * x, with broadcasting along the last dimension. Result has shape (…, 16)

Return type:

torch.Tensor