The typed feature pipeline

Every E3GNN — the default MANDALA network, the DeepH, DeepH-E3 and MACE-H presets and any custom model — is a FeaturePipeline: an ordered list of stages that read named feature slots and write new ones. Architecture names never select code paths; they select YAML files.

Anatomy of a stage

pipeline:
  inputs:                     # optional extra input contracts
    message_batch: {domain: edge, representation: index, frame: none, support: message}
  stages:
    - id: radial              # unique module name
      transform: gaussian_basis
      inputs: [edge_lengths]  # positional, never inferred by name
      outputs: [radial]       # fresh slots; existing slots cannot be overwritten
      options: {start: 0.0, stop: 6.0, n_gaussians: 32}
      role: feature           # feature | encoder | message | readout
      initialization: {seed: 17, kind: component}
  outputs: {hamiltonian: hamiltonian_native}

Slots are write-once. A stage cannot overwrite an input or a previously written slot, so residual links, parallel streams and long skip connections are explicit and the forward pass is a pure dataflow graph.
Inputs are positional. The meaning of each input is defined by its position in the transformation’s signature.
Outputs map operator names (hamiltonian, overlap, density) to slots holding matrix-valued features on the output-edge support.
Roles label stages. role: readout selects the trainable stages when freeze_backbone_train_heads_only freezes everything else.
Initialization can override a stage’s parameters with normal, uniform or zeros (with std, bound, gain) and isolate its random stream with seed; component keeps the layer’s own scheme.

Typed feature contracts

Every slot carries a FeatureSpec:

Field	Values	Meaning
`domain`	`node`, `edge`, `bond_neighbor`, `graph`, `relation`	The kind of entity each row belongs to.
`representation`	`scalar`, `irreps`, `channel_irreps`, `index`, `geometry`, `matrix`, `local_components`, `block_vectors`, `object`	How the numbers are to be interpreted.
`irreps`	e3nn string, e.g. `16x0e+8x1o`	The exact irrep layout of irrep-valued features.
`frame`	`global`, `bond`, `none`	Coordinate frame. Bond-frame components are not global irreps.
`support`	`output`, `message`, custom names	The index set of edge-like entities. Two edge features with equal shapes but different supports are not interchangeable.
`angular_basis`	permutation of `xyz`	Cartesian axis order of non-scalar irreps; global features use `xyz`.
`width`	integer	Width of non-irrep features.
`batch`	layout mapping	How a custom relation is offset when graphs are batched.

When the pipeline is compiled, every stage’s factory receives the specs of its inputs, validates them and returns the specs of its outputs. Mismatches — an edge feature where a node feature is expected, a vector where an invariant is required, two supports mixed in one sum, a readout whose irreps differ from the operator layout — raise an error at construction time.

Built-in input slots

The dataset and E3GNN.forward() provide these slots for every structure:

Slot	Contract	Content
`node_type_idx`	node, index	Species index in the resolved `species.order`.
`node_one_hot`	node, `Nx0e`	One-hot species.
`edge_type_idx`	edge, index	Index of the ordered species pair $(a_i,a_j)$, one of $N^2$.
`edge_one_hot`	edge, `N²x0e`	One-hot species pair.
`edge_index`, `edge_shift`	relation, index	Directed output edges $(i,j)$ and lattice shifts $\bm L$.
`edge_vectors`, `edge_lengths`	edge, geometry / scalar	$\bm r_{ij}^{\bm L}$ and $\\|\bm r_{ij}^{\bm L}\\|$.
`positions`, `box`	node / graph, geometry	Cartesian positions and lattice (rows), for differentiable geometry.
`node_batch`, `edge_batch`	index	Graph membership in a batch.
`selfloop_edge`	edge, index	Mask of onsite edges.
`message_edge_index`, `message_shift`, `message_vectors`, `message_lengths`, `message_to_output`	support `message`	A separately selected message graph (by default the output graph without onsite edges) and its map onto output edges (`-1` if absent).
`bond_frames`, `sub_*`	bond frame, bond neighbourhood	Local bond frames and bond-neighbourhood relations for local-coordinate architectures.

The flat MANDALA graph additionally uses edge_length_emb (Gaussian radial basis), edge_sh (spherical harmonics), pred_pair_edges_static and edge_partitions, which the model computes from the configuration.

Resolution: shapes from data

Recipes are resolved against the orbital vocabulary of the training data before a model is built (net.configuration.resolve_configuration()). Resolution fills in:

Variable	Value
`${species.order}`, `${species.count}`, `${species.pair_count}`	Species list (alphabetical or by atomic number), $N$, $N^2$.
`${species.shells}`	Shell degrees of each species, e.g. `Si: [0, 0, 1, 1, 2]`.
`${layout.coupled_irreps}`	Irreps of the packed edge vector covering every species pair.
`${layout.shell_pairs}`, `${layout.packed_dim}`, `${layout.spinful}`	Codec description consumed by `shell_tensor_decode`.
`${parameters.*}`	User-defined variables of the recipe.

layout_rules compute further shapes, e.g. tensor_product_completion finds the smallest intermediate irreps from which a tensor product can reach the coupled layout. Resolution also freezes pipeline.operator_layout — the exact species order, shell assignment and spin convention — into the checkpoint; inference verifies it before reconstructing blocks.

Inspecting a pipeline

outputs, slots = model.pipeline(x, return_trace=True)
for sid, reads, writes in model.pipeline.bindings:
    print(sid, reads, "->", writes, model.pipeline.feature_specs[writes[0]])

slots holds every intermediate tensor, which makes the trace a convenient tool for equivariance debugging: rotate the input and compare slot by slot to find the first stage that breaks the transformation law.

The registry

Transformations are registered by name:

from net.pipeline import FeatureSpec, register_transform

@register_transform("my_layer")
def build(options, inputs, mapper):
    ...                       # validate the FeatureSpecs in `inputs`
    return module, [FeatureSpec(...)]

The factory receives the stage options (a dict), the list of input specs and the shared BlockIrrepMapper, and returns a torch.nn.Module together with one spec per output. At run time the module is called with the input tensors in order; if it sets accepts_physical_derivatives = True it also receives physical_derivatives=.... Writing new layers and objectives develops this into a complete workflow.

The compiler currently accepts 47 transformation names:

add, angular_degree_expansion, bond_frames, bond_harmonics, bond_local_linear, channel_linear, concat, cutoff_radial_expansion, embedding, endpoint_gather, endpoint_irrep_edge, endpoint_irrep_node, endpoint_message_aggregation, endpoint_scalar_edge, endpoint_tensor_product_edge, endpoint_tensor_product_node, gated_scalar_message, gaussian_basis, gaussian_envelope, graph_irrep_norm, indexed_aggregation, irrep_activation, irrep_cartesian_basis, irrep_linear, irrep_residual, irrep_to_bond, local_coordinate_message, local_orbital_output, neighbor_density, one_hot, pair_distance_affine, pair_radial_angular_embedding, partitioned_blocks_to_shells, partitioned_irrep_readout, periodic_displacement, projected_attention, radial_tensor_product, scalar_gated_node, scalar_linear, scale_features, shell_tensor_decode, species_irrep_affine, species_irrep_embedding, species_tensor_product, spherical_harmonics, symmetric_product, vector_lengths