Tutorials

.gitignore

@@ -12,3 +12,4 @@ dev_notes/
 .DS_Store
 evals/.venv-torch/
 evals/results/
+.venv_tutorials

docs/guide/tensor_products.md

@@ -34,6 +34,28 @@ output = tp(left_features, edge_attributes)
 Use {class}`~e3nn_mlx.o3.ElementwiseTensorProduct` for aligned channels and
 {class}`~e3nn_mlx.o3.TensorSquare` when both inputs are the same value.
 
+## Symmetry-reduced products
+
+{class}`~e3nn_mlx.o3.ReducedTensorProducts` constructs an orthonormal
+change-of-basis tensor subject to index permutation symmetries such as
+`"ij=ji"` or `"ijk=jik=ikj"`. `filter_ir_mid` restricts every sequential
+Clebsch--Gordan coupling, including the final coupling. `filter_ir_out`
+restricts the retained final irreps. Intermediate filtering prunes contraction
+paths and is not equivalent to slicing an already-constructed output.
+
+```python
+lmax = 4
+bispectrum = o3.ReducedTensorProducts(
+    "ijk=jik=ikj",
+    i=o3.Irreps.spherical_harmonics(lmax),
+    filter_ir_mid=list(o3.Irrep.iterator(lmax=lmax)),
+    filter_ir_out=list(o3.Irrep.iterator(lmax=0)),
+)
+```
+
+This construction retains scalar symmetric triple contractions suitable for a
+bispectrum. 
+
 ## General instructions
 
 The general {class}`~e3nn_mlx.o3.TensorProduct` accepts an explicit output

docs/index.md

@@ -16,6 +16,10 @@ that is not covered by the public API below.
 
 - New to equivariance? Read [Irreducible representations](guide/irreps.md),
   then work through the [first equivariant operation](examples/getting_started.md).
+- Prefer notebook-based introductions? Explore the
+  [adapted historical e3nn tutorials](https://github.com/lamalab-org/e3nn_mlx/tree/main/tutorials)
+  on tensor types, spherical-tensor operations, and invariant atomic
+  descriptors.
 - Coming from e3nn/PyTorch? Start with the [migration guide](guide/migration.md).
 - Building a model? See [tensor products](guide/tensor_products.md), the
   [convolution example](examples/convolution.md), and [point models](examples/point_models.md).

e3nn_mlx/models/v2106/gate_points_message_passing.py

@@ -39,6 +39,8 @@ def __init__(
         irreps_edge_attr,
         fc_neurons,
         num_neighbors: float,
+        *,
+        use_custom_kernel: bool = True,
     ) -> None:
         super().__init__()
         requested_sequence = tuple(
@@ -54,12 +56,14 @@ def __init__(
         if num_neighbors <= 0:
             raise ValueError("num_neighbors must be positive")
         self.num_neighbors = float(num_neighbors)
+        self.use_custom_kernel = bool(use_custom_kernel)
         self._config = {
             "irreps_node_sequence": requested_sequence,
             "irreps_node_attr": self.irreps_node_attr,
             "irreps_edge_attr": self.irreps_edge_attr,
             "fc_neurons": self.fc_neurons,
             "num_neighbors": self.num_neighbors,
+            "use_custom_kernel": self.use_custom_kernel,
         }
 
         mx, _ = require_mlx()
@@ -107,6 +111,7 @@ def __init__(
                 gate.irreps_in,
                 self.fc_neurons,
                 self.num_neighbors,
+                use_custom_kernel=self.use_custom_kernel,
             )
             modules.append(Compose(convolution, gate))
             current = gate.irreps_out
@@ -121,6 +126,7 @@ def __init__(
                 output,
                 self.fc_neurons,
                 self.num_neighbors,
+                use_custom_kernel=self.use_custom_kernel,
             )
         )
         actual_sequence.append(output)

e3nn_mlx/models/v2106/gate_points_networks.py

@@ -58,7 +58,7 @@ def _edge_geometry(self, positions, edge_src, edge_dst):
             edge_vectors,
             normalize=True,
             normalization="component",
-            use_custom_kernel=False,
+            use_custom_kernel=self.use_custom_kernel,
         )
         edge_lengths = mx.sqrt(mx.sum(edge_vectors * edge_vectors, axis=-1))
         edge_scalars = soft_one_hot_linspace(
@@ -75,7 +75,12 @@ def _finish(self, output: IrrepsArray, batch, num_graphs: int) -> IrrepsArray:
         if num_graphs < 0:
             raise ValueError("num_graphs must be non-negative")
         if self.pool_nodes:
-            pooled = scatter_sum(output.array, batch, num_graphs) / sqrt(self.num_nodes)
+            pooled = scatter_sum(
+                output.array,
+                batch,
+                num_graphs,
+                use_custom_kernel=self.use_custom_kernel,
+            ) / sqrt(self.num_nodes)
             return IrrepsArray(self.irreps_node_output, pooled)
         return output
 
@@ -103,6 +108,8 @@ def __init__(
         layers: int = 3,
         lmax: int = 2,
         pool_nodes: bool = True,
+        *,
+        use_custom_kernel: bool = True,
     ) -> None:
         super().__init__()
         if max_radius <= 0 or num_neighbors <= 0 or num_nodes <= 0:
@@ -112,6 +119,7 @@ def __init__(
         self.max_radius = float(max_radius)
         self.num_nodes = float(num_nodes)
         self.pool_nodes = bool(pool_nodes)
+        self.use_custom_kernel = bool(use_custom_kernel)
         self.lmax = int(lmax)
         self.irreps_sh = Irreps.spherical_harmonics(self.lmax)
         self._config = {
@@ -124,6 +132,7 @@ def __init__(
             "layers": int(layers),
             "lmax": self.lmax,
             "pool_nodes": self.pool_nodes,
+            "use_custom_kernel": self.use_custom_kernel,
         }
         hidden = _hidden_irreps(int(mul), self.lmax)
         self.mp = MessagePassing(
@@ -136,6 +145,7 @@ def __init__(
             self.irreps_sh,
             [self.number_of_basis, 100],
             num_neighbors,
+            use_custom_kernel=self.use_custom_kernel,
         )
         self.irreps_in = self.mp.irreps_node_input
         self.irreps_out = self.mp.irreps_node_output
@@ -222,6 +232,8 @@ def __init__(
         layers: int = 3,
         lmax: int = 2,
         pool_nodes: bool = True,
+        *,
+        use_custom_kernel: bool = True,
     ) -> None:
         super().__init__()
         if max_radius <= 0 or num_neighbors <= 0 or num_nodes <= 0:
@@ -231,6 +243,7 @@ def __init__(
         self.max_radius = float(max_radius)
         self.num_nodes = float(num_nodes)
         self.pool_nodes = bool(pool_nodes)
+        self.use_custom_kernel = bool(use_custom_kernel)
         self.lmax = int(lmax)
         self.irreps_edge_attr = Irreps(irreps_edge_attr).remove_zero_multiplicities()
         self.irreps_sh = Irreps.spherical_harmonics(self.lmax)
@@ -247,6 +260,7 @@ def __init__(
             "layers": int(layers),
             "lmax": self.lmax,
             "pool_nodes": self.pool_nodes,
+            "use_custom_kernel": self.use_custom_kernel,
         }
         hidden = _hidden_irreps(int(mul), self.lmax)
         self.mp = MessagePassing(
@@ -259,6 +273,7 @@ def __init__(
             self.irreps_edge_attr_combined,
             [self.number_of_basis, 100],
             num_neighbors,
+            use_custom_kernel=self.use_custom_kernel,
         )
         self.irreps_node_input = self.mp.irreps_node_input
         self.irreps_node_attr = self.mp.irreps_node_attr

e3nn_mlx/models/v2106/points_convolution.py

@@ -159,7 +159,12 @@ def __call__(
             gathered = IrrepsArray(self.tp.irreps_in1, transformed.array[edge_src])
             tp_edge_attr = IrrepsArray(self.tp.irreps_in2, edge_attr.array)
             messages = self.tp(gathered, tp_edge_attr, weights)
-            aggregated = scatter_sum(messages.array, edge_dst, node_input.shape[0])
+            aggregated = scatter_sum(
+                messages.array,
+                edge_dst,
+                node_input.shape[0],
+                use_custom_kernel=self.use_custom_kernel,
+            )
         aggregated = aggregated / sqrt(self.num_neighbors)
 
         middle = IrrepsArray(self.irreps_mid, aggregated)

e3nn_mlx/ops_reduce_tensor.py

@@ -65,13 +65,19 @@ def _irrep_embeddings(irreps: Irreps, np) -> list[tuple[Irrep, Any]]:
     return embeddings
 
 
-def _couple_channels(factors: tuple[Irreps, ...], np) -> list[_CoupledChannel]:
+def _couple_channels(
+    factors: tuple[Irreps, ...],
+    np,
+    filter_ir_mid: frozenset[Irrep] | None = None,
+) -> list[_CoupledChannel]:
     channels = [_CoupledChannel(irrep, embedding) for irrep, embedding in _irrep_embeddings(factors[0], np)]
     for factor in factors[1:]:
         next_channels = []
         for channel in channels:
             for right_irrep, right_basis in _irrep_embeddings(factor, np):
                 for output_irrep in channel.irrep * right_irrep:
+                    if filter_ir_mid is not None and output_irrep not in filter_ir_mid:
+                        continue
                     coefficient = np.asarray(
                         clebsch_gordan(channel.irrep, right_irrep, output_irrep), dtype=np.float64
                     )
@@ -95,14 +101,28 @@ def _nullspace(matrix, np, tolerance: float = 1e-9):
 def _build_change_of_basis(
     formula: str,
     factor_specs: tuple[str, ...],
+    filter_ir_mid_specs: tuple[str, ...] | None = None,
+    filter_ir_out_specs: tuple[str, ...] | None = None,
 ) -> tuple[Irreps, Any]:
     import numpy as np
 
     base, constraints = _parse_formula(formula)
     factors = tuple(Irreps(spec) for spec in factor_specs)
-    channels = _couple_channels(factors, np)
+    filter_ir_mid = (
+        None
+        if filter_ir_mid_specs is None
+        else frozenset(Irrep.parse(spec) for spec in filter_ir_mid_specs)
+    )
+    filter_ir_out = (
+        None
+        if filter_ir_out_specs is None
+        else frozenset(Irrep.parse(spec) for spec in filter_ir_out_specs)
+    )
+    channels = _couple_channels(factors, np, filter_ir_mid)
     grouped: dict[Irrep, list[Any]] = {}
     for channel in channels:
+        if filter_ir_out is not None and channel.irrep not in filter_ir_out:
+            continue
         grouped.setdefault(channel.irrep, []).append(channel.basis)
 
     output_parts = []
@@ -134,15 +154,94 @@ def _build_change_of_basis(
 
 
 class ReducedTensorProducts(mlx_module_base()):
-    """Reduce a tensor product according to a permutation-symmetry formula."""
+    """Reduce a tensor product according to permutation symmetries.
+
+    Parameters
+    ----------
+    formula
+        Index formula describing the tensor symmetry. For example, ``"ij=ji"``
+        selects symmetric rank-two tensors and ``"ijk=jik=ikj"`` selects fully
+        symmetric rank-three tensors.
+    filter_ir_out
+        Optional iterable of allowed final irreps. Channels with other output
+        irreps are omitted from :attr:`irreps_out` and
+        :attr:`change_of_basis`.
+    filter_ir_mid
+        Optional iterable of irreps allowed at every sequential coupling stage,
+        including the final coupling. This prunes contraction paths before the
+        permutation-symmetry reduction and can substantially reduce basis
+        construction work.
+    **irreps
+        Irreps associated with formula indices. If only one index is supplied,
+        its irreps are used for every index in the formula.
+
+    Notes
+    -----
+    ``filter_ir_mid`` is a path constraint, not merely an output slice. It can
+    change which multiplicity channels survive the symmetry reduction.
+
+    Examples
+    --------
+    Construct scalar bispectrum contractions from spherical coefficients:
+
+    >>> from e3nn_mlx import o3
+    >>> lmax = 4
+    >>> bispectrum = o3.ReducedTensorProducts(
+    ...     "ijk=jik=ikj",
+    ...     i=o3.Irreps.spherical_harmonics(lmax),
+    ...     filter_ir_mid=list(o3.Irrep.iterator(lmax=lmax)),
+    ...     filter_ir_out=list(o3.Irrep.iterator(lmax=0)),
+    ... )
+    >>> all(part.ir.l == 0 for part in bispectrum.irreps_out)
+    True
+    """
 
-    def __init__(self, formula: str, **irreps: Irreps | str) -> None:
+    def __init__(
+        self,
+        formula: str,
+        filter_ir_out: list[Irrep | str] | None = None,
+        filter_ir_mid: list[Irrep | str] | None = None,
+        **irreps: Irreps | str,
+    ) -> None:
         super().__init__()
         base, _ = _parse_formula(formula)
         self.formula = formula
         self._labels = base
         self.irreps_in = _factor_irreps(base, irreps)
-        self.irreps_out, data = _build_change_of_basis(formula, tuple(str(value) for value in self.irreps_in))
+        try:
+            self.filter_ir_out = (
+                None
+                if filter_ir_out is None
+                else tuple(Irrep.parse(value) for value in filter_ir_out)
+            )
+        except (TypeError, ValueError) as error:
+            raise ValueError(
+                f"filter_ir_out (={filter_ir_out}) must be an iterable of Irrep values"
+            ) from error
+        try:
+            self.filter_ir_mid = (
+                None
+                if filter_ir_mid is None
+                else tuple(Irrep.parse(value) for value in filter_ir_mid)
+            )
+        except (TypeError, ValueError) as error:
+            raise ValueError(
+                f"filter_ir_mid (={filter_ir_mid}) must be an iterable of Irrep values"
+            ) from error
+        self.irreps_out, data = _build_change_of_basis(
+            formula,
+            tuple(str(value) for value in self.irreps_in),
+            (
+                None
+                if self.filter_ir_mid is None
+                else tuple(str(value) for value in self.filter_ir_mid)
+            ),
+            (
+                None
+                if self.filter_ir_out is None
+                else tuple(str(value) for value in self.filter_ir_out)
+            ),
+        )
         mx, _ = require_mlx()
         self._change_of_basis = mx.array(data, dtype=mx.float32)
 

evals/mlx_cases.py

@@ -384,6 +384,7 @@ def build_v2106_convolution(config: dict[str, Any]) -> Task:
         irreps_node,
         [config["radial"], config["radial_hidden"]],
         float(config["neighbors"]),
+        use_custom_kernel=_USE_CUSTOM_KERNELS,
     )
     module.alpha.update({"weight": 0.05 * mx.random.normal(shape=module.alpha.weight.shape)})
     node_input = mx.random.normal(shape=(config["nodes"], module.irreps_node_input.dim))
@@ -420,6 +421,7 @@ def build_v2106_message_passing(config: dict[str, Any]) -> Task:
         irreps_edge,
         [config["radial"], config["radial_hidden"]],
         float(config["neighbors"]),
+        use_custom_kernel=_USE_CUSTOM_KERNELS,
     )
     _activate_alpha(module)
     node_input = mx.random.normal(shape=(config["nodes"], module.irreps_node_input.dim))
@@ -462,6 +464,7 @@ def build_v2106_network(config: dict[str, Any]) -> Task:
         layers=config["layers"],
         lmax=config["lmax"],
         pool_nodes=True,
+        use_custom_kernel=_USE_CUSTOM_KERNELS,
     )
     _activate_alpha(module)
     positions = mx.random.normal(shape=(config["nodes"], 3))

tests/test_documentation_contract.py

@@ -69,6 +69,8 @@
         "tests/test_upstream_o3_tensor_product.py::test_upstream_tensor_product_unshared_weight_broadcast_and_validation",
     "tensor-product wrappers match explicit contractions":
         "tests/test_tensor_product_correctness.py::test_full_tensor_product_matches_explicit_tensor_product",
+    "reduced-tensor intermediate and output filters":
+        "tests/test_upstream_o3_reduce_tensor.py::test_upstream_reduced_tensor_supports_intermediate_and_output_filters",
     "lazy package import without MLX":
         "tests/test_nn_import_policy.py::test_nn_namespace_imports_without_site_packages",
     "gate-points exact documented configuration":
@@ -176,6 +178,11 @@
         "`e3nn_core` remains importable without MLX.",
         "tests/test_nn_import_policy.py::test_nn_namespace_imports_without_site_packages",
     ),
+    (
+        "docs/guide/tensor_products.md",
+        "`filter_ir_mid` restricts every sequential Clebsch--Gordan coupling, including the final coupling.",
+        "tests/test_upstream_o3_reduce_tensor.py::test_upstream_reduced_tensor_supports_intermediate_and_output_filters",
+    ),
     (
         "docs/HIGH_LEVEL_API.md",
         "High-level modules preserve the caller's style",