e3x.nn.features.rotate
- e3x.nn.features.rotate(x, wigner_d_matrix)[source]
Rotates the input features with the given Wigner D-matrix (or -matrices).
- Parameters:
x (
Union[Float[Array, '... 1 (max_degree+1)**2 num_features'], Float[Array, '... 2 (max_degree+1)**2 num_features']]) – Input features with shape (…, S, (L+1)**2, F).wigner_d_matrix (
Union[Float[Array, '(max_degree+1)**2 (max_degree+1)**2'], Float[Array, '... (max_degree+1)**2 (max_degree+1)**2']]) – Wigner D-matrix with shape ((L+1)**2, (L+1)**2) or a batch of Wigner D-matrices with shape (…, (L+1)**2, (L+1)**2).
- Return type:
- Returns:
The rotated features.