<commit_before>"""
Implementation from MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields
<NAME>, <NAME>, <NAME>, <NAME> and <NAME>
"""
from typing import Set, Optional, Callable, Tuple

import e3nn_jax as e3nn
import haiku as hk
import jax.numpy as jnp


class SymmetricTensorProduct(hk.Module):
    r"""Symmetric tensor product contraction with parameters

    Equivalent to the following code:

    ```python
    e3nn.Linear(irreps_out)(
        e3nn.concatenate([
            x,
            tensor_product(x, x),  # additionally keeping only the symmetric terms
            tensor_product(tensor_product(x, x), x),
            ...
        ])
    )
    ```

    Args:
        orders (tuple of int): orders of the tensor product
        keep_irrep_out (optional, set of Irrep): irreps to keep in the output
        get_parameter (optional, callable): function to get the parameters, by default it uses `hk.get_parameter`
            it should have the signature `get_parameter(name, shape) -> ndarray` and return a normal distribution
            with variance 1
    """

    def __init__(
        self,
        orders: Tuple[int, ...],
        keep_irrep_out: Optional[Set[e3nn.Irrep]] = None,<commit_msg>complete docstring of SymmetricTensorProduct<commit_after>"""
Implementation from MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields
<NAME>, <NAME>, <NAME>, <NAME> and <NAME>
"""
from typing import Set, Optional, Callable, Tuple

import e3nn_jax as e3nn
import haiku as hk
import jax.numpy as jnp


class SymmetricTensorProduct(hk.Module):
    r"""Symmetric tensor product contraction with parameters

    Equivalent to the following code executed in parallel on the channel dimension:

    ```python
    e3nn.Linear(irreps_out)(
        e3nn.concatenate([
            x,
            tensor_product(x, x),  # additionally keeping only the symmetric terms
            tensor_product(tensor_product(x, x), x),
            ...
        ])
    )
    ```

    Each channel has its own parameters.

    Args:
        orders (tuple of int): orders of the tensor product
        keep_irrep_out (optional, set of Irrep): irreps to keep in the output
        get_parameter (optional, callable): function to get the parameters, by default it uses `hk.get_parameter`
            it should have the signature `get_parameter(name, shape) -> ndarray` and return a normal distribution
            with variance 1
    """

    def __init__(
        self,
        orders: Tuple[int, ...],
        keep_irrep_out: Optional[Set[e3nn.Irrep]] = None,
{'fraction_non_alphanumeric': 0.08482312666413085, 'fraction_numerical': 0.004564473183720046, 'mean_word_length': 2.8506588579795022, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1356255', 'n_tokens_mistral': 753, 'n_tokens_neox': 721, 'n_words': 271}