e3tools package

Subpackages

• e3tools.nn package
  • Module contents
    • Attention
      • Attention.forward()
    • AxisToMul
      • AxisToMul.forward()
    • Conv
      • Conv.apply_per_edge()
      • Conv.forward()
    • ConvBlock
      • ConvBlock.forward()
    • EquivariantMLP
    • ExperimentalConv
    • ExperimentalTensorProduct
      • ExperimentalTensorProduct.forward()
    • ExtractIrreps
      • ExtractIrreps.forward()
    • Gate
      • Gate.forward()
    • GateWrapper
      • GateWrapper.forward()
    • Gated
      • Gated.forward()
    • LayerNorm
      • LayerNorm.forward()
    • LinearSelfInteraction
      • LinearSelfInteraction.forward()
    • MulToAxis
      • MulToAxis.forward()
    • MultiheadAttention
      • MultiheadAttention.forward()
    • ScalarMLP
    • ScaleIrreps
      • ScaleIrreps.forward()
    • SeparableConv
    • SeparableTensorProduct
      • SeparableTensorProduct.forward()
    • TransformerBlock
      • TransformerBlock.forward()

Module contents

e3tools.pack_irreps(unpacked_tuples: Iterator[Tuple[int, Irrep, Tensor]]) → Tensor

Pack fields into a single tensor.

e3tools.radius(x: Float[Tensor, 'N D'], y: Float[Tensor, 'M D'], r: float, batch_x: Int64[Tensor, 'N'] | None = None, batch_y: Int64[Tensor, 'M'] | None = None, ignore_same_index: bool = False, chunk_size: int | None = None) → Int64[Tensor, '2 E']

For each element in y find all points in x within distance r

e3tools.radius_graph(x: Float[Tensor, 'N D'], r: float, batch: Int64[Tensor, 'N'] | None = None, chunk_size: int | None = None, loop: bool = False) → Int64[Tensor, '2 E']

Returns the edge index given by applying a radial cutoff of r to x.

e3tools.scatter(src: Tensor, index: Tensor, dim: int, dim_size: int | None = None, reduce: str = 'sum') → Tensor

Collects elements at the indices index of a source tensor src.

e3tools.unpack_irreps(x: Tensor, irreps: Irreps) → Iterator[Tuple[int, Irrep, Tensor]]

Given a packed irreps vector of dimension […, irreps.dim] yield tuples (mul, ir, field) where field has dimension […, mul, 2*l+1] for each irrep in irreps