nvalchemiops.jax.segment_ops: Segment Operations#
The segment-ops module provides differentiable JAX bindings for GPU-accelerated segmented reductions and per-segment algebra.
Tip
For usage guidance, see Segment Operations.
JAX autograd bindings for segment operations (PR 3).
Mirrors the public Torch API in nvalchemiops.torch.segment_ops. Each op is
registered with jax.custom_vjp and its backward is itself wrapped in a
jax.custom_vjp so the second-order adjoint can be triggered by
jax.grad(jax.grad(...)) or jax.jacfwd(jax.jacrev(...)).
Kernel orchestration uses warp.jax_experimental.jax_callable(), which
runs the existing _launch_* Python wrappers directly on the JAX device
arrays — no host roundtrip.
Index inputs (idx) are non-differentiable — JAX returns None for them.
num_segments (number of segments) is a static argument (compile-time constant).
Public Operations#
- nvalchemiops.jax.segment_ops.segmented_sum(x, idx, num_segments)[source]#
Compute the differentiable per-segment sum of
x.Supports scalar (shape
(N,)) and vec3 (shape(N, 3)) inputs in float32 or float64. The backward pass is itself differentiable, so second-order gradients viajax.grad(jax.grad(...))are supported.- Parameters:
- Returns:
Per-segment sums;
out[s] = sum(x[i] for i where idx[i] == s).- Return type:
jax.Array, shape (num_segments,) or (num_segments, 3)
See also
nvalchemiops.jax.segment_ops.segmented_mean()Per-segment mean.
- nvalchemiops.jax.segment_ops.segmented_dot(x, y, idx, num_segments)[source]#
Compute the differentiable per-segment dot product of vec3 arrays.
For each segment
s, accumulates \(\text{out}[s] = \sum_{i:\,\text{idx}[i]=s} x[i] \cdot y[i]\). Supports float32 and float64. Second-order gradients are supported.- Parameters:
x (jax.Array, shape (N, 3)) – First vec3 operand. Must be float32 or float64.
y (jax.Array, shape (N, 3)) – Second vec3 operand. Must share dtype with
x.idx (jax.Array, shape (N,), dtype int32) – Segment index for each element. Values must lie in
[0, num_segments).num_segments (int) – Number of output segments. Treated as a static (compile-time) constant by JAX.
- Returns:
Per-segment dot products.
- Return type:
jax.Array, shape (num_segments,)
See also
nvalchemiops.jax.segment_ops.segmented_sum()Per-segment sum.
nvalchemiops.jax.segment_ops.segmented_mul()Per-element scale by per-segment scalar.
- nvalchemiops.jax.segment_ops.segmented_mul(x, y, idx, num_segments)[source]#
Scale each vec3 element by its corresponding per-segment scalar.
Computes \(\text{out}[i] = x[i] \times y[\text{idx}[i]]\) where
xis a vec3 array andyholds one scalar per segment. Supports float32 and float64. Second-order gradients are supported.- Parameters:
x (jax.Array, shape (N, 3)) – Vec3 values to scale. Must be float32 or float64.
y (jax.Array, shape (num_segments,)) – Per-segment scalar multipliers. Must share dtype with
x.idx (jax.Array, shape (N,), dtype int32) – Segment index for each element. Values must lie in
[0, num_segments).num_segments (int) – Number of segments. Treated as a static (compile-time) constant by JAX.
- Returns:
Scaled vec3 values;
out[i] = x[i] * y[idx[i]].- Return type:
jax.Array, shape (N, 3)
See also
nvalchemiops.jax.segment_ops.segmented_dot()Per-segment dot product.
- nvalchemiops.jax.segment_ops.segmented_mean(x, idx, num_segments)[source]#
Compute the differentiable per-segment mean of
x.Supports scalar (shape
(N,)) and vec3 (shape(N, 3)) inputs in float32 or float64. Per-segment element counts are computed once during the forward pass and cached as residuals so the backward never recomputes them. Second-order gradients are supported.- Parameters:
- Returns:
Per-segment means;
out[s] = mean(x[i] for i where idx[i] == s).- Return type:
jax.Array, shape (num_segments,) or (num_segments, 3)
See also
nvalchemiops.jax.segment_ops.segmented_sum()Per-segment sum.
nvalchemiops.jax.segment_ops.segmented_rms_norm()Per-segment RMS norm.
- nvalchemiops.jax.segment_ops.segmented_rms_norm(x, idx, num_segments)[source]#
Compute the differentiable per-segment RMS norm of vec3 inputs.
For each segment
s, computes \(\text{out}[s] = \sqrt{\frac{1}{|s|} \sum_{i:\,\text{idx}[i]=s} \|x[i]\|^2}\). Inverse norms and per-segment counts are precomputed during the forward pass and cached as residuals. Supports float32 and float64. Second-order gradients are supported.- Parameters:
- Returns:
Per-segment RMS norms.
- Return type:
jax.Array, shape (num_segments,)
See also
nvalchemiops.jax.segment_ops.segmented_mean()Per-segment mean.
- nvalchemiops.jax.segment_ops.segmented_matvec(v, m, idx, num_segments)[source]#
Apply a per-segment 3x3 matrix transpose to each vec3 element.
Computes \(\text{out}[i] = m[\text{idx}[i]]^\top v[i]\) where
mholds one 3x3 matrix per segment. Supports float32 and float64. Second-order gradients are supported.- Parameters:
v (jax.Array, shape (N, 3)) – Vec3 inputs. Must be float32 or float64.
m (jax.Array, shape (num_segments, 3, 3)) – Per-segment 3x3 matrices. Must share dtype with
v.idx (jax.Array, shape (N,), dtype int32) – Segment index for each element. Values must lie in
[0, num_segments).num_segments (int) – Number of segments. Treated as a static (compile-time) constant by JAX.
- Returns:
Transformed vec3 values;
out[i] = m[idx[i]]^T @ v[i].- Return type:
jax.Array, shape (N, 3)
See also
nvalchemiops.jax.segment_ops.segmented_mul()Per-element scale by per-segment scalar.