{ "schemaVersion": 1, "phaseId": "transforms", "id": "first-gradient", "title": "Your first gradient", "status": "authored", "objective": "Explain local sensitivity and the chain rule, then compare analytic, numerical, and automatic derivatives.", "prerequisites": [], "exercise": "Foundation · Define f(x) = x³. Derive f′(x), predict the values at −2, 0, and 3, then verify every prediction with JAX. Explain why the zero at x = 0 is separate from a minimum.", "evidence": "Keep predictions and analytic derivations at multiple inputs, the finite-difference step/error sweep with dtype, both argnums checks, and the scalar-objective repair. Explain why a zero gradient need not mark a minimum.", "check": "At x = 0, jax.grad(lambda x: x ** 3)(0.0) returns zero. What can you conclude?", "artifacts": { "source": "phases/02-transforms/01-your-first-gradient/docs/en.md", "script": "exercises/first-gradient.py", "notebook": "exercises/first-gradient.ipynb", "markdown": "exercises/first-gradient.md", "scriptSource": "phases/02-transforms/01-your-first-gradient/code/main.py", "notebookSource": "phases/02-transforms/01-your-first-gradient/notebooks/exercise.ipynb", "quizSource": "phases/02-transforms/01-your-first-gradient/quiz.json", "evidenceSource": "phases/02-transforms/01-your-first-gradient/outputs/evidence.md", "evidence": "exercises/first-gradient-evidence.md" }, "path": "phases/02-transforms/01-your-first-gradient", "runtime": { "kind": "cpu" }, "hardware": "CPU", "content": { "idea": "A derivative tells us how an output changes near a particular input. We can understand it before using autodiff: move the input a small amount, estimate the output change from the local slope, then compare that prediction with the function.", "code": "# Your first gradient: A derivative tells us how an output changes near a particular input.\nimport jax\nimport jax.numpy as jnp\ndef f(x):\n return x ** 2\n# Differentiate the objective to obtain `derivative` via automatic differentiation.\nderivative = jax.grad(f)\nprint(float(derivative(3.0)))\nassert jnp.allclose(derivative(3.0), 6.)\nassert jnp.allclose(derivative(0.0), 0.)", "output": "\\(6.0\\)", "exercise": "Foundation · Define \\(f(x)=x^3\\). Derive \\(f\\prime(x)\\), predict the values at \\(-2\\), \\(0\\), and \\(3\\), then verify every prediction with JAX. Explain why the zero at \\(x=0\\) is separate from a minimum.", "solution": "# Exercise solution: Foundation · Define f(x)=x^3.\ndef cubic(x):\n return x ** 3\nfor point in (-2., 0., 3.):\n assert jnp.allclose(jax.grad(cubic)(point), 3. * point ** 2)\nassert cubic(-0.1) < cubic(0.) < cubic(0.1)", "question": "For \\(f(x)=x^3\\), JAX reports \\(f\\prime(0)=0\\). What can you conclude?", "options": [ "\\(x=0\\) is a global minimum.", "The first-order local slope is zero; optimality requires more evidence.", "JAX has stopped differentiating because the output is zero." ], "answer": 1, "explanation": "The cubic is stationary at zero but takes smaller values to the left. A zero first derivative alone does not prove a minimum.", "problem": "Imagine that you can turn one dial, and a score tells you how far you are from the result you want. Which way should you turn it? We’ll begin with the curve \\(f(x)=x^2\\), where we can check the answer by hand. A derivative tells us how the score changes nearby; then JAX will calculate that derivative for us. You do not need to guess what the tool is doing.", "buildCode": "# Numerical estimate: Approximately 6.0.\ndef central_difference(function, x, h=1e-3):\n return (function(x + h) - function(x - h)) / (2 * h)\n\ndef square(x):\n return x * x\n\nestimate = central_difference(square, 3.0)\nprint(\"Finite-difference estimate:\", estimate)\nassert abs(estimate - 6.0) < 1e-8", "buildOutput": "Approximately \\(6.0\\). This estimate evaluates the function at two nearby points; it is not automatic differentiation.", "figure": "gradient", "diagnosis": "If grad reports a scalar-output error, inspect the output shape before adding an arbitrary reduction. If it reports an integer-input error, check the dtype of the differentiated argument. If a gradient is unexpectedly zero, evaluate the function at nearby inputs and verify argument selection. If finite differences disagree, record dtype and \\(h\\), sweep \\(h\\), and compare with an analytic test case. Each symptom calls for different evidence.", "objectives": [ "Predict the sign and magnitude of a derivative from a local change.", "Trace the chain rule through a composed numerical function.", "Compare an analytic derivative, a finite-difference estimate, and JAX autodiff.", "Distinguish a stationary point from a minimum and repair scalar-output/input-type errors." ], "sections": [ { "id": "guided-reasoning", "title": "Use a slope to predict a nearby change", "body": "For a separate hand calculation, let \\(f(x)=x^2\\). At \\(x=3\\), the derivative is \\(6\\). Moving to \\(3.01\\) gives a first-order prediction of \\(9+6(0.01)=9.06\\). The exact value is \\(9.0601\\); the small difference comes from curvature.\n\nIn the tangent figure, the tangent agrees with the curve at the selected point and has the same local slope. It need not follow the curve far away. Move the point and watch the slope change; then move only a small distance from that point to judge the linear approximation.\n\nAutodiff computes the derivative of the operations you wrote. It does not choose a meaningful objective for you. A gradient of the wrong loss can be numerically correct, which is why later lessons check both the formula and its derivative.", "check": { "prompt": "Does a positive derivative mean the function value itself is positive?", "answer": "No. The derivative describes local change. For \\(f(x)=x-10\\) at \\(x=0\\), the value is \\(-10\\) while the derivative is \\(1\\)." } }, { "title": "Read a slope before you compute it", "body": "Let’s start at \\(x=3\\). If we move right by \\(\\delta=0.01\\), the score changes from \\(9\\) to \\(9.0601\\). That is an increase of \\(0.0601\\). The slope predicts an increase of about \\(6\\times0.01=0.06\\), which is close.\n\nHere \\(\\delta\\) means a small change in the input, and \\(f\\prime(x)\\) means the slope at that input. Expanding the square gives the exact change below. The term \\(\\delta^2\\) is what the straight-line approximation leaves out. As the change gets smaller, that term becomes small compared with the linear term when the slope is nonzero.\n\nBefore running the code, try a negative input. At \\(x=-3\\), moving a little to the right reduces the score. The slope should therefore be negative. At zero, the slope is zero even though moving away in either direction increases the square.", "math": "\\begin{aligned}f(x+\\delta)-f(x)&=2x\\delta+\\delta^2\\\\f\\prime(x)&=2x\\end{aligned}" }, { "title": "Three ways to obtain a derivative", "body": "An analytic derivative comes from reasoning about the mathematical function. A finite difference approximates a derivative by evaluating the function at nearby inputs. Automatic differentiation applies derivative rules to the operations that execute in the program and composes those rules with the chain rule.\n\nAutodiff is not estimating a slope with a hidden step size, and it is not a symbolic algebra system simplifying a formula. Your program still runs floating-point operations. Those operations can overflow, lose precision, or evaluate a function whose mathematical derivative is not well-defined at the chosen point. Agreement among independent approaches is stronger evidence than merely printing one number." }, { "title": "Follow the chain rule through the program", "body": "Now split the calculation into two small steps. First compute \\(u=3x+1\\), then square it to get \\(y=u^2\\). A small change in \\(x\\) first changes \\(u\\); that change in \\(u\\) then changes \\(y\\). The chain rule multiplies the two local slopes.\n\nThe first slope is \\(3\\). The second is \\(2u\\). Multiplying gives \\(6(3x+1)\\). At \\(x=2\\), we have \\(u=7\\), so the derivative is \\(42\\). This gives us a hand-calculated answer for the JAX check. Autodiff follows the operations in the program; it is not estimating the derivative by repeatedly nudging the input.", "math": "\\begin{aligned}u&=3x+1,\\qquad y=u^2\\\\\\frac{dy}{dx}&=\\frac{dy}{du}\\frac{du}{dx}=2u\\cdot3\\end{aligned}" }, { "title": "Know the contract of grad", "body": "The default grad transformation differentiates the first positional argument. Its differentiated inputs must have an appropriate inexact dtype, such as a floating-point array, and its output must be scalar: shape \\(()\\), not shape \\((1,)\\). A Python float like \\(3.0\\) works; the integer \\(3\\) is a different input type.\n\nA vector of residuals is useful model output, but it does not yet specify a single objective. Reducing the residuals to a sum or a mean chooses an objective and its scaling. A Jacobian is the more appropriate object when you really want the derivatives of every vector output. We will first learn scalar objectives, then return to batching and derivative structure." }, { "title": "A zero gradient is a question, not a verdict", "body": "Both \\(x^2\\) and \\(x^3\\) have derivative zero at \\(x=0\\). The square has a minimum there; the cubic keeps decreasing toward the left. The constant function \\(x\\mapsto5\\) also has zero derivative everywhere. A gradient is local first-order information, so it does not by itself prove optimality, global quality, or a correct implementation.\n\nWhen a model produces a surprising zero gradient, ask whether the loss actually depends on the parameter, whether you differentiated the intended argument, and whether the chosen point is stationary. Checking nearby points and a simple analytic case is often more informative than immediately changing the optimizer." } ], "experiments": [ { "title": "Check the chain rule at three inputs", "prediction": "For \\(g(x)=(3x+1)^2\\), predict the derivative at \\(-1\\), \\(0\\), and \\(2\\). Explain why the sign changes.", "code": "# Experiment — Check the chain rule at three inputs: The inner derivative contributes a factor of 3.\ndef composed(x):\n return (3. * x + 1.) ** 2\nfor point in (-1., 0., 2.):\n # Differentiate the objective to obtain `observed` via automatic differentiation.\n observed = jax.grad(composed)(point)\n expected = 6. * (3. * point + 1.)\n print(\"chain rule:\", point, float(observed))\n assert jnp.allclose(observed, expected)", "output": "The three derivatives are \\(-12\\), \\(6\\), and \\(42\\).", "explanation": "The inner derivative contributes a factor of \\(3\\). Omitting that factor is a chain-rule error; matching only one point can hide a wrong expression." }, { "title": "Choose a finite-difference scale", "prediction": "For the float32 function \\(x^3\\) at \\(x=3\\), will shrinking \\(h\\) forever keep improving the estimate? Run the sweep and compare with \\(27\\).", "code": "# Experiment — Choose a finite-difference scale: A central difference has truncation error at larger h and...\ndef cubic32(x):\n return x ** 3\npoint = jnp.array(3., dtype=jnp.float32)\nfor step in (1e-1, 1e-2, 1e-3, 1e-5, 1e-7):\n estimate = (cubic32(point + step) - cubic32(point - step)) / (2. * step)\n print(\"h / estimate / error:\", step, float(estimate), float(jnp.abs(estimate - 27.)))\nassert jnp.allclose(jax.grad(cubic32)(point), 27.)", "output": "The errors need not decrease monotonically. At sufficiently small \\(h\\), rounding can erase the input change.", "explanation": "A central difference has truncation error at larger \\(h\\) and cancellation/rounding error at very small \\(h\\). The original numerical-estimate example used Python floats; this sweep deliberately uses a JAX float32 array. Do not silently compare their precision as if it were identical." } ], "practice": [ { "title": "Differentiate the parameter you intended", "difficulty": "Practice", "prompt": "For the prediction \\(\\hat y=wx\\), find the derivative with respect to \\(w\\) and then with respect to \\(x\\). Evaluate them at \\(w=2\\) and \\(x=3\\). Predict both answers before choosing `argnums`.", "hint": "grad defaults to `argnums=0.` The other input remains an input to the transformed callable.", "solution": "# Differentiate the parameter you intended (Practice): A parameter gradient and an input sensitivity answer...\ndef scalar_prediction(weight, x):\n return weight * x\nassert jnp.allclose(jax.grad(scalar_prediction, argnums=0)(2., 3.), 3.)\nassert jnp.allclose(jax.grad(scalar_prediction, argnums=1)(2., 3.), 2.)", "explanation": "A parameter gradient and an input sensitivity answer different questions even when they originate from the same function." }, { "title": "Repair an objective without hiding its meaning", "difficulty": "Challenge", "prompt": "A function returns \\([x^2,(x-2)^2]\\). Explain why ordinary grad rejects it. Build a scalar objective using their mean, derive its derivative, and check \\(x=0\\), \\(1\\), and \\(2\\). Would a sum have the same derivative magnitude?", "hint": "Write the reduction as a mathematical function first: its mean is \\(x^2-2x+2\\).", "solution": "# Repair an objective without hiding its meaning (Challenge): Mean and sum agree on the minimizer here but differ by a...\ndef two_costs(x):\n return jnp.array([x ** 2, (x - 2.) ** 2])\ndef mean_cost(x):\n return jnp.mean(two_costs(x))\nfor point in (0., 1., 2.):\n assert jnp.allclose(jax.grad(mean_cost)(point), 2. * point - 2.)\n assert jnp.allclose(jax.grad(lambda z: jnp.sum(two_costs(z)))(point), 2. * jax.grad(mean_cost)(point))", "explanation": "Mean and sum agree on the minimizer here but differ by a factor of two in gradient magnitude. Reducing output is a modeling choice, not just a syntax repair." } ], "takeaways": [ "A derivative is a local sensitivity; a gradient transformation returns a callable.", "Autodiff applies the chain rule to program operations, while finite differences approximate with input perturbations.", "Scalar objective, input dtype, and differentiated argument are part of the contract.", "A stationary point and a minimum are different claims." ], "references": [ { "title": "JAX: automatic differentiation", "url": "https://docs.jax.dev/en/latest/automatic-differentiation.html" }, { "title": "JAX: grad API and argument contract", "url": "https://docs.jax.dev/en/latest/_autosummary/jax.grad.html" } ], "visual": { "title": "A tangent describes local change", "kind": "executed", "prediction": "At \\(x=3\\), is the slope equal to the height of the curve?", "reading": "The solid curve is \\(f(x)=x^2\\). The dashed line is its tangent at the marked point \\((3,9)\\). The horizontal axis measures the input; the vertical axis shows either the function value or the tangent’s predicted value.\n\nNear the marked point the two almost coincide. Farther away they separate: at \\(x=2\\), the curve is \\(4\\) and the tangent is \\(3\\). At \\(x=0\\), the curve is \\(0\\), but the tangent has fallen to \\(-9\\).", "connection": "The derivative at \\(x=3\\) is \\(6\\), so the local prediction is \\(f(3+\\Delta x)\\approx9+6\\Delta x\\). The exact value also contains \\((\\Delta x)^2\\). That omitted term explains why the approximation is good for small changes and increasingly poor for large ones.\n\nThe negative part of the dashed line does not mean a square became negative. It means a local linear approximation was extended beyond the neighborhood where it is useful. A gradient describes local change, not the complete function.", "alt": "A tangent describes local change. The solid curve is \\(f(x)=x^2\\). The dashed line is its tangent at the marked point \\((3,9)\\). The horizontal axis measures the input; the vertical axis shows either the function value or the tangent’s predicted value. Near the marked point the two almost coincide. Farther away they separate: at \\(x=2\\), the curve is \\(4\\) and the tangent is \\(3\\). At \\(x=0\\), the curve is \\(0\\), but the tangent has fallen to \\(-9\\).", "code": "# Compute figure data for: A tangent describes local change\n# Generate a uniform grid of points in `grid`.\ngrid = jnp.linspace(-4.0, 4.0, 81)\npoint = 3.0\nslope = derivative(point)\n# Vectorize across the batch dimension without a Python loop (`visual_data`).\nvisual_data = {'kind': 'line', 'x': grid.tolist(), 'xlabel': 'input x', 'ylabel': 'function / tangent value', 'series': [{'label': 'f(x) = x squared', 'y': jax.vmap(f)(grid).tolist()}, {'label': 'tangent at x = 3', 'y': (f(point) + slope * (grid - point)).tolist()}]}\nvisual_data['markers'] = [{'x': point, 'y': float(f(point)), 'label': 'selected input'}]" }, "buildSteps": [ { "title": "Step 1: Set up imports and input tensors", "instruction": "Import the required JAX modules and define the initial inputs for your first gradient.", "code": "import jax\nimport jax.numpy as jnp", "explanation": "Establishing explicit input shapes and dtypes first makes the downstream transformation contract deterministic." }, { "title": "Step 2: Apply the core JAX transformation", "instruction": "Write the core computation and transformation step over the initialized inputs.", "code": "def f(x):\n return x ** 2\n# Differentiate the objective to obtain `derivative` via automatic differentiation.\nderivative = jax.grad(f)", "explanation": "This stage executes the primary numerical transformation and binds the intermediate outputs." }, { "title": "Step 3: Verify shapes and numerical invariants", "instruction": "Check that the resulting arrays satisfy the expected shape, dtype, and numerical tolerances.", "code": "assert jnp.allclose(derivative(3.0), 6.)\nassert jnp.allclose(derivative(0.0), 0.)", "explanation": "These assertions lock in the exact numerical contract before you run the full experiment and variations." } ] }, "minutes": 55, "contentHash": "c382307c1e6ed518fc942a959c2d6583641585dc4e4e727b7f7c5c754b8a500c", "visualArtifact": { "image": "figures/first-gradient.svg", "png": "figures/first-gradient.png", "data": "figures/first-gradient.json", "generatedAt": "2026-10-20T10:45:10.748798+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "contentHash": "c382307c1e6ed518fc942a959c2d6583641585dc4e4e727b7f7c5c754b8a500c", "kind": "executed" }, "execution": { "contentHash": "c382307c1e6ed518fc942a959c2d6583641585dc4e4e727b7f7c5c754b8a500c", "executedAt": "2026-10-20T10:45:17.292166+00:00", "environment": { "backend": "cpu", "deviceCount": 1, "python": "3.14.3", "jax": "0.9.2", "numpy": "2.4.4", "matplotlib": "3.10.8" }, "stdout": "Finite-difference estimate: 5.999999999999339\n6.0\nchain rule: -1.0 -12.0\nchain rule: 0.0 6.0\nchain rule: 2.0 42.0\nh / estimate / error: 0.1 27.009973526000977 0.009973526000976562\nh / estimate / error: 0.01 27.000045776367188 4.57763671875e-05\nh / estimate / error: 0.001 26.9975643157959 0.0024356842041015625\nh / estimate / error: 1e-05 27.0843505859375 0.0843505859375\nh / estimate / error: 1e-07 0.0 27.0\nPASS: first-gradient\n", "url": "executions/first-gradient.json" } }