<commit_before>import math
from typing import Tuple

import haiku as hk
import jax
import jax.numpy as jnp
from e3nn_jax import (
    FunctionalFullyConnectedTensorProduct,
    FunctionalLinear,
    Irreps,
    IrrepsData,
    soft_one_hot_linspace,
    spherical_harmonics,
)
from jax import lax


class Convolution(hk.Module):
    def __init__(
        self,
        irreps_out,
        irreps_sh,
        diameter: float,
        num_radial_basis: int,
        steps: Tuple[float, float, float],
        *,
        padding="SAME",
        irreps_in=None,
    ):
        super().__init__()

        self.irreps_in = Irreps(irreps_in) if irreps_in is not None else None
        self.irreps_out = Irreps(irreps_out)
        self.irreps_sh = Irreps(irreps_sh)
        self.diameter = diameter
        self.num_radial_basis = num_radial_basis
        self.steps = steps
        self.padding = padding

        with jax.ensure_compile_time_eval():
            r = self.diameter / 2

            s = math.floor(r / self.steps[0])
            x = jnp.arange(-s, s + 1.0) * self.steps[0]

            s = math.floor(r / self.steps[1])
            y = jnp.arange(-s, s + 1.0) * self.steps[1]

            s = math.floor(r / self.steps[2])
            z = jnp.arange(-s, s + 1.0) * self.steps[2]

            lattice = jnp.stack(jnp.meshgrid(x, y, z, indexing="ij"), axis=-1)  # [x, y, z, R^3]

            self.emb = soft_one_hot_linspace(
                jnp.linalg.norm(lattice, ord=2, axis=-1),
                start=0.0,
                end=self.diameter / 2,
                number=self.num_radial_basis,
                basis="smooth_finite",
                start_zero=True,
                end_zero=True,
            )  # [x, y, z, num_radial_basis]

            self.sh = spherical_harmonics(
                irreps_out=self.irreps_sh, input=lattice, normalize=True, normalization="component"
            )  # [x, y, z, irreps_sh.dim]

    def kernel(self, irreps_in, irreps_out) -> jnp.ndarray:
        # convolution
        tp = FunctionalFullyConnectedTensorProduct(irreps_in, self.irreps_sh, irreps_out)

        w = [
            hk.get_parameter(
                f"w[{i.i_in1},{i.i_in2},{i.i_out}] {tp.irreps_in1[i.i_in1]},{tp.irreps_in2[i.i_in2]},{tp.irreps_out[i.i_out]}",
                (self.num_radial_basis,) + i.path_shape,
                init=hk.initializers.RandomNormal(),
            )
            for i in tp.instructions
        ]
        w = [
            jnp.einsum("xyzk,k...->xyz...", self.emb, x)
            / (self.sh.shape[0] * self.sh.shape[1] * self.sh.shape[2])  # [x,y,z, tp_w]
            for x in w
        ]

        tp_right = tp.right
        for _ in range(3):
            tp_right = jax.vmap(tp_right, (0, 0), 0)
        k = tp_right(w, self.sh)  # [x,y,z, irreps_in.dim, irreps_out.dim]

        # self-connection
        lin = FunctionalLinear(irreps_in, irreps_out)
        w = [
            hk.get_parameter(
                f"self-connection[{ins.i_in},{ins.i_out}] {lin.irreps_in[ins.i_in]},{lin.irreps_out[ins.i_out]}",
                shape=ins.path_shape,
                init=hk.initializers.RandomNormal(),
            )
            for ins in lin.instructions
        ]
        # note that lattice[center] is always displacement zero
        k = k.at[k.shape[0] // 2, k.shape[1] // 2, k.shape[2] // 2].set(lin.matrix(w))
        return k

    def __call__(self, input: IrrepsData) -> IrrepsData:
        """
        input: [batch, x, y, z, irreps_in.dim]
        """
        if self.irreps_in is not None:
            input = IrrepsData.new(self.irreps_in, input)
        if not isinstance(input, IrrepsData):
            raise ValueError("Convolution: input should be of type IrrepsData")
<commit_msg>voxel conv, more flexible emb<commit_after>import math
from typing import Tuple

import haiku as hk
import jax
import jax.numpy as jnp
from e3nn_jax import (
    FunctionalFullyConnectedTensorProduct,
    FunctionalLinear,
    Irreps,
    IrrepsData,
    MulIrrep,
    soft_one_hot_linspace,
    spherical_harmonics,
)
from jax import lax


class Convolution(hk.Module):
    def __init__(
        self,
        irreps_out,
        irreps_sh,
        diameter: float,
        num_radial_basis: int,
        steps: Tuple[float, float, float],
        *,
        padding="SAME",
        irreps_in=None,
    ):
        super().__init__()

        self.irreps_in = Irreps(irreps_in) if irreps_in is not None else None
        self.irreps_out = Irreps(irreps_out)
        self.irreps_sh = Irreps(irreps_sh)
        self.diameter = diameter
        self.num_radial_basis = num_radial_basis
        self.steps = steps
        self.padding = padding

        with jax.ensure_compile_time_eval():
            r = self.diameter / 2

            s = math.floor(r / self.steps[0])
            x = jnp.arange(-s, s + 1.0) * self.steps[0]

            s = math.floor(r / self.steps[1])
            y = jnp.arange(-s, s + 1.0) * self.steps[1]

            s = math.floor(r / self.steps[2])
            z = jnp.arange(-s, s + 1.0) * self.steps[2]

            self.lattice = jnp.stack(jnp.meshgrid(x, y, z, indexing="ij"), axis=-1)  # [x, y, z, R^3]

            self.sh = spherical_harmonics(
                irreps_out=self.irreps_sh, input=self.lattice, normalize=True, normalization="component"
            )  # [x, y, z, irreps_sh.dim]

    def tp_weight(
        self,
        i_in1: int,
        i_in2: int,
        i_out: int,
        mul_ir_in1: MulIrrep,
        mul_ir_in2: MulIrrep,
        mul_ir_out: MulIrrep,
        path_shape: Tuple[int, ...],
    ) -> jnp.ndarray:
        number = self.num_radial_basis

        with jax.ensure_compile_time_eval():
            embedding = soft_one_hot_linspace(
                jnp.linalg.norm(self.lattice, ord=2, axis=-1),
                start=0.0,
                end=self.diameter / 2,
                number=number,
                basis="smooth_finite",
                start_zero=True,
                end_zero=True,
            )  # [x, y, z, number]

        w = hk.get_parameter(
            f"w[{i_in1},{i_in2},{i_out}] {mul_ir_in1},{mul_ir_in2},{mul_ir_out}",
            (number,) + path_shape,
            init=hk.initializers.RandomNormal(),
        )
        return jnp.einsum("xyzk,k...->xyz...", embedding, w) / (
            self.lattice.shape[0] * self.lattice.shape[1] * self.lattice.shape[2]
        )  # [x, y, z, tp_w]

    def kernel(self, irreps_in: Irreps, irreps_out: Irreps) -> jnp.ndarray:
        # convolution
        tp = FunctionalFullyConnectedTensorProduct(irreps_in, self.irreps_sh, irreps_out)

        ws = [
            self.tp_weight(
                i.i_in1,
                i.i_in2,
                i.i_out,
                tp.irreps_in1[i.i_in1],
                tp.irreps_in2[i.i_in2],
                tp.irreps_out[i.i_out],
                i.path_shape,
            )
            for i in tp.instructions
        ]

        tp_right = tp.right
        for _ in range(3):
            tp_right = jax.vmap(tp_right, (0, 0), 0)
        k = tp_right(ws, self.sh)  # [x, y, z, irreps_in.dim, irreps_out.dim]

        # self-connection
        lin = FunctionalLinear(irreps_in, irreps_out)
        w = [
            hk.get_parameter(
                f"self-connection[{ins.i_in},{ins.i_out}] {lin.irreps_in[ins.i_in]},{lin.irreps_out[ins.i_out]}",
                shape=ins.path_shape,
                init=hk.initializers.RandomNormal(),
            )
            for ins in lin.instructions
        ]
        # note that lattice[center] is always displacement zero
        k = k.at[k.shape[0] // 2, k.shape[1] // 2, k.shape[2] // 2].set(lin.matrix(w))
        return k

    def __call__(self, input: IrrepsData) -> IrrepsData:
        """
        input: [batch, x, y, z, irreps_in.dim]
        """
        if self.irreps_in is not None:
            input = IrrepsData.new(self.irreps_in, input)
        if not isinstance(input, IrrepsData):
            raise ValueError("Convolution: input should be of type IrrepsData")

{'fraction_non_alphanumeric': 0.12262628770013653, 'fraction_numerical': 0.01092217947126722, 'mean_word_length': 1.8047337278106508, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 10, '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': '2878091', 'n_tokens_mistral': 2962, 'n_tokens_neox': 2853, 'n_words': 526}