{
  "source_file": "0930",
  "source_platform": "Hashnode",
  "author": "Wesley Kambale",
  "title_listed": "NumPy vs JAX: Optimizing Performance and Functionality",
  "keywords": "",
  "domain": "机器学习",
  "content_type": "深度技术解析",
  "is_paid_flag": "否",
  "url": "https://kambale.dev/numpy-vs-jax",
  "publish_date": "2023-05-08",
  "corpus_stage": "",
  "notes": "",
  "id": "5b2d2cccdd1c",
  "collected_at": "2026-10-02T14:11:39",
  "title": "NumPy vs JAX: Optimizing Performance and Functionality",
  "body_markdown": "NumPy vs JAX: Optimizing Performance and Functionality\nA beginner's guide and comparison of NumPy and JAX for scientific computing and machine learning\nUpdated\nJuly 18, 2023\n•\n6\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\nMachine Learning\nIntroduction\nNumpy and Jax are both Python libraries that are widely used in numerical computing and scientific computing. While they have similar functionalities, there are some key differences between them that make Jax particularly useful for machine learning applications. In this tutorial, we will explore some of the differences between Numpy and Jax and provide code snippets to illustrate these differences.\nInstallation\nBoth Numpy and Jax can be installed using pip. To install Numpy, you can run the following command in your terminal:\npip install numpy\nTo install Jax, you can run the following command:\npip install jax\nNote that Jax also requires the\njaxlib\nlibrary, which can be installed using the following command:\npip install jaxlib\nFunctionality\nLet us take a look at how to write\nNumPy\nand\nJax\nwhen performing simple and complex scientific computing and machine learning.\nArray creation\nNumpy and Jax both provide functions for creating arrays. The simplest way to create an array in Numpy is to use the\narray\nfunction:\nimport\nnumpy\nas\nnp\n# create a 1D array\nnpArrOne = np.array([\n1\n,\n2\n,\n3\n])\nprint(npArrOne)\n# create a 2D array\nnpArrTwo = np.array([[\n1\n,\n2\n], [\n3\n,\n4\n]])\nprint(npArrTwo)\nIn Jax, the corresponding function is\njnp.array\n:\nimport\njax.numpy\nas\njnp\n# create a 1D array\njxArrOne = jnp.array([\n1\n,\n2\n,\n3\n])\nprint(jxArrOne)\n# create a 2D array\njxArrTwo = jnp.array([[\n1\n,\n2\n], [\n3\n,\n4\n]])\nprint(jxArrTwo)\nNote that the syntax is very similar in both libraries, but in Jax we import the numpy functions from the\njax\npackage rather than the\nnumpy\npackage.\nArray indexing\nBoth Numpy and Jax provide powerful indexing capabilities for arrays. The basic syntax is the same in both libraries:\n# indexing in numpy\nnpArr = np.array([\n1\n,\n2\n,\n3\n])\nprint(npArr[\n0\n])\n# output: 1\n# indexing in jax\njxArr = jnp.array([\n1\n,\n2\n,\n3\n])\nprint(jxArr[\n0\n])\n# output: 1\nHowever, Jax provides some additional indexing features that are particularly useful for machine learning. For example, Jax provides a function called\njnp.index_update\nthat allows you to update an element of an array by index:\n# update an element in numpy\nnpArr = np.array([\n1\n,\n2\n,\n3\n])\nnpArr[\n0\n] =\n4\nprint(npArr)\n# output: [4, 2, 3]\n# update an element in jax\njxArr = jnp.array([\n1\n,\n2\n,\n3\n])\njxArr = jnp.index_update(b,\n0\n,\n4\n)\nprint(jxArr)\n# output: [4, 2, 3]\nNote that in Jax, we cannot modify an array in-place, so we need to reassign the result of\njnp.index_update\nto the original array.\nArray broadcasting\nOne of the most powerful features of Numpy and Jax is their ability to broadcast arrays. Broadcasting allows you to perform operations on arrays with different shapes and sizes. Here is an example of broadcasting in Numpy:\n# broadcasting in numpy\nnpArrOne = np.array([[\n1\n,\n2\n], [\n3\n,\n4\n]])\nnpArrTwo = np.array([\n10\n,\n20\n])\nprint(npArrOne + npArrTwo)\n# output: [[11, 22], [13, 24]]\nIn Jax, broadcasting works in a similar way:\n# broadcasting in jax\njxArrOne = jnp.array([[\n1\n,\n2\n], [\n3\n,\n4\n]])\njxArrTwo = jnp.array([\n10\n,\n20\n])\nprint(jxArrOne + jxArrTwo)\n# output: [[11, 22], [13, 24]]\nIn the above example, Numpy and Jax are able to add the 2D array\nnpArrOne\nand the 1D array\nnpArrTwo\nby automatically broadcasting the dimensions of\nnpArrTwo\nto match the dimensions of\nnpArrOne\n.\nFor Jax, it provides some additional broadcasting features that are particularly useful for machine learning. For instance, Jax provides a function called\njnp.vmap\nthat allows you to apply a function to multiple inputs using broadcasting:\n# vmap in Jax\ndef\nadd\n(\nx, y\n):\nreturn\nx + y\n\njxArrOne = jnp.array([[\n1\n,\n2\n], [\n3\n,\n4\n]])\njxArrTwo = jnp.array([[\n10\n,\n20\n], [\n30\n,\n40\n]])\n# apply add to multiple inputs using broadcasting\njxArrThree = jnp.vmap(add)(jxArrOne, jxArrTwo)\nprint(jxArrThree)\n# output: [[11, 22], [33, 44]]\nIn the code snippet above, we define a function\nadd\nthat adds two arrays element-wise. We then use\njnp.vmap\nto apply this function to two arrays\njxArrOne\nand\njxArrTwo\n, which have the same shape. The result is a new array\njxArrThree\nthat has the same shape as\njxArrOne\nand\njxArrTwo\n, where each element of\njxArrThree\nis the sum of the corresponding elements of\njxArrOne\nand\njxArrTwo\n.\nPerformance\nOne of the key advantages of Jax over Numpy is its ability to compile code using the XLA compiler. This allows Jax to run code on GPUs and TPUs, which can significantly improve performance for large-scale machine learning applications.\nNumPy\nJAX\nHardware\nCPU\nCPU, GPU, TPU\nExecution\nSynchronously\nAsynchronously\nParallel computation\nNo\nYes\nIn the example below, we show how Jax and NumPy can be used to accelerate a simple matrix multiplication:\nNumPy:\n# matrix multiplication in numpy\nimport\nnumpy\nas\nnp\nimport\ntime\n\nnpArrOne = np.random.rand(\n1000\n,\n1000\n)\nnpArrOne = np.random.rand(\n1000\n,\n1000\n)\n\nstart = time.time()\nnpArrThree = np.dot(npArrOne, npArrOne)\nend = time.time()\n\nprint(\n\"NumPy time:\"\n, end - start)\nOutput:\nNumPy time: 0.5427792072296143\nJAX:\n# matrix multiplication in jax\nimport\njax.numpy\nas\njnp\nfrom\njax\nimport\njit\n\njxArrOne = jnp.random.rand(\n1000\n,\n1000\n)\njxArrTwo = jnp.random.rand(\n1000\n,\n1000\n)\n@jit\ndef\nmatmul\n(\njxArrOne, jxArrTwo\n):\nreturn\njnp.dot(jxArrOne, jxArrTwo)\n\nstart = time.time()\njxArrThree = matmul(jxArrOne, jxArrOne)\nend = time.time()\n\nprint(\n\"JAX time:\"\n, end - start)\nOutput:\nJAX time: 0.03486919403076172\nIn the code snippets above, we generate two random matrices NumPy and Jax of size 1000 x 1000, and we use the\nnp.dot\nfunction to perform matrix multiplication in Numpy, and the\njnp.dot\nfunction to perform matrix multiplication in Jax. We also use the\n@jit\ndecorator to compile the\nmatmul\nfunction using the XLA compiler.\nFrom the output of time given, you can see that Jax is significantly faster than NumPy in terms of performance.\nConclusion\nSummary table of the similarities and differences between NumPy and JAX\nFeature\nNumPy\nJAX\nFunctionality\nBasic array operations, linear algebra, statistics, image processing, Fourier transform\nAll of NumPy's functionality, plus automatic differentiation, just-in-time compilation, parallel execution, and stateful computations\nPerformance\nGood for simple tasks\nExcellent for complex tasks\nScalability\nGood for small to medium datasets\nExcellent for large datasets\nSo, we have explored some of the differences between Numpy and Jax, two powerful Python libraries that are widely used in numerical computing and scientific computing.\nWhile they have similar functionalities, Jax provides some additional features that are particularly useful for machine learning applications, including powerful indexing and broadcasting capabilities, as well as the ability to run code on GPUs and TPUs using the XLA compiler, which makes it particularly useful for machine learning applications.\nBy understanding the strengths and weaknesses of each library, you can choose the one that best suits your needs, and maximize the performance and functionality of your scientific computing projects.\nGood luck!\n#\nnumpy\n#\njax\n#\nmachine-learning\n#\nscientific-computing",
  "metadata": {
    "title": "NumPy vs JAX: Optimizing Performance and Functionality",
    "description": "A beginner's guide and comparison of NumPy and JAX for scientific computing and machine learning",
    "keywords": null,
    "author": null,
    "og:title": "NumPy vs JAX: Optimizing Performance and Functionality",
    "og:description": "A beginner's guide and comparison of NumPy and JAX for scientific computing and machine learning",
    "og:url": "https://kambale.dev/numpy-vs-jax",
    "og:site_name": "Wesley Kambale",
    "og:image": "https://cdn.hashnode.com/res/hashnode/image/upload/v1689698427934/a8619a7a-566a-4705-af70-355a4a440118.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": "NumPy vs JAX: Optimizing Performance and Functionality",
    "twitter:description": "A beginner's guide and comparison of NumPy and JAX for scientific computing and machine learning",
    "twitter:image": "https://cdn.hashnode.com/res/hashnode/image/upload/v1689698427934/a8619a7a-566a-4705-af70-355a4a440118.png",
    "article:published_time": "2023-05-08T17:57:39.873Z",
    "article:modified_time": "2023-07-18T16:44:32.980Z",
    "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/5b2d2cccdd1c/code/body-000.txt",
      "chars": 18,
      "language": "text"
    },
    {
      "file": "attachments/5b2d2cccdd1c/code/body-001.txt",
      "chars": 16,
      "language": "text"
    },
    {
      "file": "attachments/5b2d2cccdd1c/code/body-002.txt",
      "chars": 19,
      "language": "text"
    },
    {
      "file": "attachments/5b2d2cccdd1c/code/body-003.txt",
      "chars": 162,
      "language": "text"
    },
    {
      "file": "attachments/5b2d2cccdd1c/code/body-004.txt",
      "chars": 169,
      "language": "text"
    },
    {
      "file": "attachments/5b2d2cccdd1c/code/body-005.txt",
      "chars": 156,
      "language": "text"
    },
    {
      "file": "attachments/5b2d2cccdd1c/code/body-006.txt",
      "chars": 231,
      "language": "text"
    },
    {
      "file": "attachments/5b2d2cccdd1c/code/body-007.txt",
      "chars": 152,
      "language": "text"
    },
    {
      "file": "attachments/5b2d2cccdd1c/code/body-008.txt",
      "chars": 151,
      "language": "text"
    },
    {
      "file": "attachments/5b2d2cccdd1c/code/body-009.txt",
      "chars": 278,
      "language": "text"
    },
    {
      "file": "attachments/5b2d2cccdd1c/code/body-010.txt",
      "chars": 255,
      "language": "text"
    },
    {
      "file": "attachments/5b2d2cccdd1c/code/body-011.txt",
      "chars": 31,
      "language": "text"
    },
    {
      "file": "attachments/5b2d2cccdd1c/code/body-012.txt",
      "chars": 343,
      "language": "text"
    },
    {
      "file": "attachments/5b2d2cccdd1c/code/body-013.txt",
      "chars": 30,
      "language": "text"
    }
  ],
  "images": [
    {
      "status": "ok",
      "url": "https://cdn.hashnode.com/res/hashnode/image/upload/v1689698427934/a8619a7a-566a-4705-af70-355a4a440118.png",
      "bytes": 118129,
      "reused": true,
      "local_path": "attachments/5b2d2cccdd1c/images/000.png",
      "src": "https://cdn.hashnode.com/res/hashnode/image/upload/v1689698427934/a8619a7a-566a-4705-af70-355a4a440118.png",
      "alt": "NumPy vs JAX: Optimizing Performance and Functionality",
      "downloaded": true
    }
  ],
  "audio": [],
  "videos": [],
  "links": [
    {
      "href": "https://kambale.dev/numpy-vs-jax.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/ml-series",
      "label": "Part of series Machine Learning",
      "type": "a"
    },
    {
      "href": "http://np.dot/",
      "label": "np.dot",
      "type": "a"
    },
    {
      "href": "http://jnp.dot/",
      "label": "jnp.dot",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/tag/numpy",
      "label": "# numpy",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/tag/jax",
      "label": "# jax",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/tag/machine-learning",
      "label": "# machine-learning",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/tag/scientific-computing",
      "label": "# scientific-computing",
      "type": "a"
    }
  ],
  "repair_version": 2,
  "repaired_at": "2026-10-14T10:50:52Z",
  "source_url": "https://kambale.dev/numpy-vs-jax",
  "papers": [],
  "actual_video_files": [],
  "body_chars": 7793,
  "expected_body_images": 1,
  "discovery_status": "ok",
  "delivery_status": "ok",
  "missing_reason": "",
  "raw_html": "attachments/5b2d2cccdd1c/pages/source.html",
  "local_html": "attachments/5b2d2cccdd1c/pages/article.html",
  "body_base_path": "attachments/5b2d2cccdd1c/pages/",
  "body_markdown_path": "attachments/5b2d2cccdd1c/pages/article.md"
}