{
  "source_file": "0930",
  "source_platform": "Hashnode",
  "author": "Wesley Kambale",
  "title_listed": "Transformations That Change Everything",
  "keywords": "",
  "domain": "机器学习",
  "content_type": "深度技术解析",
  "is_paid_flag": "否",
  "url": "https://kambale.dev/transformations-that-change-everything",
  "publish_date": "2026-02-09",
  "corpus_stage": "",
  "notes": "",
  "id": "1a40d3c86d69",
  "collected_at": "2026-10-02T14:11:06",
  "title": "Automatic vectorization with vmap and gradients with grad in Jax",
  "body_markdown": "Transformations That Change Everything\nAutomatic vectorization with vmap and gradients with grad in Jax\nUpdated\nJuly 29, 2026\n•\n11\nmin read\n•\nView as Markdown\nW\nWesley Kambale\nI'm a Machine Learning Engineer passionate about building production-ready ML systems for the African market.\nWith experience in TensorFlow, Keras, and Python-based workflows, I help teams bridge the gap between machine learning research and real-world deployment—especially on resource-constrained devices.\nI'm also a Google Developer Expert in AI. I regularly speak at tech conferences including PyCon Africa, DevFest Kampala, DevFest Nairobi and more and also write technical articles on AI/ML here.\nPart of series\nJAX\nLast week, we learned that JAX makes code fast through JIT compilation. We took a matrix multiplication from 2 seconds to 0.001 seconds with a single decorator.\nBut speed isn't JAX's only trick. The real power of JAX lies in its\ntransformations\n; functions that take functions and return new functions with different behavior.\nToday we're covering the two transformations that make JAX indispensable for machine learning:\njax.vmap\n: Automatic vectorization. Write code for one example, run it on a million.\njax.grad\n: Automatic differentiation. Get gradients of any function, for free.\nBy the end of this article, you'll understand why JAX developers almost never write for-loops, and you'll have trained multiple machine learning models in parallel without writing any loop at all.\nThe Problem with Python Loops\nLet's start with why loops are the enemy.\nWhen you write a Python for-loop, the interpreter does a surprising amount of work for each iteration:\nresults = []\nfor x in data:\n    # For EACH iteration, Python must:\n    # 1. Fetch x from memory\n    # 2. Check the type of x\n    # 3. Look up what \"+\" means for that type\n    # 4. Execute the addition\n    # 5. Append to the list (which may require memory reallocation)\n    results.append(x + 1)\nIf\ndata\nhas a million elements, Python performs those administrative steps a million times. The actual math (\nx + 1\n) takes nanoseconds. The overhead takes microseconds. You're spending 99% of your time on bookkeeping.\nIn deep learning, we process batches: 64 images, 128 sentences, 256 audio clips at once. If you loop through them in Python, your GPU sits idle while Python shuffles paperwork.\nNumPy helps by pushing operations into C:\nresults = data + 1  # Vectorized, fast\nBut what if your function is more complex than addition? What if it involves multiple steps, conditionals, or nested operations? You'd have to manually rewrite everything to handle batches, adding batch dimensions everywhere and keeping track of which axis is which.\nThis is where\njax.vmap\nchanges the game.\njax.vmap: Automatic Vectorization\nvmap\nstands for \"vectorizing map.\" It takes a function written for a single example and transforms it into a function that operates on batches.\nThe Basic Pattern\nimport jax\nimport jax.numpy as jnp\n\n# A function that works on ONE number\ndef square(x):\n    return x ** 2\n\n# Transform it to work on MANY numbers\nbatched_square = jax.vmap(square)\n\n# Now use it\nnumbers = jnp.array([1, 2, 3, 4, 5])\nresult = batched_square(numbers)\nprint(result)  # [1, 4, 9, 16, 25]\n\"But wait,\" you might say, \"I could just write\nnumbers ** 2\ndirectly.\"\nTrue. The power of\nvmap\nshows up with complex functions:\ndef complex_operation(x):\n    \"\"\"A function with multiple steps.\"\"\"\n    a = jnp.sin(x)\n    b = jnp.exp(-x ** 2)\n    c = a * b + jnp.log(1 + jnp.abs(x))\n    return c\n\n# Without vmap, you'd need to think about broadcasting at each step\n# With vmap, you just wrap it\nbatched_complex = jax.vmap(complex_operation)\n\nx_batch = jnp.linspace(-3, 3, 1000)\nresults = batched_complex(x_batch)\nYou write the function thinking about one input.\nvmap\nhandles the batch dimension for you.\nMultiple Arguments: The\nin_axes\nParameter\nReal functions have multiple arguments.\nin_axes\ntells\nvmap\nwhich arguments to map over and which to broadcast.\nConsider a dot product:\ndef dot_product(weights, features):\n    return jnp.dot(weights, features)\nDifferent scenarios require different\nin_axes\n:\n# Scenario 1: One set of weights, many feature vectors\n# weights: don't map (None), features: map over axis 0\nbatch_predict = jax.vmap(dot_product, in_axes=(None, 0))\n\nweights = jnp.array([1.0, 2.0, 3.0])\nfeatures = jnp.array([\n    [1.0, 0.0, 0.0],\n    [0.0, 1.0, 0.0],\n    [0.0, 0.0, 1.0],\n])\n\nresults = batch_predict(weights, features)\nprint(results)  # [1., 2., 3.]\n# Scenario 2: Many weight sets, one feature vector (ensemble of models)\n# weights: map over axis 0, features: don't map (None)\nensemble_predict = jax.vmap(dot_product, in_axes=(0, None))\n\nmany_weights = jnp.array([\n    [1.0, 0.0, 0.0],\n    [0.0, 1.0, 0.0],\n    [0.0, 0.0, 1.0],\n])\nsingle_features = jnp.array([1.0, 2.0, 3.0])\n\nresults = ensemble_predict(many_weights, single_features)\nprint(results)  # [1., 2., 3.]\n# Scenario 3: Many weights, many features (parallel evaluation)\n# Both: map over axis 0\nparallel_predict = jax.vmap(dot_product, in_axes=(0, 0))\n\nresults = parallel_predict(many_weights, features)\nprint(results)  # [1., 2., 3.]\nThe rule is simple:\n0\nmeans \"iterate over the first axis of this argument\"\nNone\nmeans \"broadcast this argument to all iterations\"\nYou can use other integers for different axes\nNested vmap\nYou can stack\nvmap\ncalls for multi-dimensional batching:\ndef single_multiply(a, b):\n    return a * b\n\n# Map over rows, then over columns\ndouble_batched = jax.vmap(jax.vmap(single_multiply))\n\nmatrix_a = jnp.array([[1, 2], [3, 4]])\nmatrix_b = jnp.array([[10, 20], [30, 40]])\n\nresult = double_batched(matrix_a, matrix_b)\nprint(result)\n# [[10, 40],\n#  [90, 160]]\njax.grad: Automatic Differentiation\nThe other transformation that makes JAX essential for ML is\njax.grad\n. It computes gradients automatically.\nThe Basic Pattern\ndef f(x):\n    return x ** 2\n\n# grad returns a NEW FUNCTION that computes the derivative\ndf_dx = jax.grad(f)\n\nprint(f(3.0))      # 9.0\nprint(df_dx(3.0))  # 6.0 (derivative of x² is 2x, and 2*3=6)\nThis works for any function, no matter how complex:\ndef messy_function(x):\n    return jnp.sin(x) * jnp.exp(-x ** 2) + jnp.tanh(x)\n\ngradient_fn = jax.grad(messy_function)\n\n# The gradient at x=1.0\nprint(gradient_fn(1.0))  # -0.5047...\nYou didn't write any derivative rules. JAX traced through your function and computed the gradient automatically using the chain rule.\njax.value_and_grad: Get Both at Once\nIn training loops, you need both the loss value (to log it) and the gradients (to update parameters).\njax.value_and_grad\ngives you both in one pass:\ndef loss_fn(params, x, y):\n    prediction = params * x\n    return (prediction - y) ** 2\n\n# Returns (loss_value, gradients)\nloss_and_grad_fn = jax.value_and_grad(loss_fn)\n\nparams = 1.0\nx, y = 2.0, 6.0  # We want params=3.0 so that 3*2=6\n\nloss, grad = loss_and_grad_fn(params, x, y)\nprint(f\"Loss: {loss}\")  # 16.0 (because (1*2 - 6)² = 16)\nprint(f\"Grad: {grad}\")  # -16.0\nGradients with Respect to Specific Arguments\nBy default,\ngrad\ndifferentiates with respect to the first argument. Use\nargnums\nto change this:\ndef f(x, y):\n    return x ** 2 + x * y\n\n# Gradient with respect to x (first argument, default)\ndf_dx = jax.grad(f, argnums=0)\n\n# Gradient with respect to y (second argument)\ndf_dy = jax.grad(f, argnums=1)\n\n# Gradients with respect to both\ndf_both = jax.grad(f, argnums=(0, 1))\n\nx, y = 2.0, 3.0\nprint(f\"df/dx: {df_dx(x, y)}\")  # 2*2 + 3 = 7\nprint(f\"df/dy: {df_dy(x, y)}\")  # 2\nprint(f\"Both:  {df_both(x, y)}\")  # (7.0, 2.0)\nCombining Transforms: The Real Power\nJAX transforms compose. You can combine\njit\n,\nvmap\n, and\ngrad\nfreely:\ndef loss_single(params, x, y):\n    \"\"\"Loss for a single data point.\"\"\"\n    pred = params[0] * x + params[1]  # Linear: y = mx + b\n    return (pred - y) ** 2\n\n# Stack the transforms:\n# 1. grad: compute gradients with respect to params\n# 2. vmap: do this for a batch of (x, y) pairs\n# 3. jit: compile the whole thing\n\nbatched_grad_fn = jax.jit(\n    jax.vmap(\n        jax.grad(loss_single),\n        in_axes=(None, 0, 0)  # Same params, batch of x, batch of y\n    )\n)\n\nparams = jnp.array([1.0, 0.0])  # Initial guess: y = 1*x + 0\nx_batch = jnp.array([1.0, 2.0, 3.0])\ny_batch = jnp.array([2.0, 4.0, 6.0])  # True relationship: y = 2x\n\n# Get gradients for each example in the batch\ngrads_per_example = batched_grad_fn(params, x_batch, y_batch)\nprint(\"Gradients per example:\")\nprint(grads_per_example)\n\n# Average them for a batch gradient\nbatch_grad = jnp.mean(grads_per_example, axis=0)\nprint(f\"Batch gradient: {batch_grad}\")\nThis pattern,\njit(vmap(grad(...)))\n, is the backbone of efficient training in JAX.\nProject: Parallel Linear Regression\nLet's train multiple models simultaneously without a single Python loop during training.\nScenario\n: We have housing price data from three different cities. Each city has different price dynamics, so we want to train a separate linear regression model for each.\nStep 1: Generate Synthetic Data\nimport jax\nimport jax.numpy as jnp\nfrom jax import random\n\n# Seed for reproducibility\nkey = random.PRNGKey(42)\n\n# True parameters for 3 cities\n# City 0: price = 50 * size + 100\n# City 1: price = 30 * size + 200  \n# City 2: price = 80 * size + 50\ntrue_slopes = jnp.array([50.0, 30.0, 80.0])\ntrue_intercepts = jnp.array([100.0, 200.0, 50.0])\n\nn_cities = 3\nn_samples = 100\n\n# Generate features (house sizes) for each city\nkey, subkey = random.split(key)\nX = random.uniform(subkey, (n_cities, n_samples, 1), minval=10, maxval=100)\n\n# Generate targets (prices) with some noise\nkey, subkey = random.split(key)\nnoise = random.normal(subkey, (n_cities, n_samples)) * 50\n\n# Y[i] = true_slopes[i] * X[i] + true_intercepts[i] + noise[i]\nY = (X[:, :, 0] * true_slopes[:, None] + \n     true_intercepts[:, None] + \n     noise)\n\nprint(f\"X shape: {X.shape}\")  # (3, 100, 1)\nprint(f\"Y shape: {Y.shape}\")  # (3, 100)\nStep 2: Define the Model for ONE City\nWe write everything as if we only have one city:\ndef predict(params, x):\n    \"\"\"Predict price for one city's houses.\"\"\"\n    slope, intercept = params[0], params[1]\n    return x[:, 0] * slope + intercept\n\ndef loss_fn(params, x, y):\n    \"\"\"MSE loss for one city.\"\"\"\n    predictions = predict(params, x)\n    return jnp.mean((predictions - y) ** 2)\n\ndef train_step(params, x, y, learning_rate):\n    \"\"\"One gradient descent step for one city.\"\"\"\n    loss, grads = jax.value_and_grad(loss_fn)(params, x, y)\n    new_params = params - learning_rate * grads\n    return new_params, loss\nStep 3: Vectorize Across Cities\nNow we use\nvmap\nto run training for all three cities in parallel:\n# Vectorize the training step\n# params: axis 0 (each city has its own params)\n# x: axis 0 (each city has its own data)\n# y: axis 0 (each city has its own targets)\n# learning_rate: None (same for all)\nparallel_train_step = jax.jit(\n    jax.vmap(train_step, in_axes=(0, 0, 0, None))\n)\n\n# Initialize parameters for all 3 cities\n# Shape: (3, 2) - 3 cities, 2 params each (slope, intercept)\nkey, subkey = random.split(key)\nparams = random.normal(subkey, (n_cities, 2))\n\nprint(\"Initial parameters:\")\nfor i in range(n_cities):\n    print(f\"  City {i}: slope={params[i, 0]:.2f}, intercept={params[i, 1]:.2f}\")\nStep 4: Train All Models in Parallel\nlearning_rate = 0.0001\nn_epochs = 2000\n\nprint(\"\\nTraining...\")\nfor epoch in range(n_epochs):\n    # This single line trains ALL THREE models simultaneously\n    params, losses = parallel_train_step(params, X, Y, learning_rate)\n    \n    if epoch % 500 == 0:\n        print(f\"Epoch {epoch:4d} | Losses: {losses}\")\n\nprint(\"\\nFinal learned parameters:\")\nprint(f\"{'City':<6} {'True Slope':<12} {'Learned Slope':<15} {'True Intercept':<15} {'Learned Intercept':<15}\")\nprint(\"-\" * 70)\nfor i in range(n_cities):\n    print(f\"{i:<6} {true_slopes[i]:<12.1f} {params[i, 0]:<15.2f} {true_intercepts[i]:<15.1f} {params[i, 1]:<15.2f}\")\nExpected output:\nFinal learned parameters:\nCity   True Slope   Learned Slope   True Intercept   Learned Intercept  \n----------------------------------------------------------------------\n0      50.0         49.87           100.0            102.34         \n1      30.0         29.92           200.0            201.15         \n2      80.0         79.78           50.0             52.89\nWe trained three separate models, and there's not a single Python for-loop in the training logic. The\nparallel_train_step\nfunction processes all cities in one fused GPU kernel.\nWhy This Matters\nThe pattern we just used scales to serious applications:\nHyperparameter search\n: Train 100 models with different learning rates simultaneously. Pick the best one.\nEnsemble methods\n: Train 10 models with different random seeds. Average their predictions for more robust results.\nPer-user personalization\n: Train a tiny model for each of your 10,000 users.\nvmap\nhandles the parallelization.\nBayesian methods\n: Sample 1000 parameter configurations from a posterior distribution and evaluate all of them at once.\nThe key insight is that\nvmap\ndoesn't just save you from writing loops; it enables computations that would be impractical with sequential processing.\nExercises\nGradient verification\n: Use\njax.grad\nto compute the derivative of\nf(x) = sin(x)\n. Plot it alongside\ncos(x)\nto verify they match.\nThe ensemble challenge\n: Modify the project to train 10 models on the\nsame\ncity but with different random initializations. Use\nvmap\nover the params axis only (\nin_axes=(0, None, None, None)\n). Check if they all converge to similar values.\nSecond derivatives\n:\njax.grad\nreturns a function; and you can take its gradient too. Compute the second derivative of\nf(x) = x³\nand verify it equals\n6x\n.\nQuick Reference\nimport jax\nimport jax.numpy as jnp\n\n# vmap: Automatic Vectorization\ndef single_fn(x):\n    return x ** 2\n\nbatched_fn = jax.vmap(single_fn)\nresults = batched_fn(jnp.array([1, 2, 3]))\n\n# With multiple arguments\ndef dot(w, x):\n    return jnp.dot(w, x)\n\n# Shared weights, batched inputs\nbatch_dot = jax.vmap(dot, in_axes=(None, 0))\n\n# grad: Automatic Differentiation\ndef loss(params):\n    return params ** 2\n\ngrad_fn = jax.grad(loss)\ngradient = grad_fn(3.0)  # 6.0\n\n# Get both value and gradient\nloss_val, grad_val = jax.value_and_grad(loss)(3.0)\n\n# Gradient with respect to specific argument\ndef f(x, y):\n    return x * y\n\ndf_dy = jax.grad(f, argnums=1)\n\n# Combining Transforms\nfast_batched_grad = jax.jit(jax.vmap(jax.grad(loss_fn), in_axes=(None, 0)))\nWhat's Next\nWe've now covered the core JAX transforms:\njit\nfor speed,\nvmap\nfor batching, and\ngrad\nfor gradients. These three tools are enough to train neural networks from scratch.\nBut writing raw JAX for complex models gets tedious. Next week, we'll introduce\nFlax NNX';\na neural network library that gives you PyTorch-style ergonomics while keeping all the power of JAX transformations.\nWe'll build our first real neural network: a CNN for image classification.\nResources\nAutomatic Vectorization\nJAX AI Stack Guide\n#\njax",
  "metadata": {
    "title": "Automatic vectorization with vmap and gradients with grad in Jax",
    "description": "The core JAX transforms: jit for speed, vmap for batching, and grad for gradients. These three tools are enough to train neural networks from scratch.",
    "keywords": null,
    "author": null,
    "og:title": "Automatic vectorization with vmap and gradients with grad in Jax",
    "og:description": "The core JAX transforms: jit for speed, vmap for batching, and grad for gradients. These three tools are enough to train neural networks from scratch.",
    "og:url": "https://kambale.dev/transformations-that-change-everything",
    "og:site_name": "Wesley Kambale",
    "og:image": "https://cdn.hashnode.com/res/hashnode/image/upload/v1770659671216/1cee86fb-2dc9-42ed-a66c-5294c8c11bc1.png",
    "og:image:width": "1200",
    "og:image:height": "630",
    "og:type": "article",
    "twitter:card": "summary_large_image",
    "twitter:creator": "https://twitter.com/WesKambale",
    "twitter:title": "Automatic vectorization with vmap and gradients with grad in Jax",
    "twitter:description": "The core JAX transforms: jit for speed, vmap for batching, and grad for gradients. These three tools are enough to train neural networks from scratch.",
    "twitter:image": "https://cdn.hashnode.com/res/hashnode/image/upload/v1770659671216/1cee86fb-2dc9-42ed-a66c-5294c8c11bc1.png",
    "article:published_time": "2026-02-09T17:57:28.781Z",
    "article:modified_time": "2026-07-29T00:59:06.354Z",
    "article:author": "https://kambale.dev/@kambale"
  },
  "paywall": {
    "listed_paid": false,
    "detected": false,
    "marker": "",
    "public_only": false,
    "policy": "付费/会员内容只保留页面公开可见部分，不绕过付费墙"
  },
  "crawl": {
    "success": true,
    "status_code": 200,
    "error": ""
  },
  "code_blocks": [
    {
      "file": "attachments/1a40d3c86d69/code/body-000.txt",
      "chars": 295,
      "language": "text"
    },
    {
      "file": "attachments/1a40d3c86d69/code/body-001.txt",
      "chars": 39,
      "language": "text"
    },
    {
      "file": "attachments/1a40d3c86d69/code/body-002.txt",
      "chars": 300,
      "language": "text"
    },
    {
      "file": "attachments/1a40d3c86d69/code/body-003.txt",
      "chars": 382,
      "language": "text"
    },
    {
      "file": "attachments/1a40d3c86d69/code/body-004.txt",
      "chars": 74,
      "language": "text"
    },
    {
      "file": "attachments/1a40d3c86d69/code/body-005.txt",
      "chars": 369,
      "language": "text"
    },
    {
      "file": "attachments/1a40d3c86d69/code/body-006.txt",
      "chars": 416,
      "language": "text"
    },
    {
      "file": "attachments/1a40d3c86d69/code/body-007.txt",
      "chars": 228,
      "language": "text"
    },
    {
      "file": "attachments/1a40d3c86d69/code/body-008.txt",
      "chars": 302,
      "language": "text"
    },
    {
      "file": "attachments/1a40d3c86d69/code/body-009.txt",
      "chars": 195,
      "language": "text"
    },
    {
      "file": "attachments/1a40d3c86d69/code/body-010.txt",
      "chars": 181,
      "language": "text"
    },
    {
      "file": "attachments/1a40d3c86d69/code/body-011.txt",
      "chars": 370,
      "language": "text"
    },
    {
      "file": "attachments/1a40d3c86d69/code/body-012.txt",
      "chars": 423,
      "language": "text"
    },
    {
      "file": "attachments/1a40d3c86d69/code/body-013.txt",
      "chars": 920,
      "language": "text"
    },
    {
      "file": "attachments/1a40d3c86d69/code/body-014.txt",
      "chars": 899,
      "language": "text"
    },
    {
      "file": "attachments/1a40d3c86d69/code/body-015.txt",
      "chars": 528,
      "language": "text"
    },
    {
      "file": "attachments/1a40d3c86d69/code/body-016.txt",
      "chars": 601,
      "language": "text"
    },
    {
      "file": "attachments/1a40d3c86d69/code/body-017.txt",
      "chars": 618,
      "language": "text"
    },
    {
      "file": "attachments/1a40d3c86d69/code/body-018.txt",
      "chars": 367,
      "language": "text"
    },
    {
      "file": "attachments/1a40d3c86d69/code/body-019.txt",
      "chars": 750,
      "language": "text"
    }
  ],
  "images": [
    {
      "status": "ok",
      "url": "https://cdn.hashnode.com/res/hashnode/image/upload/v1770659671216/1cee86fb-2dc9-42ed-a66c-5294c8c11bc1.png",
      "bytes": 76606,
      "reused": true,
      "local_path": "attachments/1a40d3c86d69/images/000.png",
      "src": "https://cdn.hashnode.com/res/hashnode/image/upload/v1770659671216/1cee86fb-2dc9-42ed-a66c-5294c8c11bc1.png",
      "alt": "Transformations That Change Everything",
      "downloaded": true
    }
  ],
  "audio": [],
  "videos": [],
  "links": [
    {
      "href": "https://kambale.dev/transformations-that-change-everything.md",
      "label": "View as Markdown",
      "type": "a"
    },
    {
      "href": "https://hashnode.com/@kambale",
      "label": "W",
      "type": "a"
    },
    {
      "href": "https://twitter.com/WesKambale",
      "label": "",
      "type": "a"
    },
    {
      "href": "https://linkedin.com/in/WesKambale",
      "label": "",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/series/jax-series",
      "label": "Part of series JAX",
      "type": "a"
    },
    {
      "href": "https://jaxstack.ai/",
      "label": "JAX AI Stack Guide",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/tag/jax",
      "label": "# jax",
      "type": "a"
    }
  ],
  "repair_version": 2,
  "repaired_at": "2026-10-14T10:50:48Z",
  "source_url": "https://kambale.dev/transformations-that-change-everything",
  "papers": [],
  "actual_video_files": [],
  "body_chars": 14899,
  "expected_body_images": 1,
  "discovery_status": "ok",
  "delivery_status": "ok",
  "missing_reason": "",
  "raw_html": "attachments/1a40d3c86d69/pages/source.html",
  "local_html": "attachments/1a40d3c86d69/pages/article.html",
  "body_base_path": "attachments/1a40d3c86d69/pages/",
  "body_markdown_path": "attachments/1a40d3c86d69/pages/article.md"
}