Inversion#

class cuequivariance_torch.Inversion#

Inversion layer for \(O(3)\) representations.

Parameters:

  • irreps (Irreps) – The irreducible representations of the tensor to invert.

  • layout (IrrepsLayout, optional) – The memory layout of the tensor, cue.ir_mul is preferred.

  • layout_in (IrrepsLayout, optional) – The layout of the input irreducible representations, by default layout.

  • layout_out (IrrepsLayout, optional) – The layout of the output irreducible representations, by default layout.

  • device (torch.device, optional) – The device to use for the linear layer.

  • math_dtype (torch.dtype, optional) – The dtype to use for the math operations, by default it follows the dtype of the input tensors.

  • method (str, optional) – The method to use for the linear layer, by default “uniform_1d” (using a CUDA kernel) if all segments have the same shape, otherwise “naive” (using a PyTorch implementation).

  • use_fallback (bool, optional, deprecated) – Whether to use a “fallback” implementation, now maps to method: If True the “naive” method is used. If False the “uniform_1d” method is used (make sure all segments have the same shape).

Forward Pass

forward(x)#

Apply the inversion layer.

Parameters:

x (Tensor)

Return type:

Tensor

__init__(
irreps,
*,
layout=None,
layout_in=None,
layout_out=None,
device=None,
math_dtype=None,
use_fallback=None,
method=None,
)#

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Parameters:
classmethod __new__(*args, **kwargs)#