MolecularDiffusion.modules.layers.e3x.indexed

Sparse indexed (segment) ops. Port of the e3x/ops/indexed.py surface.

Only the sparse dst_idx/src_idx path is ported; DiTMC never uses the dense adjacency form.

Functions

gather_dst(→ torch.Tensor)	inputs[dst_idx] along the leading axis.
gather_src(→ torch.Tensor)	inputs[src_idx] along the leading axis.
indexed_max(→ torch.Tensor)	Segment max over the leading axis; empty segments give -inf.
indexed_softmax(→ torch.Tensor)	Softmax over each dst_idx segment.
indexed_sum(→ torch.Tensor)	Segment sum over the leading axis.
segment_mean(→ torch.Tensor)	Segment mean; empty segments give 0 (jraph.segment_mean semantics).

Module Contents

MolecularDiffusion.modules.layers.e3x.indexed.gather_dst(inputs: torch.Tensor, dst_idx: torch.Tensor) → torch.Tensor

inputs[dst_idx] along the leading axis.

MolecularDiffusion.modules.layers.e3x.indexed.gather_src(inputs: torch.Tensor, src_idx: torch.Tensor) → torch.Tensor

inputs[src_idx] along the leading axis.

MolecularDiffusion.modules.layers.e3x.indexed.indexed_max(inputs: torch.Tensor, dst_idx: torch.Tensor, num_segments: int) → torch.Tensor

Segment max over the leading axis; empty segments give -inf.

MolecularDiffusion.modules.layers.e3x.indexed.indexed_softmax(inputs: torch.Tensor, dst_idx: torch.Tensor, num_segments: int, multiplicative_mask: torch.Tensor | None = None) → torch.Tensor

Softmax over each dst_idx segment.

The per-segment maximum is subtracted for numerical stability under ``stop_gradient`` (detach()), and an optional multiplicative mask is applied to the raw exponentials before normalization – both exactly as e3x does.

MolecularDiffusion.modules.layers.e3x.indexed.indexed_sum(inputs: torch.Tensor, dst_idx: torch.Tensor, num_segments: int) → torch.Tensor

Segment sum over the leading axis.

MolecularDiffusion.modules.layers.e3x.indexed.segment_mean(inputs: torch.Tensor, segment_ids: torch.Tensor, num_segments: int) → torch.Tensor

Segment mean; empty segments give 0 (jraph.segment_mean semantics).