Skip to content

← JAX Profiling & Debugging step 9 of 12

Medium Primitives

Replace a Python Sum Loop with Array Operations

Jump to the task

Express the reduction as array operations

For a vector x, summing its squared entries takes two conceptual steps: square elementwise, then reduce the resulting vector to a scalar.

import jax
import jax.numpy as jnp

x = jnp.array([1., 2., 3.])
squares = x ** 2                # [1., 4., 9.]
total = jnp.sum(squares)        # scalar 14

This is not the square of the sum, which would be 36 for this input. Nor is it the Euclidean norm, which would be sqrt(14).

A static Python loop can work

def loop_sum(x):
    total = 0
    for xi in x:
        total = total + xi ** 2
    return total

With a known leading dimension, JAX can trace this loop by unrolling its iterations. It is valid code, but longer inputs add more repeated work to the traced program and can increase tracing and compilation costs. Small static loops can be useful; they are not inherently bugs.

jnp.sum(x ** 2) expresses the same mathematical reduction compactly. It is not necessarily one jaxpr equation or one device kernel. Fusion and speed depend on the compiler and backend. Reduction order can also change floating-point rounding compared with a sequential sum.

Choose a construct for the dependency

Use native array operations for this exercise. vmap is useful for applying a custom function independently to rows or elements, but it does not perform the final sum for you. A stateful sequence can use lax.scan when its length is static. A dynamic counted loop can use lax.fori_loop, and a runtime stopping condition can use lax.while_loop. These alternatives have different differentiation limits.

Python range on a dynamic traced count cannot determine how many iterations to execute during tracing. That is different from iterating over this exercise's statically shaped vector. scan's length is not a general solution for a dynamic traced count.

The diagram models cache reuse for one function with fixed dtype: equal shapes can reuse a trace, while a new shape needs another one. It does not compare loop sizes, operation counts or kernel fusion.

Loading visualization…

Problem

Implement safe_loop_sum_squares(x) without a Python loop. Compute the elementwise squares and reduce them with a JAX sum.

X is a one-dimensional float32 JAX array. Return a scalar JAX array. [3, 4] returns 25; [-3, 4] also returns 25. An empty input returns zero. The function must work under jax.jit. Numerical output tests cannot prove that you avoided a Python loop, so inspect the implementation too.