JAX: Difference between revisions

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.

}}

~~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 ==

Line 21:

Line 19:

== 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>:

# NVIDIA drivers are unfree

Line 59:

Line 57:

[[Category:Applications]]

[[Category:Python]]

@@ 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.
 }}
-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 ==
@@ Line 21: / Line 19: @@
 == 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>
 # NVIDIA drivers are unfree
@@ Line 59: / Line 57: @@
 [[Category:Applications]]
+[[Category:Python]]