irreps_add

cuequivariance_jax.ir_dict.irreps_add(x, y)

Element-wise addition of two dict[Irrep, Array] representations.

Parameters:

• x (dict[Irrep, Array]) – First dictionary with arrays of shape (…, mul, ir.dim).

• y (dict[Irrep, Array]) – Second dictionary with same keys and compatible shapes.

Returns:

Dictionary with element-wise sum of corresponding arrays.

Raises:

AssertionError – If the dictionaries have different keys.

Return type:

dict[Irrep, Array]