e3nn_mlx.scatter_sum

e3nn_mlx.scatter_sum(source, index, dim_size=None, *, use_custom_kernel=False, jvp_safe=False)[source]

Sum rows of source into rows selected by one-dimensional index.

Set jvp_safe=True for forward-mode differentiation with respect to source. This path requires eager, fixed indices and implements the same reduction with a sorted prefix sum because MLX 0.31’s indexed-add primitive has no JVP rule.

Parameters:

  • dim_size (int | None)

  • use_custom_kernel (bool)

  • jvp_safe (bool)