e3j.utils.options.Convolution

class e3j.utils.options.Convolution

Evaluation strategy for Convolution.

Values:

  • UNFUSED – Naive JAX implementation.

  • FUSED_CUDA – Dispatch to custom CUDA kernel via XLA-FFI (requires e3j_ops).

  • FUSED_MOSAIC_TPU – Single fused Pallas Mosaic-TPU kernel (gather + TP + mixing + scatter), TRAILING_CHANNELS layout. Requires a TPU backend.

Methods

__getitem__(name)

Return the member matching name.

__iter__()

Return members in definition order.

__len__()

Return the number of members (no aliases)

Attributes

UNFUSED

FUSED_CUDA

FUSED_MOSAIC_TPU