mindchemistry.e3.o3.TensorProduct

class mindchemistry.e3.o3.TensorProduct(irreps_in1, irreps_in2=None, irreps_out=None, instructions='full', dtype=float32, irrep_norm='component', path_norm='element', weight_init='normal', weight_mode='inner', core_mode='ncon', ncon_dtype=float32)

Versatile tensor product operator of two input Irreps and a output Irreps, that sends two tensors into a tensor and keep the geometric tensor properties. This class integrates different typical usages: TensorSquare, FullTensorProduct, FullyConnectedTensorProduct, ElementwiseTensorProduct and Linear.

A TensorProduct class defines an algebraic structure with equivariance. Ones the TensorProduct object is created and initialized, the algorithm is determined. For any given two legal input tensors, this object will provide a output tensor. If the object do not have learnable weights, the output tensor is deterministic. When the learnable weights are introduced, this operator will correspond to a general bilinear, equivariant operation, as a generalization of the standard tensor product.

If irreps_in2 is not specified, it will be assigned as irreps_in1, corresponding to TensorSquare. If irreps_out is not specified, this operator will account all possible output irreps. If both irreps_out and instructions are not specified, this operator is the standard tensor product without any learnable weights, corresponding to FullTensorProduct.

Each output irrep should satisfy:

\[\| l_1 - l_2 \| \leq l_{out} \leq \| l_1 + l_2 \| p_1 p_2 = p_{out}\]

Parameters

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

• irreps_in2 (Union[str, Irrep, Irreps, None]) – Irreps for the second input. Default: None. If irreps_in2 is None, irreps_in2 will be assigned as '0e' in 'linear' instructions, or be assigned as irreps_in1 in otherwise, corresponding to TensorSquare.

• irreps_out (Union[str, Irrep, Irreps, None]) – Irreps for the output in 'connect' and custom instructions, or filter irreps for the output in otherwise. If irreps_out is None, irreps_out will be the full tensor product irreps (including all possible paths). Default: None.

• instructions (Union[str, List[Tuple[int, int, int, str, bool, (float)]]]) –

  List of tensor product path instructions. Default: 'full'. For str in {'full', 'connect', 'element', 'linear', 'mearge'}, the instructions are constructed automatically according to the different modes:

  • 'full': each output irrep for every pair of input irreps — is created and returned independently. The outputs are not mixed with each other. Corresponding to the standard tensor product FullTensorProduct if irreps_out is not specified.

  • 'connect': each output is a learned weighted sum of compatible paths. This allows the operator to produce outputs with any multiplicity. Corresponding to FullyConnectedTensorProduct.

  • 'element': the irreps are multiplied one-by-one. The inputs will be split and that the multiplicities of the outputs match with the multiplicities of the input. Corresponding to ElementwiseTensorProduct.

  • 'linear': linear operation equivariant on the first irreps, while the second irreps is set to be '0e'. This can be regarded as the geometric tensors version of teh dense layer. Corresponding to Linear.

  • 'merge': Automatically build 'uvu' mode instructions with trainable parameters. The irreps_out here plays the role of output filters.

  For List[Tuple[int, int, int, str, bool, (float)]], the instructions are constructed manually.

  Each instruction contain a tuple: (indice_one, indice_two, i_out, mode, has_weight, (optional: path_weight)). Each instruction puts in1[indice_one] \(\otimes\) in2[indice_two] into out[i_out].

  • indice_one, indice_two, i_out: int, the index of the irrep in irreps for irreps_in1, irreps_in2 and irreps_out correspondingly.

  • mode: str in {'uvw', 'uvu', 'uvv', 'uuw', 'uuu', 'uvuv'}, the way of the multiplicities of each path are treated. 'uvw' is the fully mixed mode.

  • has_weight: bool, True if this path should have learnable weights, otherwise False.

  • path_weight:float, a multiplicative weight to apply to the output of this path. Defaults: 1.0.

• irrep_norm (str) –

  {'component', 'norm'}, the assumed normalization of the input and output representations. Default: 'component'.

  • 'norm': \(\| x \| = \| y \| = 1 \Longrightarrow \| x \otimes y \| = 1\)

• path_norm (str) –

  {'element', 'path'}, the normalization method of path weights. Default: 'element'.

  • 'element': each output is normalized by the total number of elements (independently of their paths).

  • 'path': each path is normalized by the total number of elements in the path, then each output is normalized by the number of paths.

• weight_init (str) – {'zeros', 'ones', 'truncatedNormal', 'normal', 'uniform', 'he_uniform', 'he_normal', 'xavier_uniform'}, the initial method of weights. Default: 'normal'.

• weight_mode (str) –

  {'inner', 'share', 'custom'} determine the weights' mode. Default: 'inner'.

  • 'inner': weights will initialized in the tensor product internally.

  • 'share': weights should given manually without batch dimension.

  • 'custom': weights should given manually with batch dimension.

• core_mode (str) – {'ncon', 'einsum'} determine the core computation mode. Default: 'ncon'.

• dtype (mindspore.dtype) – The type of input tensor. Default: mindspore.float32 .

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

Inputs:

• x (Tensor) - The shape of Tensor is (..., irreps_in1.dim)

• y (Tensor) - The shape of Tensor is (..., irreps_in2.dim)

• weight (Tensor) - Tensor or list of Tensor, optional required if internal_weights is False. The shape of Tensor is (self.weight_numel,) if shared_weights is True. The shape of Tensor is (..., self.weight_numel) if shared_weights is False or list of tensors of shapes weight_shape / (...) + weight_shape. Use self.instructions to know what are the weights used for. The shape of Tensor is (..., irreps_out.dim).

Outputs:

• outputs (Tensor) - The shape of Tensor is (..., irreps_out.dim).

Raises

• ValueError – If irreps_out is not legal.

• ValueError – If the connection mode is not in ['uvw', 'uvu', 'uvv', 'uuw', 'uuu', 'uvuv'].

• ValueError – If the degree of inputs and output do not match.

• ValueError – If the parity of inputs and output do not match.

• ValueError – If the multiplicity of inputs and output do not match.

• ValueError – If the connection mode is 'uvw', but has_weight is False.

• ValueError – If the connection mode is 'uuw' and has_weight is False, but the multiplicity is not equal to 1.

• ValueError – If the initial method is not supported.

• ValueError – If the number of input tensors is not match to the number of input irreps.

Supported Platforms:

Ascend

Examples

>>> import mindspore as ms
>>> from mindchemistry.e3.o3 import TensorProduct
Standard tensor product:
>>> tp1 = TensorProduct('2x1o+4x0o', '1x1o+3x0e')
TensorProduct [full] (2x1o+4x0o x 1x1o+3x0e -> 2x0e+12x0o+6x1o+2x1e+4x1e+2x2e)
>>> v1 = ms.Tensor(np.linspace(1., 2., tp1.irreps_in1.dim), dtype=ms.float32)
>>> v2 = ms.Tensor(np.linspace(2., 3., tp1.irreps_in2.dim), dtype=ms.float32)
>>> tp1(v1, v2).shape
(1, 60)
Elementwise tensor product:
>>> tp2 = TensorProduct('2x2e+4x1o', '3x1e+3x0o')
TensorProduct [element] (2x2e+1x1o+3x1o x 2x1e+1x1e+3x0o -> 2x1e+2x2e+2x3e+1x0o+1x1o+1x2o+3x1e)
>>> tp2.instructions
[(0, 0, 0, 'uuu', False), (0, 0, 1, 'uuu', False), (0, 0, 2, 'uuu', False), (1, 1, 3, 'uuu', False),
(1, 1, 4, 'uuu', False), (1, 1, 5, 'uuu', False), (2, 2, 6, 'uuu', False)]
Custom tensor product with learnable weights:
>>> tp3 = TensorProduct(
...     '3x2o+2x1o', '2x2e+4x1o+5x0e', '2x3o+8x1e+10x1o',
...     [
...         (0,0,0,'uvv',True),
...         (1,0,0,'uuu',True),
...         (1,1,1,'uvuv',True),
...         (1,2,2,'uvw',True)
...     ]
... )
TensorProduct [custom] (3x2o+2x1o x 2x2e+4x1o+5x0e -> 2x3o+8x1e+10x1o)
>>> [w.shape for w in tp3.weights]
[(3, 2), (2,), (2, 4), (2, 5, 10)]
Linear operation with an output filter:
>>> tp4 = TensorProduct('2x1o', irreps_out='5x2e+4x1e+7x1o', instructions='connect')
TensorProduct [linear] (2x2e+3x1o+3x0e x 1x0e -> 3x2e+5x1o+2x0e)
>>> v1 = ms.Tensor(np.linspace(1., 2., tp4.irreps_in1.dim), dtype=ms.float32)
>>> tp4(v1).shape
(1, 32)