Conditional computation on the device

Let the input choose the calculation

For a vector whose sum is positive, square its entries. Otherwise, negate them. Input [2, -1] should produce [4, 1]; input [-2, 1] should produce [2, -1]. We want one compiled function to accept either input and make the decision from its current tensor values.

A Java if cannot obtain that runtime decision by reading a traced tensor with .item(): Chapter 3 showed why that host read is rejected during capture. ControlFlow.cond expresses the decision as a region in the tensor program. It accepts a scalar tensor predicate, two tensor functions and an operand.

Choose Run all, then change the inputs and branch expressions. This notebook uses the bundled jax4j 0.2.1 CPU backend in the emulated browser JVM. The CPU backend interprets the structured regions; native GPU backends need their own lowering. All examples on this page execute on CPU. Reloading discards edits and session state.

1. Build a scalar predicate in tensor space

Predicates are F32 tensors containing one value. Zero means false; any nonzero value means true. To ask specifically whether the sum is positive, we use sum().sign().maximum(0): a negative sign becomes zero rather than accidentally counting as true.

Both branch functions receive the same logical operand shape. Their outputs must have matching shapes so the surrounding calculation has a consistent signature regardless of the runtime decision.

Expect [4, 1]: the sum is 1, so the square branch is selected. We use Jit.apply to trace this whole expression and manage its predicate temporaries. The ordinary input and returned output remain explicitly owned resources.

Try this: change the input to [-2, 1]. Its sum is -1, so expect [2, -1]. Now try [2, -2]: the sum is zero and the negation branch gives [-2, 2]. Both the predicate and the boundary decision are part of the function you wrote; choose them to match your actual problem.

2. Replay a different decision without another shape

Compilation can retain the conditional region while new runtime values choose its branch. Both inputs below have shape [2], F32 and the same placement. Their opposite sums should produce different decisions with one cached plan.

Expect [4, 1], [2, -1] and one plan. Change the negative input's values while keeping its shape and predict its branch from the new sum. The scalar predicate is evaluated from each call's input; it is not a Java boolean frozen by the first capture.

The captured program contains both branch descriptions and executes the selected region at runtime. That does not mean the Java lambdas are called anew for every compiled replay. Their role is to build the branch programs. We will make that capture behavior visible in the final cell.

3. A negative predicate is still true

Because the predicate convention is nonzero, a literal -1 selects the true branch. This example keeps the predicate separate from the operand so the meaning is unambiguous. The branches add or subtract 100 from [2, 3].

Expect [102, 103]. Change the predicate to 0 and expect [-98, -97]. Then try 1. The particular nonzero magnitude does not scale the output or blend the branches; it selects the true region. This is a conditional, not a weighted mixture of two results.

Keep the predicate finite and deliberately constructed. If your intent is “positive,” use an expression with that meaning rather than assuming a negative number is false. Also keep the predicate scalar: a vector of per-example choices is a batching or elementwise-selection question, not this single conditional's one-value input.

4. Differentiate the selected smooth branch

For the positive-sum case, the gradient of the summed output is 2x, giving [4, -2]. For the negative-sum case, negation gives [-1, -1]. Differentiation follows the selected branch's tensor calculation; it does not supply a smooth derivative of the discrete switching decision.

Expect [4, -2] and [-1, -1]. Try another input with a sum comfortably above or below zero and calculate its branch derivative by hand. Inputs on the switching surface need separate mathematical care: the complete function can be nonsmooth or even discontinuous there. A returned branch derivative alone does not resolve that issue.

5. Every branch must fit the result signature

An inactive branch still needs a valid program and a compatible output shape. The next cell deliberately returns a two-element vector from one branch and a one-element sum from the other. A true predicate does not exempt the false branch from the capture-time shape check.

Replace the false branch with t -> t.neg() and the shapes match again. Then change the predicate to zero and predict the result. Device regions in this release support F32 operations; keep branch outputs on the intended placement and use supported region operations. A method existing on Tensor does not by itself establish support inside a captured region.

6. Observe both Java branch builders

We deliberately put a counter in each builder. With predicate zero, the result comes from negation, but both counters should report one. This is why branch builders should be pure tensor descriptions rather than a place to send messages, update application state or perform per-decision logging.

Expect [-2, 1] and builder counts 1, 1. Those counts describe Java program construction, not execution counts of the numerical regions. Changing the predicate changes the selected result, while both regions are still captured for a direct cond call.

Continue with scan, while loops and structured state when the computation needs iteration. Revisit capture rules when choosing between fixed Java configuration and runtime tensor decisions.