JAX - Revision history

Pigs: /* Example shell.nix with GPU support */ syntax highlight code block, additional crosslinking and formatting

2025-12-09T19:11:28Z

Example shell.nix with GPU support: syntax highlight code block, additional crosslinking and formatting

← Older revision		Revision as of 19:11, 9 December 2025
Line 41:		Line 41:
	You can test that JAX is using the GPU as intended with		You can test that JAX is using the GPU as intended with

	# after version 0.8.0		<syntaxHighlight lang=bash>
	python -c "import jax.extend as ex; print(ex.backend.get_backend().platform)"		# after version 0.8.0
			python -c "import jax.extend as ex; print(ex.backend.get_backend().platform)"

	# before version 0.8.0		# before version 0.8.0
	python -c "from jax.lib import xla_bridge; print(xla_bridge.get_backend().platform)"		python -c "from jax.lib import xla_bridge; print(xla_bridge.get_backend().platform)"
			</syntaxHighlight>

	It should print either "cpu", "gpu", or "tpu".		It should print either <code>cpu</code>, <code>gpu</code>, or <code>tpu</code>.

	Note ~~that hydra~~ may not cache <code>jaxlibWithCuda</code> builds on cache.nixos.org since CUDA is "unfree." @samuela publishes builds on a public cachix [https://app.cachix.org/cache/ploop#pull ploop] cache. These are periodically built and pushed from [https://github.com/samuela/nixpkgs-upkeep/ nixpkgs-upkeep].		{{Note\| [[Hydra]] may not cache <code>jaxlibWithCuda</code> builds on cache.nixos.org since CUDA is [[unfree software]]. @samuela publishes builds on a public cachix [https://app.cachix.org/cache/ploop#pull ploop] cache. These are periodically built and pushed from [https://github.com/samuela/nixpkgs-upkeep/ nixpkgs-upkeep].}}

	== FAQ ==		== FAQ ==

Arunoruto: jax.lib.xla_bridge module will be removed in JAX v0.9.0; all its APIs were deprecated and removed by JAX v0.8.0

2025-12-09T14:31:21Z

jax.lib.xla_bridge module will be removed in JAX v0.9.0; all its APIs were deprecated and removed by JAX v0.8.0

← Older revision		Revision as of 14:31, 9 December 2025
Line 41:		Line 41:
	You can test that JAX is using the GPU as intended with		You can test that JAX is using the GPU as intended with

			# after version 0.8.0
			python -c "import jax.extend as ex; print(ex.backend.get_backend().platform)"

			# before version 0.8.0
	python -c "from jax.lib import xla_bridge; print(xla_bridge.get_backend().platform)"		python -c "from jax.lib import xla_bridge; print(xla_bridge.get_backend().platform)"

Pigs: Remove out of date note, add link to nixpkgs and other wiki pages

2025-05-16T17:10:23Z

Remove out of date note, add link to nixpkgs and other wiki pages

← Older revision		Revision as of 17:10, 16 May 2025
Line 1:		Line 1:
	[https://github.com/google/jax JAX] is a framework for program transformation, esp. for automatic differentiation and machine learning. It's available ~~on Nix/NixOS~~ in the <code>python3Packages.{jax, jaxlib, jaxlibWithCuda}</code> packages.		[https://github.com/google/jax JAX] is a framework for program transformation, esp. for automatic differentiation and machine learning. It's available in [[Nixpkgs]] in the <code>python3Packages.{{{nixos:package\|python3Packages.jax\|jax}}, {{nixos:package\|python3%20jaxlib\|jaxlib}}, {{nixos:package\|python3*.jaxlibWithCuda\|jaxlibWithCuda}}}</code> packages.

	{{tip\|1='''Cache''': Using the [https://app.cachix.org/cache/nix-community nix-community cache] is recommended! It will save you valuable time and electrons. Getting set up should be as simple as <code>cachix use nix-community</code>. See the [[CUDA]] wiki page for more info.		{{tip\|1='''Cache''': Using the [https://app.cachix.org/cache/nix-community nix-community cache] is recommended! It will save you valuable time and electrons. Getting set up should be as simple as <code>cachix use nix-community</code>. See the [[CUDA]] wiki page for more info.
	}}		}}

	~~NOTE: JAX requires Python 3.9, the current version of <code>python3</code> in nixpkgs (as of 9/4/2021). JAX is currently only packaged for x86_64-linux (send a PR for your platform!).~~

	== Example shell.nix, CPU only ==		== Example shell.nix, CPU only ==
Line 21:		Line 19:

	== Example shell.nix with GPU support ==		== Example shell.nix with GPU support ==
	JAX defers execution to the jaxlib library for execution. In order to use GPU support you'll need a NVIDIA GPU and OpenGL. In your <code>/etc/nixos/configuration.nix</code>:		JAX defers execution to the jaxlib library for execution. In order to use GPU support you'll need a [[NVIDIA]] GPU and [[OpenGL]]. In your <code>/etc/nixos/configuration.nix</code>:
	<syntaxHighlight lang=nix>		<syntaxHighlight lang=nix>
	# NVIDIA drivers are unfree		# NVIDIA drivers are unfree
Line 59:		Line 57:

	[[Category:Applications]]		[[Category:Applications]]
			[[Category:Python]]

Nim65s: cache : cuda-maintainers -> nix-community

2025-05-16T13:12:11Z

cache : cuda-maintainers -> nix-community

← Older revision		Revision as of 13:12, 16 May 2025
Line 1:		Line 1:
	[https://github.com/google/jax JAX] is a framework for program transformation, esp. for automatic differentiation and machine learning. It's available on Nix/NixOS in the <code>python3Packages.{jax, jaxlib, jaxlibWithCuda}</code> packages.		[https://github.com/google/jax JAX] is a framework for program transformation, esp. for automatic differentiation and machine learning. It's available on Nix/NixOS in the <code>python3Packages.{jax, jaxlib, jaxlibWithCuda}</code> packages.

	'''Cache''': Using the [https://app.cachix.org/cache/~~cuda~~-~~maintainers#pull cuda~~-~~maintainers~~ cache] is recommended! It will save you valuable time and electrons. Getting set up should be as simple as <code>cachix use ~~cuda~~-~~maintainers~~</code>. See the [[CUDA]] wiki page for more info.		{{tip\|1='''Cache''': Using the [https://app.cachix.org/cache/nix-community nix-community cache] is recommended! It will save you valuable time and electrons. Getting set up should be as simple as <code>cachix use nix-community</code>. See the [[CUDA]] wiki page for more info.
			}}

	NOTE: JAX requires Python 3.9, the current version of <code>python3</code> in nixpkgs (as of 9/4/2021). JAX is currently only packaged for x86_64-linux (send a PR for your platform!).		NOTE: JAX requires Python 3.9, the current version of <code>python3</code> in nixpkgs (as of 9/4/2021). JAX is currently only packaged for x86_64-linux (send a PR for your platform!).

imported>Samuela at 19:27, 31 March 2022

2022-03-31T19:27:48Z

← Older revision		Revision as of 19:27, 31 March 2022
Line 1:		Line 1:
	[https://github.com/google/jax JAX] is a framework for program transformation, esp. for automatic differentiation and machine learning. It's available on Nix/NixOS in the <code>python3Packages.{jax, jaxlib, jaxlibWithCuda}</code> packages.		[https://github.com/google/jax JAX] is a framework for program transformation, esp. for automatic differentiation and machine learning. It's available on Nix/NixOS in the <code>python3Packages.{jax, jaxlib, jaxlibWithCuda}</code> packages.

			'''Cache''': Using the [https://app.cachix.org/cache/cuda-maintainers#pull cuda-maintainers cache] is recommended! It will save you valuable time and electrons. Getting set up should be as simple as <code>cachix use cuda-maintainers</code>. See the [[CUDA]] wiki page for more info.

	NOTE: JAX requires Python 3.9, the current version of <code>python3</code> in nixpkgs (as of 9/4/2021). JAX is currently only packaged for x86_64-linux (send a PR for your platform!).		NOTE: JAX requires Python 3.9, the current version of <code>python3</code> in nixpkgs (as of 9/4/2021). JAX is currently only packaged for x86_64-linux (send a PR for your platform!).

imported>Samuela at 21:21, 31 January 2022

2022-01-31T21:21:36Z

← Older revision		Revision as of 21:21, 31 January 2022
Line 44:		Line 44:
	It should print either "cpu", "gpu", or "tpu".		It should print either "cpu", "gpu", or "tpu".

	Note that hydra may not cache `jaxlibWithCuda` builds on cache.nixos.org since CUDA is "unfree." @samuela publishes builds on a public cachix [https://app.cachix.org/cache/ploop#pull ploop] cache. These are periodically built and pushed from [https://github.com/samuela/nixpkgs-upkeep/ nixpkgs-upkeep].		Note that hydra may not cache <code>jaxlibWithCuda</code> builds on cache.nixos.org since CUDA is "unfree." @samuela publishes builds on a public cachix [https://app.cachix.org/cache/ploop#pull ploop] cache. These are periodically built and pushed from [https://github.com/samuela/nixpkgs-upkeep/ nixpkgs-upkeep].

	== FAQ ==		== FAQ ==

imported>Samuela at 21:20, 31 January 2022

2022-01-31T21:20:22Z

← Older revision		Revision as of 21:20, 31 January 2022
Line 6:		Line 6:
	<syntaxHighlight lang=nix>		<syntaxHighlight lang=nix>
	let		let
	pkgs = import (fetchTarball("https://github.com/NixOS/nixpkgs/archive/~~eac6215607e4ccceb9631b01ee8f8903a6e82e02~~.tar.gz")) {};		# Last updated 01/31/2022. Check status.nixos.org for updates.
			pkgs = import (fetchTarball("https://github.com/NixOS/nixpkgs/archive/376934f4b7ca6910b243be5fabcf3f4228043725.tar.gz")) {};
	in pkgs.mkShell {		in pkgs.mkShell {
	buildInputs = with pkgs; [		buildInputs = with pkgs; [
Line 27:		Line 28:
	<syntaxHighlight lang=nix>		<syntaxHighlight lang=nix>
	let		let
	pkgs = import (fetchTarball("https://github.com/NixOS/nixpkgs/archive/~~eac6215607e4ccceb9631b01ee8f8903a6e82e02~~.tar.gz")) {};		# Last updated 01/31/2022. Check status.nixos.org for updates.
			pkgs = import (fetchTarball("https://github.com/NixOS/nixpkgs/archive/376934f4b7ca6910b243be5fabcf3f4228043725.tar.gz")) {};
	in pkgs.mkShell {		in pkgs.mkShell {
	buildInputs = with pkgs; [		buildInputs = with pkgs; [

imported>Samuela at 21:18, 31 January 2022

2022-01-31T21:18:50Z

← Older revision		Revision as of 21:18, 31 January 2022
Line 45:		Line 45:

	== FAQ ==		== FAQ ==
			=== How do I package JAX libraries? ===
			Never ever ever put <code>jaxlib</code> in <code>propagatedBuildInputs</code>. However, it may live happily in <code>buildInputs</code> or <code>checkInputs</code>. See https://github.com/NixOS/nixpkgs/pull/156808 for context.

	=== RuntimeError: Unknown: no kernel image is available for execution on the device ===		=== RuntimeError: Unknown: no kernel image is available for execution on the device ===
	This usually indicates that you have a driver version that is too old for the CUDA toolkit version the package is built with. The easiest fix is to set the environment variable <code>XLA_FLAGS="--xla_gpu_force_compilation_parallelism=1"</code>. Also consider upgrading your CUDA driver.		This usually indicates that you have a driver version that is too old for the CUDA toolkit version the package is built with. The easiest fix is to set the environment variable <code>XLA_FLAGS="--xla_gpu_force_compilation_parallelism=1"</code>. Also consider upgrading your CUDA driver.

imported>Samuela at 21:15, 31 January 2022

2022-01-31T21:15:00Z

← Older revision		Revision as of 21:15, 31 January 2022
Line 24:		Line 24:
	hardware.opengl.enable = true;		hardware.opengl.enable = true;
	</syntaxHighlight>		</syntaxHighlight>
	Then you can use the ~~jaxlib~~ package by setting the <code>cudaSupport</code> parameter:		Then you can use the <code>jaxlibWithCuda</code> package (equivalent to setting the <code>cudaSupport</code> parameter):
	<syntaxHighlight lang=nix>		<syntaxHighlight lang=nix>
	let		let
	~~# Fork with jax/jaxlib. See https://github.com/NixOS/nixpkgs/pull/134894.~~
	pkgs = import (fetchTarball("https://github.com/NixOS/nixpkgs/archive/eac6215607e4ccceb9631b01ee8f8903a6e82e02.tar.gz")) {};		pkgs = import (fetchTarball("https://github.com/NixOS/nixpkgs/archive/eac6215607e4ccceb9631b01ee8f8903a6e82e02.tar.gz")) {};
	in pkgs.mkShell {		in pkgs.mkShell {

imported>Samuela at 21:13, 31 January 2022

2022-01-31T21:13:21Z

← Older revision		Revision as of 21:13, 31 January 2022
Line 1:		Line 1:
	[https://github.com/google/jax JAX] is a framework for program transformation, esp. for automatic differentiation and machine learning. It's available on Nix/NixOS in the <code>python3Packages.{jax, jaxlib}</code> packages.		[https://github.com/google/jax JAX] is a framework for program transformation, esp. for automatic differentiation and machine learning. It's available on Nix/NixOS in the <code>python3Packages.{jax, jaxlib, jaxlibWithCuda}</code> packages.

	NOTE: JAX requires Python 3.9, the current version of <code>python3</code> in nixpkgs (as of 9/4/2021). JAX is currently only packaged for x86_64-linux (send a PR for your platform!).		NOTE: JAX requires Python 3.9, the current version of <code>python3</code> in nixpkgs (as of 9/4/2021). JAX is currently only packaged for x86_64-linux (send a PR for your platform!).
Line 33:		Line 33:
	python3		python3
	python3Packages.jax		python3Packages.jax
	(python3Packages.~~jaxlib.override { cudaSupport = true; })~~		python3Packages.jaxlibWithCuda
	];		];
	}		}
Line 42:		Line 42:

	It should print either "cpu", "gpu", or "tpu".		It should print either "cpu", "gpu", or "tpu".

			Note that hydra may not cache `jaxlibWithCuda` builds on cache.nixos.org since CUDA is "unfree." @samuela publishes builds on a public cachix [https://app.cachix.org/cache/ploop#pull ploop] cache. These are periodically built and pushed from [https://github.com/samuela/nixpkgs-upkeep/ nixpkgs-upkeep].

	== FAQ ==		== FAQ ==