mindscience.e3nn.o3.TensorSquare

class mindscience.e3nn.o3.TensorSquare(irreps_in, irreps_out=None, filter_ir_out=None, ncon_dtype=mindspore.float32, **kwargs)[source]

Compute the square tensor product of a tensor.

Equivalent to TensorProduct with irreps_in2=None and instructions='full' or 'connect'. For details, see mindscience.e3nn.o3.TensorProduct.

If irreps_out is given, this operation is fully connected. If irreps_out is not given, the operation has no parameter and is like full tensor product.

Parameters

  • irreps_in (Union[str, Irrep, Irreps]) – Irreps for the input.

  • irreps_out (Union[str, Irrep, Irreps, None], optional) – Irreps for the output. Default: None.

  • filter_ir_out (Union[str, Irrep, Irreps, None], optional) – Filter to select only specific Irrep of the output. Default: None.

  • ncon_dtype (mindspore.dtype, optional) – The type of input tensors of ncon computation module. Default: mindspore.float32.

Raises

ValueError – If both irreps_out and filter_ir_out are not None.

Examples

>>> from mindscience.e3nn.o3 import TensorSquare
>>> TensorSquare('2x1o', irreps_out='5x2e+4x1e+7x1o')
TensorProduct [connect] (2x1o x 2x1o -> 5x2e+4x1e)
>>> TensorSquare('2x1o+3x0e', filter_ir_out='5x2o+4x1e+2x0e')
TensorProduct [full] (2x1o+3x0e x 2x1o+3x0e -> 4x0e+9x0e+4x1e)