Find a square root by averaging two wrong guesses

Start with 1. Divide 2 by it, average those two numbers, and repeat. The guesses go 1 → 1.5 → 1.4166667 → 1.4142157. Four updates are enough to get close to the best square root of two that a Java float can hold.

Why should averaging do that? Imagine a rectangle with area two. If one side is 1, the other must be 2. Neither side is the square root we want, but the answer lies between them. Their average gives us the next guess.

This is the Babylonian method for square roots. Plain Java arithmetic and a loop are enough to implement it; a Java runtime suffices. This notebook also uses jax4j, our Java tensor library, to show how the same calculation becomes an array operation.

Keep the area, change the sides

For a rectangle of area a, choosing a positive width x fixes its height at a/x. A square has equal sides, so we would like x and a/x to agree. Take their average as a new width:

next = (x + a/x) / 2

Then recompute the height from the original area. We are not averaging both sides independently: that would change the area. Every rectangle in the picture still has area two.

For a = 2 and x = 1, the first update is (1 + 2)/2 = 1.5. The next is (3/2 + 4/3)/2 = 17/12. Let's check those steps before bringing arrays into it.

Those values are 1.500000000000 and 1.416666666667, with the latter rounded for printing. This first cell uses Java double. The tensor calculation below uses F32, the same precision as Java float, so its last digits will differ.

Why use jax4j here?

For these few numbers, an ordinary Java loop would be shorter and simpler. The first cell's calculation already uses only standard Java; its jax4j dependency and imports prepare the later cells. We use this small problem because the arithmetic is easy to check while learning the array API.

A tensor is an array with a shape. Our expression divides, adds, and averages corresponding entries, handling one number or a batch without an explicit loop over the elements. The Roots record groups the two tensor inputs. compileTree captures that update and reuses its execution plan as the guesses change, provided the inputs remain compatible.

Only the update is captured. The outer iteration, stopping checks, and printed diagnostics remain ordinary Java. Closing old tensors introduces resource ownership, at the cost of extra code compared with Java arrays.

Those are the features this example demonstrates. It establishes no speed advantage; it uses neither GPU execution nor automatic differentiation. For a calculation this small, the tensor setup is extra work.

Our stopping rule uses the relative residual: |x*x − a| / a. It measures how closely squaring the guess recovers the requested number, relative to that number's size. We stop when every entry is within one part in a million, or when the iteration cap is reached. The residual is calculated in Java double from the downloaded F32 guesses, so its own squaring does not hide the remaining error through F32 rounding.

The loop closes the previous guess after each update and the final guess when it finishes. The compiled expression manages its intermediate tensors; the Roots record itself does not close anything. See tensor lifetimes for the ownership rules and the first JIT lesson for more on capture.

Watch two converge

Start with the same width of 1. The trace includes step zero, before any updates. Its final column tells us whether the guessed square is close enough to two.

With these defaults, the solver reaches its tolerance after four updates. Try a starting guess of 10, then 0.1f. A guess below the answer jumps above it on the first update. A very small guess can jump quite far: dividing by a small positive number makes a large number.

Why the last few steps improve so quickly

Call the exact root r, so a = r*r. A little algebra gives:

next − r = (x + r*r/x)/2 − r
         = (x − r)² / (2*x)

The denominator is positive and the numerator is a square. In exact arithmetic, the next guess therefore lands at or above the root, even if the current guess is below it. Once above the root, the average moves downward toward it.

The square in the numerator also explains the speed near the answer. When the denominator changes little, making the current error ten times smaller makes the next error roughly a hundred times smaller. This is called quadratic convergence. It describes the behavior near the root; a terrible initial guess may take several steps to get there.

This graph shows the default calculation, continued beyond the stopping point to reveal the rounding floor. Each vertical interval represents a factor of ten. Edited inputs print their own traces; this reference graphic does not change with them.

Four square roots in one array

Here are four deliberately chosen inputs: a fraction, an irrational-root case, and two perfect squares. Starting all guesses at 1 makes their different journeys visible. The batch stops when its slowest entry meets the tolerance.

Change one starting guess and rerun this cell. The other entries keep their own arithmetic paths, though the shared stopping rule may give them extra updates. Change the cap to 1 to see a run that stops before it succeeds. The returned array is then only the current approximation; read the message before treating it as a solved result.

Now compare each answer with Math.sqrt. That library call is our independent reference, not part of the iteration. We compare it with the actual F32 input value, including any rounding when that input was created.

When another update buys nothing

For √2, the float returned by the default run is approximately 1.4142135382. The double reference is approximately 1.4142135624. The difference is about 2.42e-8. That does not mean the loop needs another hundred steps: the next F32 update returns the same float.

The final cell deliberately ignores our stopping tolerance and performs eight updates. This is also the calculation behind the convergence graphic.

The mathematical error identity assumes exact arithmetic. Rounded division and addition eventually take over, so the formula does not promise endless improvement in a computer. A residual tolerance and a finite cap give the loop an honest way to finish.

Keep the inputs moderate for this experiment. Zero needs separate handling; a zero guess divides by zero. Negative numbers have no real square root, and negative guesses follow the negative-root branch. We exclude both here. Even finite positive floats can overflow during division or addition at extreme scales. A square-root routine intended for every floating-point input needs scaling and more edge-case handling.

If you know derivatives, apply Newton's update x − f(x)/f′(x) to f(x) = x² − a. Since f′(x) = 2x, it simplifies to exactly our average. The UBC calculus text on Newton's method develops that connection. The rectangles give us a way to understand this particular update before learning the general method.

For ordinary Java code that needs a square root, use Math.sqrt. Here, writing out the update lets you see where the answer comes from, which guesses cost extra steps, and where the available digits run out.