API reference

The API follows the upstream namespace organization where an MLX equivalent is implemented. Narrative guides explain intended composition; this reference is generated from the current Python signatures and docstrings.

• O(3) operations
  • Representations
  • Tensor products
  • Spherical harmonics and grids
  • Rotations
• Neural-network modules
  • e3nn_mlx.nn.Activation
  • e3nn_mlx.nn.BatchNorm
  • e3nn_mlx.nn.Dropout
  • e3nn_mlx.nn.Extract
  • e3nn_mlx.nn.ExtractIr
  • e3nn_mlx.nn.FullyConnectedNet
  • e3nn_mlx.nn.Gate
  • e3nn_mlx.nn.Identity
  • e3nn_mlx.nn.NormActivation
  • e3nn_mlx.nn.S2Activation
  • e3nn_mlx.nn.SO3Activation
• Mathematical helpers
  • e3nn_mlx.math.soft_unit_step
  • e3nn_mlx.math.soft_one_hot_linspace
  • e3nn_mlx.math.smooth_cutoff
• Graph operations
  • e3nn_mlx.radius_graph
  • e3nn_mlx.scatter_sum
• Models
  • Gate-points 2102
  • Modular v2106 models
• Representation-aware arrays
  • e3nn_mlx.IrrepsArray