release: create release-0.2.1

CHANGELOG

@@ -1,5 +1,17 @@
 # Changelog
 
+## Release 0.2.1
+
+- Reducing memory requirements for Hessian predictions in batched inference drastically.
+- Optimizing eSEN implementation, introducing (1) Quaternion method to compute Wigner
+  D matrices, and (2) `use_remat_edgewise` flag to rematerialize `Edgewise`.
+- Improving file reading efficiency in `Hdf5Reader`.
+- Fixing `CombinedGraphDataset` by catching `StopIteration` runtime error.
+- Fixing `MSEHessianLoss` to use correct axes for aggregation.
+- Updating advanced simulation and Hessian tutorial notebooks.
+- Fixing memory bottleneck in ViSNet, adding rematerialisation option with `use_remat`.
+- Upgrading e3j dependency to minimum version of 0.1.0b3.
+
 ## Release 0.2.0
 
 - Introducing breaking API updates in multiple places, see Migration guide in

Dockerfile

@@ -8,7 +8,10 @@ RUN pip install "mlip[cuda13]" huggingface_hub notebook
 
 RUN wget https://raw.githubusercontent.com/instadeepai/mlip/refs/heads/main/tutorials/simulation_tutorial.ipynb \
          https://raw.githubusercontent.com/instadeepai/mlip/refs/heads/main/tutorials/model_training_tutorial.ipynb \
-         https://raw.githubusercontent.com/instadeepai/mlip/refs/heads/main/tutorials/model_addition_tutorial.ipynb
+         https://raw.githubusercontent.com/instadeepai/mlip/refs/heads/main/tutorials/model_addition_tutorial.ipynb \
+         https://raw.githubusercontent.com/instadeepai/mlip/refs/heads/main/tutorials/advanced_simulation_tutorial.ipynb \
+         https://raw.githubusercontent.com/instadeepai/mlip/refs/heads/main/tutorials/moe_training_and_inference_tutorial.ipynb \
+         https://raw.githubusercontent.com/instadeepai/mlip/refs/heads/main/tutorials/hessian_model_training_tutorial.ipynb
 
 EXPOSE 8888
 

README.md

@@ -88,6 +88,9 @@ To install the **CUDA 13 version of JAX and e3j** binaries alongside *mlip*, run
 pip install "mlip[cuda13]"
 ```
 
+The alias `cuda` can be used instead of `cuda13` for backwards
+compatibility with *mlip* v1.
+
 To install the CUDA 12 version instead, run:
 
 ```bash

docs/source/api_reference/data/batching_helpers.rst

@@ -2,7 +2,7 @@
 
 .. module:: mlip.graph.batching_helpers
 
-Batching Helpers
+Batching helpers
 ================
 
 .. autofunction:: batch_graphs

docs/source/api_reference/data/dataset_builder_config.rst

@@ -0,0 +1,8 @@
+.. _dataset_configs:
+
+.. module:: mlip.data.configs
+
+Dataset builder config
+======================
+
+.. autoclass:: GraphDatasetBuilderConfig

docs/source/api_reference/data/index.rst

@@ -16,7 +16,7 @@ This module contains all code related to dataset reading and processing.
     chemical_systems_readers/index
     chemical_system
     dataset_info
-    dataset_configs
+    dataset_builder_config
     graph_dataset
     combined_graph_dataset
     data_prefetching

docs/source/api_reference/models/esen.rst

@@ -1,7 +1,7 @@
 .. _esen:
 
-ESEN
-======
+eSEN
+====
 
 .. module:: mlip.models.esen.network
 

docs/source/installation/index.rst

@@ -19,6 +19,9 @@ Use this command instead to install the GPU-compatible version:
 
 **This command installs the CUDA 13 version of JAX.**
 
+Note that the alias `cuda` can be used instead of `cuda13` for backwards compatibility
+with *mlip* v1.
+
 To install the CUDA 12 version instead, run:
 
 .. code-block:: bash

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mlip"
-version = "0.2.0"
+version = "0.2.1"
 description = "Machine Learning Interatomic Potentials in JAX"
 license-files = [
     "LICENSE"
@@ -40,29 +40,32 @@ dependencies = [
     "h5py >=3.12.1",
     "tqdm >=4.0",
     "tqdm-loggable >=0.2",
-    "e3j >=0.1.0b0, <0.1.1",
+    "e3j >=0.1.0b3, <0.1.1",
 ]
 
 [project.optional-dependencies]
 cuda13 = [
     "jax[cuda13] >= 0.8.0",
     "jaxlib >= 0.8.0",
-    "e3j_ops >= 0.1.0b0, < 0.1.1"
+    "e3j_ops >= 0.1.0b3, < 0.1.1"
+]
+cuda = [
+    "mlip[cuda13]"
 ]
 cuda13_local = [
     "jax[cuda13_local] >= 0.8.0",
     "jaxlib >= 0.8.0",
-    "e3j_ops >= 0.1.0b0, < 0.1.1"
+    "e3j_ops >= 0.1.0b3, < 0.1.1"
 ]
 cuda12 = [
     "jax[cuda12] >=0.4.32",
     "jaxlib >=0.4.32",
-    "e3j_ops >= 0.1.0b0, < 0.1.1"
+    "e3j_ops >= 0.1.0b3, < 0.1.1"
 ]
 cuda12_local = [
     "jax[cuda12_local] >=0.4.32",
     "jaxlib >=0.4.32",
-    "e3j_ops >= 0.1.0b0, < 0.1.1"
+    "e3j_ops >= 0.1.0b3, < 0.1.1"
 ]
 tpu = [
     "jax[tpu]>=0.5.2",

src/mlip/data/chemical_systems_readers/hdf5_reader.py

@@ -44,17 +44,25 @@ def load(self) -> ChemicalSystems:
             all_systems.extend(self._load_single_file(filepath))
         return all_systems
 
+    @staticmethod
+    def _visit_groups(file):
+        groups = []
+
+        def append_group(_, item):
+            if isinstance(item, h5py.Group) and item.parent.name == "/":
+                groups.append(item)
+
+        file.visititems(append_group)
+        return groups
+
     def _read_file(self, filepath: str | os.PathLike) -> ChemicalSystems:
         """Read structures from an HDF5 file and convert each to a
         :class:`~mlip.data.chemical_system.ChemicalSystem`."""
         with h5py.File(filepath, "r") as h5file:
-            struct_names = list(h5file.keys())
+            groups = self._visit_groups(h5file)
             if self.num_to_load:
-                struct_names = struct_names[: self.num_to_load]
-            return [
-                self._hdf5_row_to_chemical_system(h5file[struct_name])
-                for struct_name in struct_names
-            ]
+                groups = groups[: self.num_to_load]
+            return [self._hdf5_row_to_chemical_system(group) for group in groups]
 
     def _hdf5_row_to_chemical_system(
         self,

src/mlip/data/helpers/combined_graph_dataset.py

@@ -185,7 +185,11 @@ def __iter__(self):
             self.state = self.state.replace(
                 num_graphs_processed=self.state.num_graphs_processed + 1
             )
-            yield next(iterators[ds_idx])
+            try:
+                item = next(iterators[ds_idx])
+            except StopIteration:
+                continue
+            yield item
         self.state = self.state.replace(rng=next_rng, num_graphs_processed=jnp.int32(0))
 
     def _get_regular_interleaving_indices(

src/mlip/data/helpers/hessian_utils.py

@@ -240,23 +240,55 @@ def get_hessian_processing_functions() -> tuple[
     return pad_systems_hessians, process_graph_hessian
 
 
-def single_graph_hessian_from_batch(
+def request_all_hessian_rows_batched(batched_graph: Graph) -> Graph:
+    """Set `sample_hessian_rows` to request the full Hessian for all graphs in a batch.
+
+    Rather than jacrev operating on `total_padded_nodes*3` independent outputs, this
+    mimics the subsampled-rows path used in `HessianPredictor`, such that jacrev
+    differentiates `max_n_atoms*3` summed outputs instead, requiring less memory.
+
+    R = max(n_atoms_per_graph) * 3. For graphs with fewer atoms, the extra row slots
+    are redirected to a padding graph, so they contribute nothing to the Jacobian sum.
+
+    Hessians computed with this setting of `sample_hessian_rows` will be of shape
+    (total_nodes, R, 3).
+    """
+    padding_mask = batched_graph.graph_mask()
+    batch_size = batched_graph.num_graphs
+    n_node = batched_graph.n_node[padding_mask]
+
+    R = int(n_node.max()) * 3
+    graph_starts_3 = np.concatenate([[0], np.cumsum(n_node[:-1])]) * 3
+    # Rows beyond the real atom count are redirected to padding graph
+    padding_force_idx = int(n_node.sum()) * 3
+
+    row_offsets = np.arange(R)[None, :]  # (1, R)
+    valid = row_offsets < (n_node * 3)[:, None]  # (n_real_graphs, R)
+    sampled_rows = np.where(
+        valid, graph_starts_3[:, None] + row_offsets, padding_force_idx
+    )
+    # Pad to (batch_size, R) to include a slot for the padding graph
+    sampled_rows = np.pad(sampled_rows, ((0, batch_size - len(n_node)), (0, 0)))
+
+    return batched_graph.replace_globals(sample_hessian_rows=jnp.array(sampled_rows))
+
+
+def single_graph_hessian_from_subsampled_batch(
     batch_hessian: Array, system_start: int, system_end: int
 ) -> Array:
-    """Retrieve the Hessian of a single graph from the batched graph Hessian.
-
-    Args:
-        batch_hessian (Array): Hessian of the whole batched graph.
-        system_start (int): Index of the first node of the system.
-        system_end (int): Index of the last node of the system.
+    """Retrieve a single graph's Hessian from the subsampled-rows batch output.
 
-    Returns:
-    The retrieved Hessian of a single graph from a batched graph Hessian.
+    After `request_all_hessian_rows_batched`, `nodes.hessian` has shape
+    (total_nodes, R, 3). We extract the per-system (n_atoms, 3, n_atoms, 3) Hessian.
     """
-    graph_hessian = batch_hessian[
-        system_start:system_end,
-        :,
-        system_start:system_end,
-        :,
-    ]
-    return graph_hessian
+    n_atoms = system_end - system_start
+
+    # (n_atoms, n_atoms*3, 3): local[col_i, r, col_j] = H_g[r//3, r%3, col_i, col_j]
+    # Rows beyond n_atoms*3 are zero (padding forces) and are excluded with :n_atoms*3
+    local = batch_hessian[system_start:system_end, : n_atoms * 3]
+
+    # Split R=n_atoms*3 -> (n_atoms, 3): axes become (col_i, row_i, row_j, col_j)
+    local = local.reshape(n_atoms, n_atoms, 3, 3)
+
+    # Transpose to (row_i, row_j, col_i, col_j)
+    return local.transpose(1, 2, 0, 3)

src/mlip/data/single_graph_dataset_builder.py

@@ -300,8 +300,13 @@ def _get_median_num_neighbors_and_max_total_edges_long_range(
 
     @staticmethod
     def _find_max_n_node(graphs: list[Graph], batch_size: int) -> int:
-        """Compute `max_n_node`, resizing if the largest graph
-        exceeds the batch capacity."""
+        """Compute `max_n_node`, and resize if largest graph exceeds capacity.
+
+        `max_n_node` is computed as the median number of nodes across all graphs, which
+        is then multiplied by `batch_size` to determine the number of nodes returned in
+        a batch (+1 for padding). If this estimate isn't sufficient to fit the largest
+        graph, the value of `max_n_node` is adjusted to make it sufficient.
+        """
         max_n_node, max_num_atoms = (
             SingleGraphDatasetBuilder._get_median_and_max_n_node(graphs)
         )
@@ -313,14 +318,20 @@ def _find_max_n_node(graphs: list[Graph], batch_size: int) -> int:
 
     @staticmethod
     def _find_max_n_edge(graphs: list[Graph], max_n_node: int, batch_size: int) -> int:
-        """Compute `max_n_edge`, resizing if the largest graph
-        exceeds the batch capacity."""
-        median_n_nei, max_total_edges = (
+        """Compute `max_n_edge`, and resize if largest graph exceeds capacity.
+
+        `max_n_edge` is computed as ceil(median_neighbors_per_node * max_n_node / 2),
+        giving an estimate of the typical edge count for a graph of max_n_node atoms.
+        This is multiplied by batch_size * 2 to determine the total edge capacity of a
+        batch (+1 for padding). If this estimate isn't sufficient to fit the largest
+        graph, the value of `max_n_edge` is adjusted to make it sufficient.
+        """
+        median_n_nei_per_node, max_total_edges = (
             SingleGraphDatasetBuilder._get_median_num_neighbors_and_max_total_edges(
                 graphs
             )
         )
-        max_n_edge = median_n_nei * max_n_node // 2
+        max_n_edge = int(np.ceil(median_n_nei_per_node * max_n_node / 2))
 
         if max_n_edge * batch_size * 2 < max_total_edges:
             logger.debug("Largest graph does not fit into batch -> resizing it.")

src/mlip/inference/batched_inference.py

@@ -25,10 +25,11 @@
     GraphDatasetBuilderConfig,
     SingleGraphDatasetBuilder,
 )
-from mlip.data.helpers.hessian_utils import single_graph_hessian_from_batch
-from mlip.data.helpers.type_aliases import (
-    GraphDatasetLike,
+from mlip.data.helpers.hessian_utils import (
+    request_all_hessian_rows_batched,
+    single_graph_hessian_from_subsampled_batch,
 )
+from mlip.data.helpers.type_aliases import GraphDatasetLike
 from mlip.graph import Graph
 from mlip.models import ForceField
 from mlip.typing import Prediction
@@ -63,12 +64,9 @@ def run_inference_on_a_single_batch(
     batch_hessians = []
     batch_partial_charges = []
 
-    # Requesting full Hessian explicitly will
-    # result in computing full Hessian matrices
-    # in case `HessianPredictor` is used. If a different
-    # predictor class (e.g., ConservativePredictor) is used,
-    # this won't have any impact.
-    batch = batch.request_full_hessian()
+    # Use the row-summing trick: jacrev over n_atoms*3 summed outputs instead of
+    # total_padded_nodes*3 independent outputs, saving ~batch_size backward passes.
+    batch = request_all_hessian_rows_batched(batch)
     output = jitted_force_field_fun(batch)
     mask = batch.graph_mask()
 
@@ -82,7 +80,7 @@ def run_inference_on_a_single_batch(
             batch_energies.append(float(output.energy[i]))
 
             if output.hessian is not None:
-                graph_hessian = single_graph_hessian_from_batch(
+                graph_hessian = single_graph_hessian_from_subsampled_batch(
                     output.hessian, graph_start, graph_end
                 )