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.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']

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