e3j.utils.options.Aggregation

class e3j.utils.options.Aggregation

Aggregation method for the sparse reduction step.

Values:

  • SCATTER – jax.lax.scatter_add — JAX-native, CPU and GPU compatible.

  • SCATTER_1 – Custom CUDA scatter_add_1 kernel (requires e3j_ops).

  • SPARSE – Matmul with a sparse BCOO target matrix.

  • DENSE – Matmul with a dense target matrix.

Methods

__getitem__(name)

Return the member matching name.

__iter__()

Return members in definition order.

__len__()

Return the number of members (no aliases)

Attributes

SCATTER

SCATTER_1

SPARSE

DENSE