{
  "source_file": "0930",
  "source_platform": "Hashnode",
  "author": "Wesley Kambale",
  "title_listed": "Why JAX? The NumPy You Know, But Faster",
  "keywords": "",
  "domain": "机器学习",
  "content_type": "深度技术解析",
  "is_paid_flag": "否",
  "url": "https://kambale.dev/why-jax-the-numpy-you-know-but-faster",
  "publish_date": "2026-02-02",
  "corpus_stage": "",
  "notes": "",
  "id": "01e58bea53bc",
  "collected_at": "2026-10-02T14:11:07",
  "title": "Why JAX? The NumPy You Know, But Faster",
  "body_markdown": "Why JAX? The NumPy You Know, But Faster\nHow JAX uses XLA to strap a rocket engine onto your math operations.\nUpdated\nJuly 29, 2026\n•\n10\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\nIf you've been doing machine learning in Python for any length of time, you've written code like this:\nimport numpy as np\n\nx = np.random.randn(1000, 1000)\ny = np.random.randn(1000, 1000)\nresult = np.dot(x, y)\nNumPy is comfortable. It's the first thing we reach for when we need to do math. But here's the uncomfortable truth: for the kind of work we're about to do: training neural networks, processing millions of samples, running on GPUs and TPUs, NumPy is holding us back.\nNot because NumPy is bad. It's genuinely excellent at what it was designed for. The problem is that NumPy was designed in 2005, before GPUs became the workhorses of machine learning, before TPUs existed, and before we needed to compute gradients of functions with millions of parameters.\nJAX is what NumPy would look like if we designed it today, knowing what we know now.\nBy the end of this article, we'll have:\nUnderstood\nwhy\nJAX is faster (not just that it is)\nWritten our first JAX code and seen the speedup ourselves\nLearned the one mental shift that trips up everyone coming from NumPy or PyTorch\nBuilt a working benchmark that proves the difference\nLet's get into it.\nThe Problem with \"Normal\" Python\nWhen you run Python code, the interpreter reads your instructions one line at a time, translates each one to machine code, executes it, then moves to the next line. This is called\ninterpretation\n.\nLine 1 → Translate → Execute\nLine 2 → Translate → Execute\nLine 3 → Translate → Execute\n...\nFor a script that processes a CSV file or serves a web page, this is fine. The overhead of interpretation is negligible compared to the actual work being done.\nBut matrix multiplication? That's different. When we multiply two 5000×5000 matrices, we're doing 125 billion floating-point operations. The \"translate → execute\" overhead for each operation adds up fast.\nNumPy helps by pushing the heavy lifting into compiled C code. When you call\nnp.dot\n()\n, Python hands off the work to optimized BLAS libraries that run at near-hardware speed. That's why NumPy is fast\ner\nthan pure Python.\nBut there's still a problem:\nPython is still orchestrating the operations\n. Every time you chain NumPy calls together:\nnp.dot\n()\n, then\nnp.sum()\n, then\nnp.exp()\n, Python has to:\nCall into C\nWait for the result\nCopy the result back to Python\nCall into C again for the next operation\nEach of those handoffs has overhead. And when you're doing this millions of times in a training loop, it adds up.\nHow JAX Fixes This: XLA Compilation\nJAX takes a different approach. Instead of executing operations one at a time, JAX can\ncompile your entire function\ninto a single optimized program using XLA (Accelerated Linear Algebra).\nHere's what that means in practice:\nRead entire function → Analyze → Optimize → Fuse operations → Execute once\nWhen JAX compiles a function, it:\nFuses operations\n: Instead of computing\na + b\n, storing the result, then computing\nresult * c\n, XLA fuses these into a single kernel that does\n(a + b) * c\nwithout intermediate storage.\nEliminates dead code\n: If you compute something but never use it, XLA removes it entirely.\nOptimizes memory access\n: XLA reorders operations to minimize cache misses and memory transfers.\nTargets your hardware\n: The same JAX code compiles to optimized instructions for CPU, GPU, or TPU without you changing anything.\nThe result? Code that runs 10x, 100x, sometimes 1000x faster than the NumPy equivalent.\nSetting Up\nLet's stop talking and start coding. We'll use Google Colab for this, it's free and gives us access to GPUs.\n# Install the JAX AI stack (includes JAX, Flax, Optax, and friends)\n!pip install -q jax-ai-stack\nNow let's verify our setup:\nimport jax\nimport jax.numpy as jnp\nimport numpy as np\nimport time\n\nprint(f\"JAX version: {jax.__version__}\")\nprint(f\"Available devices: {jax.devices()}\")\nIf you're on a GPU runtime, you should see something like:\nJAX version: 0.8.0\nAvailable devices: [CudaDevice(id=0)]\nIf you see\nCpuDevice\n, that's fine too; JAX still provides speedups on CPU through XLA compilation.\nThe First Mental Shift: Immutability\nBefore we benchmark anything, we need to talk about the one thing that trips up\neveryone\ncoming from NumPy or PyTorch.\nIn JAX, arrays are immutable. You cannot modify them in place.\nIn NumPy, this is perfectly normal:\n# NumPy: Mutable arrays\narr = np.zeros(5)\narr[0] = 42  # Modify in place\nprint(arr)   # [42. 0. 0. 0. 0.]\nIn JAX, this will raise an error:\n# JAX: This will FAIL\narr = jnp.zeros(5)\narr[0] = 42  #TypeError: JAX arrays are immutable\nWhy? Because immutability is what makes JAX's optimizations possible. If arrays can be modified from anywhere in your code, the compiler can't safely reorder operations or run them in parallel. By guaranteeing that arrays never change, JAX can aggressively optimize your code.\nSo how do we update arrays? We use the\n.at[].set()\nsyntax, which\nreturns a new array\nwith the modification:\n# JAX: The correct way\narr = jnp.zeros(5)\nnew_arr = arr.at[0].set(42)\n\nprint(arr)      # [0. 0. 0. 0. 0.] — Original unchanged\nprint(new_arr)  # [42. 0. 0. 0. 0.] — New array with the update\nThis feels wasteful at first, are we really copying the entire array just to change one element? In practice, no. JAX and XLA are smart enough to optimize this. But conceptually, you should think of it as creating a new array.\nHere's the full set of\n.at[]\noperations:\nx = jnp.array([1, 2, 3, 4, 5])\n\n# Set a value\nx.at[0].set(10)         # [10, 2, 3, 4, 5]\n\n# Add to a value\nx.at[0].add(10)         # [11, 2, 3, 4, 5]\n\n# Multiply a value\nx.at[0].multiply(10)    # [10, 2, 3, 4, 5]\n\n# Works with slices too\nx.at[1:3].set(99)       # [1, 99, 99, 4, 5]\nCommit this to memory. You'll use it constantly.\nThe Benchmark: NumPy vs JAX vs JAX+JIT\nNow let's prove that JAX is actually faster. We'll multiply two large matrices; this is the core operation in neural networks (every linear layer is a matrix multiplication).\nStep 1: Create the Data\n# Matrix size\nsize = 3000\n\n# Create random matrices with NumPy\nx_np = np.random.normal(size=(size, size)).astype(np.float32)\ny_np = np.random.normal(size=(size, size)).astype(np.float32)\n\n# Convert to JAX arrays\nx_jax = jnp.array(x_np)\ny_jax = jnp.array(y_np)\n\nprint(f\"Matrix shape: {x_np.shape}\")\nprint(f\"Total elements per matrix: {size * size:,}\")\nprint(f\"Operations for multiplication: {size ** 3:,}\")\nOutput:\nMatrix shape: (3000, 3000)\nTotal elements per matrix: 9,000,000\nOperations for multiplication: 27,000,000,000\nThat's 27 billion operations. Let's see who can do it fastest.\nStep 2: Benchmark NumPy\ndef matmul_numpy(x, y):\n    return np.dot(x, y)\n\n# Warmup\n_ = matmul_numpy(x_np, y_np)\n\n# Timed run\nstart = time.perf_counter()\nresult_np = matmul_numpy(x_np, y_np)\nnumpy_time = time.perf_counter() - start\n\nprint(f\"NumPy time: {numpy_time:.4f} seconds\")\nResult:\nNumPy time: 0.6294 seconds\nStep 3: Benchmark JAX (without JIT)\ndef matmul_jax(x, y):\n    return jnp.dot(x, y)\n\n# Warmup\n_ = matmul_jax(x_jax, y_jax).block_until_ready()\n\n# Timed run\nstart = time.perf_counter()\nresult_jax = matmul_jax(x_jax, y_jax).block_until_ready()\njax_time = time.perf_counter() - start\n\nprint(f\"JAX time (no JIT): {jax_time:.4f} seconds\")\nImportant\n: We call\n.block_until_ready()\nbecause JAX operations are\nasynchronous\n. When you call\njnp.dot\n()\n, JAX immediately returns a \"future\" and continues executing Python code while the GPU works in the background. Without\nblock_until_ready()\n, we'd be timing how fast JAX can\ndispatch\nthe operation, not how fast it actually\nruns\n.\nResult:\nJAX time (no JIT): 0.0094 seconds\nStep 4: Benchmark JAX with JIT Compilation\nHere's where JAX shows its true power. We add a single decorator:\n@jax.jit\ndef matmul_jax_jit(x, y):\n    return jnp.dot(x, y)\n\n# First call: JAX traces and compiles the function\n# This takes a moment, so we don't include it in the benchmark\nprint(\"Compiling...\")\n_ = matmul_jax_jit(x_jax, y_jax).block_until_ready()\nprint(\"Done.\")\n\n# Timed run (using the compiled version)\nstart = time.perf_counter()\nresult_jit = matmul_jax_jit(x_jax, y_jax).block_until_ready()\njit_time = time.perf_counter() - start\n\nprint(f\"JAX time (with JIT): {jit_time:.4f} seconds\")\nThe first call to a\n@jax.jit\nfunction is slow because JAX is\ntracing\nyour function, figuring out what operations it contains, and then\ncompiling\nit with XLA. Subsequent calls use the compiled version and are extremely fast.\nStep 5: Compare Results\nprint(\"RESULTS\")\nprint(f\"NumPy:          {numpy_time:.4f} seconds\")\nprint(f\"JAX (no JIT):   {jax_time:.4f} seconds\")\nprint(f\"JAX (with JIT): {jit_time:.4f} seconds\")\nprint(f\"\\nSpeedup (JIT vs NumPy): {numpy_time / jit_time:.1f}x\")\nTypical results on a Colab GPU:\nRESULTS\nNumPy:          0.6294 seconds\nJAX (no JIT):   0.0094 seconds\nJAX (with JIT): 0.0198 seconds\n\nSpeedup (JAX vs NumPy):     67.2x\nSpeedup (JIT vs NumPy):     31.8x\nSpeedup (JIT vs JAX):       0.5x\nWhat Just Happened?\nLet's break down why the JIT version is so much faster:\nHardware acceleration\n: JAX moved the computation to the GPU, which has thousands of cores optimized for parallel math.\nXLA compilation\n: Even on GPU, the JIT version is faster than raw JAX because XLA fuses operations and optimizes memory access patterns.\nNo Python overhead\n: Once compiled, the function runs entirely in native code. Python is only involved in dispatching the call.\nThe key insight is that\n@jax.jit\ndoesn't just run your code on a GPU; it fundamentally changes\nhow\nyour code runs.\nThe Randomness Trap (Bonus Lesson)\nThere's one more gotcha that catches everyone early on. Try this:\n@jax.jit\ndef broken_random():\n    return np.random.randn(5)  # Using NumPy's random\n\nresult1 = broken_random()\nresult2 = broken_random()\nprint(f\"First call:  {result1}\")\nprint(f\"Second call: {result2}\")\nYou'll notice that\nresult1\nand\nresult2\nare\nidentical\n. The random numbers got \"baked in\" during compilation.\nJAX requires\nexplicit random state\nmanagement. Here's the correct way:\nfrom jax import random\n\n@jax.jit\ndef correct_random(key):\n    return random.normal(key, shape=(5,))\n\n# Create a PRNG key\nkey = random.PRNGKey(42)\n\n# Split the key for each use\nkey, subkey1 = random.split(key)\nresult1 = correct_random(subkey1)\n\nkey, subkey2 = random.split(key)\nresult2 = correct_random(subkey2)\n\nprint(f\"First call:  {result1}\")\nprint(f\"Second call: {result2}\")\nNow you get different random numbers each time. We'll cover this pattern in depth when we build neural networks, but for now, just remember:\nnever use\nnp.random\ninside JIT-compiled functions.\nExercises\nBefore moving on, try these:\nBreak the rules\n: Try to modify a JAX array in place (\nx[0] = 1\n). Read the error message carefully: JAX errors are verbose but informative.\nVary the size\n: Run the benchmark with different matrix sizes (1000, 2000, 5000). How does the speedup change?\nChain operations\n: Write a function that does multiple operations (\njnp.dot\n, then\njnp.sum\n, then\njnp.exp\n). Compare JIT vs non-JIT. The speedup should be even larger because XLA fuses the operations.\nWhat's Next\nWe've established\nwhy\nJAX is fast and seen the proof. But speed is only half the story. Next week, we'll explore\ntransformations\n; the features that make JAX genuinely different from NumPy, not just faster.\nSpecifically, we'll cover:\njax.vmap\n: Automatic vectorization that eliminates for-loops\njax.grad\n: Automatic differentiation that makes backpropagation trivial\nThese two functions are why JAX has become the framework of choice for machine learning research. Once you understand them, you'll never look at NumPy the same way again.\nQuick Reference\nimport jax\nimport jax.numpy as jnp\nfrom jax import random\n\n# Basic array operations (same as NumPy)\nx = jnp.array([1, 2, 3])\ny = jnp.zeros((3, 3))\nz = jnp.dot(a, b)\n\n# Updating arrays (immutable style)\nnew_x = x.at[0].set(99)\nnew_x = x.at[1:].add(10)\n\n# JIT compilation\n@jax.jit\ndef fast_function(x):\n    return jnp.dot(x, x.T)\n\n# Explicit randomness\nkey = random.PRNGKey(0)\nkey, subkey = random.split(key)\nsamples = random.normal(subkey, shape=(100,))\n\n# Block for accurate timing\nresult = fast_function(x).block_until_ready()\nNext week\n:\nTransformations That Change Everything—Automatic Vectorization with vmap and Gradients with grad\nResources:\nJAX Documentation\nJAX AI Stack\nNotebook\n#\njax\n#\nxla\n#\nnumpy",
  "metadata": {
    "title": "Why JAX? The NumPy You Know, But Faster",
    "description": "How JAX uses XLA to strap a rocket engine onto your math operations.",
    "keywords": null,
    "author": null,
    "og:title": "Why JAX? The NumPy You Know, But Faster",
    "og:description": "How JAX uses XLA to strap a rocket engine onto your math operations.",
    "og:url": "https://kambale.dev/why-jax-the-numpy-you-know-but-faster",
    "og:site_name": "Wesley Kambale",
    "og:image": "https://cdn.hashnode.com/res/hashnode/image/upload/v1770040892040/a8adacb0-31ca-4dd9-bc27-d29584712241.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": "Why JAX? The NumPy You Know, But Faster",
    "twitter:description": "How JAX uses XLA to strap a rocket engine onto your math operations.",
    "twitter:image": "https://cdn.hashnode.com/res/hashnode/image/upload/v1770040892040/a8adacb0-31ca-4dd9-bc27-d29584712241.png",
    "article:published_time": "2026-02-02T14:15:09.495Z",
    "article:modified_time": "2026-07-29T00:58:44.048Z",
    "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/01e58bea53bc/code/body-000.txt",
      "chars": 106,
      "language": "text"
    },
    {
      "file": "attachments/01e58bea53bc/code/body-001.txt",
      "chars": 91,
      "language": "text"
    },
    {
      "file": "attachments/01e58bea53bc/code/body-002.txt",
      "chars": 75,
      "language": "text"
    },
    {
      "file": "attachments/01e58bea53bc/code/body-003.txt",
      "chars": 97,
      "language": "text"
    },
    {
      "file": "attachments/01e58bea53bc/code/body-004.txt",
      "chars": 153,
      "language": "text"
    },
    {
      "file": "attachments/01e58bea53bc/code/body-005.txt",
      "chars": 57,
      "language": "text"
    },
    {
      "file": "attachments/01e58bea53bc/code/body-006.txt",
      "chars": 106,
      "language": "text"
    },
    {
      "file": "attachments/01e58bea53bc/code/body-007.txt",
      "chars": 91,
      "language": "text"
    },
    {
      "file": "attachments/01e58bea53bc/code/body-008.txt",
      "chars": 191,
      "language": "text"
    },
    {
      "file": "attachments/01e58bea53bc/code/body-009.txt",
      "chars": 282,
      "language": "text"
    },
    {
      "file": "attachments/01e58bea53bc/code/body-010.txt",
      "chars": 406,
      "language": "text"
    },
    {
      "file": "attachments/01e58bea53bc/code/body-011.txt",
      "chars": 110,
      "language": "text"
    },
    {
      "file": "attachments/01e58bea53bc/code/body-012.txt",
      "chars": 254,
      "language": "text"
    },
    {
      "file": "attachments/01e58bea53bc/code/body-013.txt",
      "chars": 27,
      "language": "text"
    },
    {
      "file": "attachments/01e58bea53bc/code/body-014.txt",
      "chars": 297,
      "language": "text"
    },
    {
      "file": "attachments/01e58bea53bc/code/body-015.txt",
      "chars": 34,
      "language": "text"
    },
    {
      "file": "attachments/01e58bea53bc/code/body-016.txt",
      "chars": 491,
      "language": "text"
    },
    {
      "file": "attachments/01e58bea53bc/code/body-017.txt",
      "chars": 231,
      "language": "text"
    },
    {
      "file": "attachments/01e58bea53bc/code/body-018.txt",
      "chars": 203,
      "language": "text"
    },
    {
      "file": "attachments/01e58bea53bc/code/body-019.txt",
      "chars": 203,
      "language": "text"
    },
    {
      "file": "attachments/01e58bea53bc/code/body-020.txt",
      "chars": 378,
      "language": "text"
    },
    {
      "file": "attachments/01e58bea53bc/code/body-021.txt",
      "chars": 528,
      "language": "text"
    }
  ],
  "images": [
    {
      "status": "ok",
      "url": "https://cdn.hashnode.com/res/hashnode/image/upload/v1770040892040/a8adacb0-31ca-4dd9-bc27-d29584712241.png",
      "bytes": 66809,
      "reused": true,
      "local_path": "attachments/01e58bea53bc/images/000.png",
      "src": "https://cdn.hashnode.com/res/hashnode/image/upload/v1770040892040/a8adacb0-31ca-4dd9-bc27-d29584712241.png",
      "alt": "Why JAX? The NumPy You Know, But Faster",
      "downloaded": true
    }
  ],
  "audio": [],
  "videos": [],
  "links": [
    {
      "href": "https://kambale.dev/why-jax-the-numpy-you-know-but-faster.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": "http://np.dot/",
      "label": "np.dot",
      "type": "a"
    },
    {
      "href": "http://jnp.dot/",
      "label": "jnp.dot",
      "type": "a"
    },
    {
      "href": "https://jax.readthedocs.io/",
      "label": "JAX Documentation",
      "type": "a"
    },
    {
      "href": "https://jaxstack.ai/",
      "label": "JAX AI Stack",
      "type": "a"
    },
    {
      "href": "https://colab.research.google.com/drive/1f-qg4vlfBHSQSdEdhfgDdRA4k1dPtAqW?usp=sharing",
      "label": "Notebook",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/tag/jax",
      "label": "# jax",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/tag/xla",
      "label": "# xla",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/tag/numpy",
      "label": "# numpy",
      "type": "a"
    }
  ],
  "repair_version": 2,
  "repaired_at": "2026-10-14T10:50:47Z",
  "source_url": "https://kambale.dev/why-jax-the-numpy-you-know-but-faster",
  "papers": [],
  "actual_video_files": [],
  "body_chars": 12952,
  "expected_body_images": 1,
  "discovery_status": "ok",
  "delivery_status": "ok",
  "missing_reason": "",
  "raw_html": "attachments/01e58bea53bc/pages/source.html",
  "local_html": "attachments/01e58bea53bc/pages/article.html",
  "body_base_path": "attachments/01e58bea53bc/pages/",
  "body_markdown_path": "attachments/01e58bea53bc/pages/article.md"
}