k3-node 1.0.0__py3-none-any.whl

k3_node/applications/materials/io.py

@@ -0,0 +1,230 @@
+"""Checkpoint downloading and pretrained weight loading utilities for MatGL models."""
+
+from __future__ import annotations
+
+import os
+import json
+import logging
+from pathlib import Path
+from typing import Optional, Union, Dict, Any, List
+import numpy as np
+import torch
+import keras
+
+logger = logging.getLogger(__name__)
+
+HF_MATGL_ORG = "materialyze"
+
+KNOWN_PRETRAINED_MODELS = [
+    "CHGNet-PES-MatPES-PBE-2025.2.10",
+    "CHGNet-PES-MatPES-r2SCAN-2025.2.10",
+    "M3GNet-Eform-MP-2018.6.1",
+    "M3GNet-Eform-MP-2019.4.1",
+    "M3GNet-PES-ANI-1x-Subset",
+    "M3GNet-PES-MatPES-PBE-2025.2",
+    "M3GNet-PES-MatPES-r2SCAN-2025.2",
+    "MEGNet-BandGap-mfi-MP-2019.4.1",
+    "MEGNet-Eform-MP-2018.6.1",
+    "QET-PES-MatPES-PBE-2025.2",
+    "QET-PES-MatPES-r2SCAN-2025.2",
+    "QET-PES-MatQ",
+    "SO3Net-PES-ANI-1x-Subset",
+    "TensorNet-PES-ANI-1x-Subset",
+    "TensorNet-PES-MatPES-PBE-2025.2",
+    "TensorNet-PES-MatPES-PBE-2025.2-m",
+    "TensorNet-PES-MatPES-r2SCAN-2025.2",
+    "TensorNet-PES-MatPES-r2SCAN-2025.2-m",
+]
+
+
+def get_available_pretrained_models() -> List[str]:
+    """Return list of available pretrained materials models."""
+    try:
+        from huggingface_hub import HfApi
+        api = HfApi()
+        names = []
+        for model_info in api.list_models(author=HF_MATGL_ORG):
+            repo_id = str(getattr(model_info, "id", "") or getattr(model_info, "modelId", "") or "")
+            if "/" in repo_id:
+                names.append(repo_id.split("/", 1)[1])
+        if names:
+            return sorted(names)
+    except Exception as e:
+        logger.debug("Hugging Face API listing unavailable, returning static list: %s", e)
+    return sorted(KNOWN_PRETRAINED_MODELS)
+
+
+def download_matgl_checkpoint(
+    name_or_repo: str,
+    folder: str = "checkpoints",
+    log: bool = True,
+) -> Dict[str, str]:
+    """Download matgl model files (model.json, state.pt) from Hugging Face Hub."""
+    from huggingface_hub import hf_hub_download
+
+    repo_id = name_or_repo if "/" in name_or_repo else f"{HF_MATGL_ORG}/{name_or_repo}"
+    if log:
+        print(f"Downloading checkpoint files for {repo_id}...")
+
+    os.makedirs(folder, exist_ok=True)
+    f_json = hf_hub_download(repo_id=repo_id, filename="model.json", cache_dir=folder)
+    f_state = hf_hub_download(repo_id=repo_id, filename="state.pt", cache_dir=folder)
+
+    return {
+        "model.json": f_json,
+        "state.pt": f_state,
+        "repo_id": repo_id,
+    }
+
+
+def _assign_dense_weights(keras_layer, weight_tensor, bias_tensor=None):
+    """Assign PyTorch Linear weight/bias to Keras Dense layer."""
+    w_np = weight_tensor.detach().cpu().numpy()
+    if len(w_np.shape) == 2:
+        # PyTorch is [out_features, in_features], Keras is [in_features, out_features]
+        w_np = np.transpose(w_np, (1, 0))
+    weights = [w_np]
+    if bias_tensor is not None:
+        b_np = bias_tensor.detach().cpu().numpy()
+        weights.append(b_np)
+    keras_layer.set_weights(weights)
+
+
+def load_matgl_weights(
+    model: keras.Model,
+    state_dict_or_path: Union[str, Dict[str, torch.Tensor]],
+    log: bool = False,
+) -> int:
+    """Load PyTorch checkpoint weights into multi-backend Keras 3 MatGL model."""
+    if isinstance(state_dict_or_path, (str, Path)):
+        state = torch.load(state_dict_or_path, map_location="cpu", weights_only=True)
+    else:
+        state = state_dict_or_path
+
+    # Strip potential 'model.' prefix
+    cleaned_state = {}
+    for k, v in state.items():
+        if k.startswith("model."):
+            cleaned_state[k[6:]] = v
+        else:
+            cleaned_state[k] = v
+
+    loaded_count = 0
+
+    # 1. Bond expansion centers/width if present
+    if hasattr(model, "bond_expansion") and hasattr(model.bond_expansion, "rbf"):
+        rbf = model.bond_expansion.rbf
+        for k in ("bond_expansion.rbf.centers", "bond_expansion.centers"):
+            if k in cleaned_state and hasattr(rbf, "centers"):
+                rbf.centers.assign(cleaned_state[k].detach().cpu().numpy())
+                loaded_count += 1
+        for k in ("bond_expansion.rbf.width", "bond_expansion.width"):
+            if k in cleaned_state and hasattr(rbf, "width"):
+                rbf.width.assign(cleaned_state[k].detach().cpu().numpy())
+                loaded_count += 1
+
+    # 2. Embedding block
+    if hasattr(model, "embedding"):
+        emb = model.embedding
+        if hasattr(emb, "layer_node_embedding") and "embedding.layer_node_embedding.weight" in cleaned_state:
+            w = cleaned_state["embedding.layer_node_embedding.weight"].detach().cpu().numpy()
+            emb.layer_node_embedding.set_weights([w])
+            loaded_count += 1
+        if hasattr(emb, "emb") and "embedding.emb.weight" in cleaned_state:
+            w = cleaned_state["embedding.emb.weight"].detach().cpu().numpy()
+            emb.emb.set_weights([w])
+            loaded_count += 1
+
+    # 3. Traverse model sublayers and match with state keys
+    for name, sublayer in model.__dict__.items():
+        if isinstance(sublayer, keras.layers.Layer):
+            # Check for direct linear weight matches
+            w_key = f"{name}.weight"
+            b_key = f"{name}.bias"
+            if w_key in cleaned_state:
+                _assign_dense_weights(sublayer, cleaned_state[w_key], cleaned_state.get(b_key))
+                loaded_count += 1
+
+    if log:
+        print(f"Loaded {loaded_count} weight tensors into {model.__class__.__name__}")
+
+    return loaded_count
+
+
+def load_model(name_or_path: str, **kwargs) -> keras.Model:
+    """Convenience factory to download/load and instantiate any MatGL model."""
+    from .megnet import MEGNet
+    from .m3gnet import M3GNet
+    from .tensornet import TensorNet
+    from .chgnet import CHGNet
+    from .so3net import SO3Net
+    from .grace import GRACE
+    from .qet import QET
+    from .wrappers import TransformedTargetModel
+
+    if os.path.exists(name_or_path) and os.path.isdir(name_or_path):
+        f_json = os.path.join(name_or_path, "model.json")
+        f_state = os.path.join(name_or_path, "state.pt")
+    else:
+        files = download_matgl_checkpoint(name_or_path, **kwargs)
+        f_json = files["model.json"]
+        f_state = files["state.pt"]
+
+    with open(f_json) as f:
+        meta = json.load(f)
+
+    cls_name = meta.get("@class")
+    init_kwargs = meta.get("kwargs", {})
+
+    # Handle TransformedTargetModel
+    if cls_name == "TransformedTargetModel":
+        inner_info = init_kwargs.get("model", {})
+        inner_cls = inner_info.get("@class", "MEGNet")
+        inner_args = inner_info.get("init_args", {})
+        transformer_info = init_kwargs.get("target_transformer", {})
+        mean = float(transformer_info.get("mean", 0.0))
+        std = float(transformer_info.get("std", 1.0))
+
+        if inner_cls == "MEGNet":
+            base_model = MEGNet(**{k: v for k, v in inner_args.items() if k in (
+                "dim_node_embedding", "dim_edge_embedding", "dim_state_embedding",
+                "nblocks", "cutoff"
+            )})
+        elif inner_cls == "M3GNet":
+            base_model = M3GNet(**{k: v for k, v in inner_args.items() if k in (
+                "dim_node_embedding", "dim_edge_embedding", "nblocks", "cutoff", "threebody_cutoff"
+            )})
+        elif inner_cls == "TensorNet":
+            base_model = TensorNet(**{k: v for k, v in inner_args.items() if k in (
+                "units", "nblocks", "num_rbf", "cutoff"
+            )})
+        elif inner_cls == "CHGNet":
+            base_model = CHGNet(**{k: v for k, v in inner_args.items() if k in (
+                "dim_atom_embedding", "dim_bond_embedding", "cutoff", "threebody_cutoff"
+            )})
+        else:
+            base_model = MEGNet()
+
+        load_matgl_weights(base_model, f_state)
+        return TransformedTargetModel(model=base_model, mean=mean, std=std)
+
+    elif cls_name == "MEGNet":
+        model = MEGNet()
+    elif cls_name == "M3GNet":
+        model = M3GNet()
+    elif cls_name == "TensorNet":
+        model = TensorNet()
+    elif cls_name == "CHGNet":
+        model = CHGNet()
+    elif cls_name == "SO3Net":
+        model = SO3Net()
+    elif cls_name == "GRACE":
+        model = GRACE()
+    elif cls_name == "QET":
+        model = QET()
+    else:
+        model = MEGNet()
+
+    load_matgl_weights(model, f_state)
+    return model
+

k3_node/applications/materials/m3gnet.py

@@ -0,0 +1,462 @@
+"""Multi-backend Keras 3 implementation of M3GNet."""
+
+from __future__ import annotations
+
+from typing import Sequence, Optional, Union, Tuple, Dict, Any, Literal
+import keras
+from keras import layers, ops
+import numpy as np
+
+from .core import MLP, GatedMLP, EmbeddingBlock, scatter_add, scatter_mean, infer_num_graphs
+from .basis import (
+    BondExpansion,
+    SphericalBesselWithHarmonics,
+    compute_pair_vector_and_distance,
+    compute_theta_and_phi,
+    cosine_cutoff,
+)
+from .readout import WeightedAtomReadOut, ReduceReadOut, Set2SetReadOut
+
+
+class ThreeBodyInteractions(layers.Layer):
+    """Three-body bond angular update using directed line-graph message passing.
+
+    Example:
+        ```python
+        import numpy as np
+        import keras
+        from k3_node.models import ThreeBodyInteractions
+
+        edge_index = np.array([[0, 1, 1, 2, 2, 3, 3, 0], [1, 0, 2, 1, 3, 2, 0, 3]])  # 4 atoms, 8 bonds
+        node_feat = np.random.rand(4, 16).astype("float32")
+        edge_feat = np.random.rand(8, 16).astype("float32")
+        line_edge_index = np.array([[0, 1, 2, 3], [1, 2, 3, 0]])  # bond pairs forming 4 angles
+        three_basis = np.random.rand(4, 9).astype("float32")  # angular basis of each triplet
+        three_cutoff = np.ones(8, dtype="float32")  # smooth cutoff weight of each bond
+
+        layer = ThreeBodyInteractions(
+            update_network_atom=keras.layers.Dense(9, activation="sigmoid"),
+            update_network_bond=keras.layers.Dense(16),
+        )
+        print(tuple(layer(edge_index, line_edge_index, three_basis, three_cutoff, node_feat, edge_feat).shape))  # (8, 16)
+        ```
+    """
+
+    def __init__(
+        self,
+        update_network_atom: layers.Layer,
+        update_network_bond: layers.Layer,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.update_network_atom = update_network_atom
+        self.update_network_bond = update_network_bond
+
+    def call(
+        self,
+        edge_index,
+        line_edge_index,
+        three_basis,
+        three_cutoff,
+        node_feat,
+        edge_feat,
+    ):
+        src_bond = ops.cast(line_edge_index[0], "int32")
+        dst_bond = ops.cast(line_edge_index[1], "int32")
+        edge_dst_atom = ops.cast(edge_index[1], "int32")
+        num_bonds = ops.shape(edge_feat)[0]
+
+        # Destination atom of destination bond
+        end_atom_idx = ops.take(edge_dst_atom, dst_bond, axis=0)
+        updated_atoms = self.update_network_atom(node_feat)
+        end_atom_features = ops.take(updated_atoms, end_atom_idx, axis=0)
+
+        basis = three_basis * end_atom_features
+        weights = ops.take(three_cutoff, src_bond, axis=0) * ops.take(three_cutoff, dst_bond, axis=0)
+        basis = basis * ops.expand_dims(weights, axis=-1)
+
+        new_bonds = scatter_add(basis, src_bond, num_segments=num_bonds)
+        return edge_feat + self.update_network_bond(new_bonds)
+
+
+class M3GNetGraphConv(layers.Layer):
+    """M3GNet graph convolution layer: two-body edge and node updates.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import M3GNetGraphConv
+
+        edge_index = np.array([[0, 1, 1, 2, 2, 3, 3, 0], [1, 0, 2, 1, 3, 2, 0, 3]])  # 4 atoms, 8 bonds
+        node_feat = np.random.rand(4, 16).astype("float32")
+        edge_feat = np.random.rand(8, 16).astype("float32")
+        rbf = np.random.rand(8, 9).astype("float32")  # radial basis of each bond
+
+        conv = M3GNetGraphConv(degree=9, edge_dims=[48, 16, 16], node_dims=[48, 16, 16])
+        edge_out, node_out, state_out = conv(edge_index, edge_feat, node_feat, None, rbf, num_nodes=4)
+        print(tuple(edge_out.shape), tuple(node_out.shape))  # (8, 16) (4, 16)
+        ```
+    """
+
+    def __init__(
+        self,
+        degree: int,
+        edge_dims: Sequence[int],
+        node_dims: Sequence[int],
+        state_dims: Optional[Sequence[int]] = None,
+        include_state: bool = False,
+        activation: str = "swish",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.include_state = include_state
+        self.edge_update_func = GatedMLP(in_feats=edge_dims[0], dims=edge_dims[1:])
+        self.edge_weight_func = layers.Dense(edge_dims[-1], use_bias=False)
+
+        self.node_update_func = GatedMLP(in_feats=node_dims[0], dims=node_dims[1:])
+        self.node_weight_func = layers.Dense(node_dims[-1], use_bias=False)
+
+        if include_state and state_dims is not None:
+            self.state_update_func = MLP(state_dims, activation=activation, activate_last=True)
+        else:
+            self.state_update_func = None
+
+    def call(
+        self,
+        edge_index,
+        edge_feat,
+        node_feat,
+        state_feat,
+        rbf,
+        batch=None,
+        edge_batch=None,
+        num_nodes: Optional[int] = None,
+        num_graphs: Optional[int] = None,
+    ):
+        src = ops.cast(edge_index[0], "int32")
+        dst = ops.cast(edge_index[1], "int32")
+
+        vi = ops.take(node_feat, src, axis=0)
+        vj = ops.take(node_feat, dst, axis=0)
+
+        if self.include_state and state_feat is not None:
+            if edge_batch is not None:
+                u_edge = ops.take(state_feat, edge_batch, axis=0)
+            else:
+                num_e = ops.shape(edge_feat)[0]
+                u_edge = ops.broadcast_to(state_feat, (num_e, ops.shape(state_feat)[-1]))
+            edge_inputs = ops.concatenate([vi, vj, edge_feat, u_edge], axis=-1)
+        else:
+            edge_inputs = ops.concatenate([vi, vj, edge_feat], axis=-1)
+
+        # 1. Edge update
+        edge_update = self.edge_update_func(edge_inputs) * self.edge_weight_func(rbf)
+        edge_feat_new = edge_feat + edge_update
+
+        # 2. Node update
+        node_update = self.node_update_func(edge_inputs) * self.node_weight_func(rbf)
+        node_update_sum = scatter_add(node_update, src, num_segments=num_nodes)
+        node_feat_new = node_feat + node_update_sum
+
+        # 3. State update
+        state_feat_new = state_feat
+        if self.include_state and self.state_update_func is not None and state_feat is not None:
+            if batch is not None:
+                uv = scatter_mean(node_feat_new, batch, num_segments=num_graphs)
+            else:
+                uv = ops.mean(node_feat_new, axis=0, keepdims=True)
+            state_inputs = ops.concatenate([state_feat, uv], axis=-1)
+            state_feat_new = self.state_update_func(state_inputs)
+
+        return edge_feat_new, node_feat_new, state_feat_new
+
+
+class M3GNetBlock(layers.Layer):
+    """M3GNet block wrapping M3GNetGraphConv with optional dropout.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import M3GNetBlock
+
+        edge_index = np.array([[0, 1, 1, 2, 2, 3, 3, 0], [1, 0, 2, 1, 3, 2, 0, 3]])  # 4 atoms, 8 bonds
+        node_feat = np.random.rand(4, 16).astype("float32")
+        edge_feat = np.random.rand(8, 16).astype("float32")
+        rbf = np.random.rand(8, 9).astype("float32")  # radial basis of each bond
+
+        block = M3GNetBlock(degree=9, conv_hiddens=[16], dim_node_feats=16, dim_edge_feats=16)
+        edge_out, node_out, state_out = block(edge_index, edge_feat, node_feat, None, rbf, num_nodes=4)
+        print(tuple(edge_out.shape), tuple(node_out.shape))  # (8, 16) (4, 16)
+        ```
+    """
+
+    def __init__(
+        self,
+        degree: int,
+        conv_hiddens: Sequence[int],
+        dim_node_feats: int,
+        dim_edge_feats: int,
+        dim_state_feats: int = 0,
+        include_state: bool = False,
+        activation: str = "swish",
+        dropout: float = 0.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.include_state = include_state
+        edge_in = 2 * dim_node_feats + dim_edge_feats + (dim_state_feats if include_state else 0)
+        node_in = 2 * dim_node_feats + dim_edge_feats + (dim_state_feats if include_state else 0)
+        state_in = dim_state_feats + dim_node_feats if include_state else 0
+
+        self.conv = M3GNetGraphConv(
+            degree=degree,
+            edge_dims=[edge_in, *conv_hiddens, dim_edge_feats],
+            node_dims=[node_in, *conv_hiddens, dim_node_feats],
+            state_dims=[state_in, *conv_hiddens, dim_state_feats] if include_state else None,
+            include_state=include_state,
+            activation=activation,
+        )
+        self.dropout = layers.Dropout(dropout) if dropout > 0.0 else None
+
+    def call(
+        self,
+        edge_index,
+        edge_feat,
+        node_feat,
+        state_feat,
+        rbf,
+        batch=None,
+        edge_batch=None,
+        num_nodes: Optional[int] = None,
+        num_graphs: Optional[int] = None,
+        training=None,
+    ):
+        edge_feat, node_feat, state_feat = self.conv(
+            edge_index, edge_feat, node_feat, state_feat, rbf,
+            batch=batch, edge_batch=edge_batch, num_nodes=num_nodes, num_graphs=num_graphs
+        )
+        if self.dropout is not None:
+            edge_feat = self.dropout(edge_feat, training=training)
+            node_feat = self.dropout(node_feat, training=training)
+        return edge_feat, node_feat, state_feat
+
+
+class M3GNet(keras.Model):
+    """M3GNet materials potential model supporting 3-body angles and multibackend training.
+
+    Example:
+        ```python
+        import numpy as np
+        from k3_node.models import M3GNet
+
+        # A 4-atom structure: positions, bonds (listed in both directions) and atomic numbers
+        structure = {
+            "pos": np.array([[0.0, 0.0, 0.0], [1.0, 0.5, 0.0], [0.5, 1.2, 0.8], [1.5, 1.5, 1.0]], dtype="float32"),
+            "edge_index": np.array([[0, 1, 1, 2, 2, 3, 3, 0], [1, 0, 2, 1, 3, 2, 0, 3]]),
+            "line_edge_index": np.array([[0, 1, 2, 3], [1, 2, 3, 0]]),  # bond pairs forming angles
+            "node_type": np.array([6, 8, 1, 6]),  # atomic numbers
+            "batch": np.zeros(4, dtype="int32"),  # all atoms belong to structure 0
+            "state_attr": np.zeros((1, 2), dtype="float32"),  # global state features
+        }
+
+        model = M3GNet(dim_node_embedding=16, dim_edge_embedding=16, nblocks=2, units=16, max_n=3, max_l=3)
+        energy = model(structure)  # predicted property (e.g. energy) of the structure
+        print(tuple(energy.shape))  # (1,)
+        ```
+    """
+
+    def __init__(
+        self,
+        dim_node_embedding: int = 64,
+        dim_edge_embedding: int = 64,
+        dim_state_embedding: int = 0,
+        ntypes_state: Optional[int] = None,
+        max_n: int = 3,
+        max_l: int = 3,
+        nblocks: int = 3,
+        rbf_type: Literal["Gaussian", "SphericalBessel"] = "SphericalBessel",
+        is_intensive: bool = True,
+        readout_type: Literal["set2set", "weighted_atom", "reduce_atom"] = "weighted_atom",
+        cutoff: float = 5.0,
+        threebody_cutoff: float = 4.0,
+        units: int = 64,
+        ntargets: int = 1,
+        include_state: bool = False,
+        activation_type: str = "swish",
+        dropout: float = 0.0,
+        ntypes_node: int = 95,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.cutoff = cutoff
+        self.threebody_cutoff = threebody_cutoff
+        self.include_state = include_state
+        self.is_intensive = is_intensive
+
+        self.bond_expansion = BondExpansion(
+            max_l=max_l,
+            max_n=max_n,
+            cutoff=cutoff,
+            rbf_type=rbf_type,
+            smooth=False,
+        )
+
+        degree_rbf = max_n * max_l if rbf_type.lower() == "sphericalbessel" else 100
+        self.embedding = EmbeddingBlock(
+            degree_rbf=degree_rbf,
+            dim_node_embedding=dim_node_embedding,
+            dim_edge_embedding=dim_edge_embedding,
+            ntypes_node=ntypes_node,
+            ntypes_state=ntypes_state,
+            include_state=include_state,
+            dim_state_embedding=dim_state_embedding if include_state else None,
+            activation=activation_type,
+        )
+
+        self.sbf_shf = SphericalBesselWithHarmonics(
+            max_n=max_n,
+            max_l=max_l,
+            cutoff=threebody_cutoff,
+        )
+        degree_3body = max_n * max_l
+
+        self.three_body_interactions = []
+        self.blocks = []
+        for _ in range(nblocks):
+            self.three_body_interactions.append(
+                ThreeBodyInteractions(
+                    update_network_atom=MLP([dim_node_embedding, degree_3body], activation=activation_type, activate_last=False),
+                    update_network_bond=GatedMLP(in_feats=degree_3body, dims=[dim_edge_embedding], activate_last=False),
+                )
+            )
+            self.blocks.append(
+                M3GNetBlock(
+                    degree=degree_rbf,
+                    conv_hiddens=[units, units],
+                    dim_node_feats=dim_node_embedding,
+                    dim_edge_feats=dim_edge_embedding,
+                    dim_state_feats=dim_state_embedding if include_state else 0,
+                    include_state=include_state,
+                    activation=activation_type,
+                    dropout=dropout,
+                )
+            )
+
+        if readout_type == "weighted_atom":
+            self.readout = WeightedAtomReadOut(dim_node_embedding, dims=[units, units, ntargets], activation=activation_type)
+        elif readout_type == "set2set":
+            self.readout = Set2SetReadOut(dim_node_embedding)
+            self.final_mlp = MLP([2 * dim_node_embedding, units, ntargets], activation=activation_type, activate_last=False)
+        else:
+            self.readout = ReduceReadOut(op="mean")
+            self.final_mlp = MLP([dim_node_embedding, units, ntargets], activation=activation_type, activate_last=False)
+
+    def _unpack_inputs(self, inputs):
+        if isinstance(inputs, dict):
+            pos = inputs.get("pos")
+            edge_index = inputs.get("edge_index")
+            node_type = inputs.get("node_type", inputs.get("z"))
+            line_edge_index = inputs.get("line_edge_index", None)
+            state_attr = inputs.get("state_attr", None)
+            pbc_offshift = inputs.get("pbc_offshift", None)
+            batch = inputs.get("batch", None)
+            num_graphs = inputs.get("num_graphs", None)
+            return pos, edge_index, node_type, line_edge_index, state_attr, pbc_offshift, batch, num_graphs
+        elif isinstance(inputs, (tuple, list)):
+            pos = inputs[0]
+            edge_index = inputs[1]
+            node_type = inputs[2]
+            line_edge_index = inputs[3] if len(inputs) > 3 else None
+            state_attr = inputs[4] if len(inputs) > 4 else None
+            pbc_offshift = inputs[5] if len(inputs) > 5 else None
+            batch = inputs[6] if len(inputs) > 6 else None
+            num_graphs = inputs[7] if len(inputs) > 7 else None
+            return pos, edge_index, node_type, line_edge_index, state_attr, pbc_offshift, batch, num_graphs
+        return inputs, None, None, None, None, None, None, None
+
+    def call(
+        self,
+        inputs,
+        edge_index=None,
+        node_type=None,
+        line_edge_index=None,
+        state_attr=None,
+        pbc_offshift=None,
+        batch=None,
+        num_graphs=None,
+        training=None,
+    ):
+        if edge_index is None:
+            (
+                pos,
+                edge_index,
+                node_type_in,
+                line_edge_index_in,
+                state_attr_in,
+                pbc_offshift_in,
+                batch_in,
+                num_graphs_in,
+            ) = self._unpack_inputs(inputs)
+            if node_type is None:
+                node_type = node_type_in
+            if line_edge_index is None:
+                line_edge_index = line_edge_index_in
+            if state_attr is None:
+                state_attr = state_attr_in
+            if pbc_offshift is None:
+                pbc_offshift = pbc_offshift_in
+            if batch is None:
+                batch = batch_in
+            if num_graphs is None:
+                num_graphs = num_graphs_in
+        else:
+            pos = inputs
+
+        num_nodes = ops.shape(pos)[0]
+        if batch is None:
+            batch = ops.zeros((num_nodes,), dtype="int32")
+        else:
+            batch = ops.cast(batch, "int32")
+        n_graphs = infer_num_graphs(batch=batch, num_graphs=num_graphs, state_attr=state_attr)
+
+        src = ops.cast(edge_index[0], "int32")
+        edge_batch = ops.take(batch, src, axis=0)
+
+        # 1. 2-body pair distances and expansion
+        _, bond_dists = compute_pair_vector_and_distance(pos, edge_index, pbc_offshift)
+        edge_attr = self.bond_expansion(bond_dists)
+
+        # 2. Embeddings
+        node_feat, edge_feat, state_feat = self.embedding(node_type, edge_attr, state_attr)
+
+        # 3. 3-body expansion if line_edge_index is available
+        three_cutoff = cosine_cutoff(bond_dists, self.threebody_cutoff)
+        if line_edge_index is not None and ops.shape(line_edge_index)[1] > 0:
+            theta, phi = compute_theta_and_phi(pos, edge_index, line_edge_index, pbc_offshift)
+            src_bonds = ops.cast(line_edge_index[0], "int32")
+            r_triplets = ops.take(bond_dists, src_bonds, axis=0)
+            three_basis = self.sbf_shf(r_triplets, theta, phi)
+        else:
+            three_basis = None
+
+        # 4. Message passing blocks
+        for i in range(len(self.blocks)):
+            if three_basis is not None and line_edge_index is not None:
+                edge_feat = self.three_body_interactions[i](
+                    edge_index, line_edge_index, three_basis, three_cutoff, node_feat, edge_feat
+                )
+            edge_feat, node_feat, state_feat = self.blocks[i](
+                edge_index, edge_feat, node_feat, state_feat, edge_attr,
+                batch=batch, edge_batch=edge_batch, num_nodes=num_nodes, num_graphs=n_graphs,
+                training=training,
+            )
+
+        # 5. Readout
+        if isinstance(self.readout, WeightedAtomReadOut):
+            output = self.readout(node_feat, batch=batch, num_graphs=n_graphs)
+        else:
+            pooled = self.readout(node_feat, batch=batch, num_graphs=n_graphs)
+            output = self.final_mlp(pooled)
+
+        return ops.squeeze(output, axis=-1)
+