{"id":148907,"date":"2023-11-15T13:04:30","date_gmt":"2023-11-15T18:04:30","guid":{"rendered":"http:\/\/www.bu.edu\/tech\/?page_id=148907"},"modified":"2024-01-31T09:40:38","modified_gmt":"2024-01-31T14:40:38","slug":"jax","status":"publish","type":"page","link":"https:\/\/www.bu.edu\/tech\/support\/research\/software-and-programming\/common-languages\/python\/python-ml\/jax\/","title":{"rendered":"JAX"},"content":{"rendered":"<p><a name=\"top\"><\/a><\/p>\n<div style=\"width:100%;\">\n<div style=\"float:right;width:45%;\"><img loading=\"lazy\" src=\"\/tech\/files\/2023\/12\/jax_new.png\" alt=\"\" width=\"600\" height=\"356\" class=\"alignnone size-full wp-image-149628\" \/><\/div>\n<div style=\"float:none;\"><a href=\"https:\/\/github.com\/google\/jax\">Jax<\/a> is a library for high performance numerical computing with GPUs. It has a wide range of capabilities, from autodifferentiation to a GPU-enabled version of the Numpy library.<\/div>\n<p>&nbsp;<br \/>\n<strong>Sections<\/strong><\/p>\n<ul>\n<li><a href=\"#install\">Install for Yourself<\/a>\n<ul>\n<li><a href=\"#venv\">Python Virtual Environment<\/a><\/li>\n<li><a href=\"#conda\">Conda Environment<\/a><\/li>\n<\/ul>\n<\/li>\n<\/ul>\n<p>&nbsp;<\/p>\n<p><a name=\"install\"><\/a><\/p>\n<h2>Install for Yourself<\/h2>\n<p>It&#8217;s recommended that you install Jax in your own virtual environment or conda environment. The sections below explain how to install CPU and GPU versions of Jax into your environment. These instructions have been taken from the <a href=\"https:\/\/jax.readthedocs.io\/en\/latest\/installation.html\">Jax installation page<\/a>.<\/p>\n<p><a name=\"venv\"><\/a><\/p>\n<h3>Python Virtual Environment<\/h3>\n<ol>\n<li>Create and activate your virtualenv as explained on the <a href=\"https:\/\/www.bu.edu\/tech\/support\/research\/software-and-programming\/common-languages\/python\/python-installs\/virtualenv\/\">Python Installs page<\/a>.<\/li>\n<li>Install the CPU-Only <u>or<\/u> CPU+GPU versions of Jax with <code><b>pip<\/b><\/code>:<br \/>\n<h4>CPU-Only<\/h4>\n<pre><code class=\"code-block\"><span class=\"prompt\">(my_newenv) [rcs@scc1 ~]<\/span> <span class=\"command\">pip install jax[cpu]<\/span><\/code><\/pre>\n<p><u style=\"color:red\";><center><\/p>\n<h2 style=\"color:red\";>OR<\/h2>\n<p><\/center><\/u><\/p>\n<h4>CPU + GPU<\/h4>\n<pre><code class=\"code-block\"><span class=\"prompt\">(my_newenv) [rcs@scc1 ~]<\/span> <span class=\"command\">pip install --upgrade \"jax[cuda12_pip]\" -f https:\/\/storage.googleapis.com\/jax-releases\/jax_cuda_releases.html<\/span><\/code><\/pre>\n<\/li>\n<\/ol>\n<p>&nbsp;<\/p>\n<p><a name=\"conda\"><\/a><\/p>\n<h3>Conda Environment<\/h3>\n<ol>\n<li>Create and activate your conda env as explained on the <a href=\"https:\/\/www.bu.edu\/tech\/support\/research\/software-and-programming\/common-languages\/python\/python-installs\/conda\/\">Miniconda Installs page<\/a>.<\/li>\n<li>Install the CPU-Only <u>or<\/u> CPU+GPU versions of Jax with <code><b>conda install<\/b><\/code>:<br \/>\n<h4>CPU-Only<\/h4>\n<pre><code class=\"code-block\"><span class=\"prompt\">(my_conda_env) [rcs@scc1 ~]<\/span> <span class=\"command\">conda install jax -c conda-forge<\/span><\/code><\/pre>\n<p><u style=\"color:red\";><center><\/p>\n<h2 style=\"color:red\";>OR<\/h2>\n<p><\/center><\/u><\/p>\n<h4>CPU + GPU<\/h4>\n<pre><code class=\"code-block\"><span class=\"prompt\">(my_conda_env) [rcs@scc1 ~]<\/span> <span class=\"command\">conda install jaxlib=*=*cuda* jax cuda-nvcc -c conda-forge -c nvidia<\/span><\/code><\/pre>\n<\/li>\n<\/ol>\n<p><a href=\"#top\">Back to top<\/a><\/p>\n<p><script>\r\nconst date = new Date(document.lastModified);\r\ndocument.getElementById(\"demo\").innerHTML = date;\r\n<\/script><\/p>\n","protected":false},"excerpt":{"rendered":"<p>Jax is a library for high performance numerical computing with GPUs. It has a wide range of capabilities, from autodifferentiation to a GPU-enabled version of the Numpy library. &nbsp; Sections Install for Yourself Python Virtual Environment Conda Environment &nbsp; Install for Yourself It&#8217;s recommended that you install Jax in your own virtual environment or conda&#8230;<\/p>\n","protected":false},"author":1692,"featured_media":0,"parent":147628,"menu_order":4,"comment_status":"closed","ping_status":"closed","template":"","meta":[],"_links":{"self":[{"href":"https:\/\/www.bu.edu\/tech\/wp-json\/wp\/v2\/pages\/148907"}],"collection":[{"href":"https:\/\/www.bu.edu\/tech\/wp-json\/wp\/v2\/pages"}],"about":[{"href":"https:\/\/www.bu.edu\/tech\/wp-json\/wp\/v2\/types\/page"}],"author":[{"embeddable":true,"href":"https:\/\/www.bu.edu\/tech\/wp-json\/wp\/v2\/users\/1692"}],"replies":[{"embeddable":true,"href":"https:\/\/www.bu.edu\/tech\/wp-json\/wp\/v2\/comments?post=148907"}],"version-history":[{"count":17,"href":"https:\/\/www.bu.edu\/tech\/wp-json\/wp\/v2\/pages\/148907\/revisions"}],"predecessor-version":[{"id":150185,"href":"https:\/\/www.bu.edu\/tech\/wp-json\/wp\/v2\/pages\/148907\/revisions\/150185"}],"up":[{"embeddable":true,"href":"https:\/\/www.bu.edu\/tech\/wp-json\/wp\/v2\/pages\/147628"}],"wp:attachment":[{"href":"https:\/\/www.bu.edu\/tech\/wp-json\/wp\/v2\/media?parent=148907"}],"curies":[{"name":"wp","href":"https:\/\/api.w.org\/{rel}","templated":true}]}}