JAX

JAX on Autograd ja XLA yhdistettynä suorituskykyistä koneoppimistutkimusta varten.

Uutiset

25.6.2026 JAX on nyt saatavilla Roihu-GPU:ssa, ja moduulin nimi on muutettu muotoon python-jax.

Saatavuus

Tällä hetkellä tuetut JAX-versiot:

Version	Module	Roihu-GPU	Puhti	Mahti	LUMI	Notes
0.10.2	`python-jax/0.10`	X	-	-	-	Default on Roihu-GPU
0.5.0	`jax/0.5.0`		default	default	-	all packages
	`jax/0.5.0-small`		X	X	-	framework only
0.4.38	`jax/0.4.38`		X	X	default*	all packages
	`jax/0.4.38-small`		X	X	X*	framework only
0.4.30	`jax/0.4.30`		X	X	X*	all packages
	`jax/0.4.30-small`		X	X	X*	framework only
0.4.23	`jax/0.4.23-py3.9`		X	X	X*
0.4.20	`jax/0.4.20`		X	X	X*
0.4.18	`jax/0.4.18`		-	-	X*
0.4.14	`jax/0.4.14`		X	X	-
0.4.13	`jax/0.4.13`		X	X	-
0.4.1	`jax/0.4.1`		X	X	-
0.3.13	`jax/0.3.13`		X	X	-

Moduulit sisältävät Pythonille JAXin, jossa on GPU-tuki CUDA:n/ROCm:n kautta, sekä suuren joukon muita Python-paketteja, joita käytetään yleisesti yhdessä JAXin kanssa. Python-versio vaihtelee moduulijulkaisujen välillä. Moduulit, joissa JAX-versio on enintään 0.4.23, käyttävät Python 3.9:ää, myöhemmät moduulit käyttävät Python 3.12:ta.

LUMIn versiot, jotka on merkitty tähdellä "*", ovat edelleen kokeellisia ja niiden tuki on rajallinen. Osa alla kuvatuista ominaisuuksista ei välttämättä toimi niissä. Huomaa, että JAX on saatavilla myös LUMI Software Libraryssa, jota ylläpitää LUMI User Support Team (ei CSC:n tutkimuksen asiakastuki).

Versiosta 0.4.30 alkaen JAX-moduulista on kaksi versiota:

"small"-versio, joka sisältää JAXin CUDA 12.2 GPU -tuella sekä PyTorchin ja Tensorflown vain CPU:lle tarkoitetut versiot, jotta niiden datanlataustyökaluja voidaan käyttää. Näiden nimeämiskäytäntö on jax/<version>-small, tai voit käyttää jax/small ladataksesi oletusversion (uusimman).
"full"-versio, joka sisältää monia yleisesti käytettyjä JAXiin perustuvia koneoppimispaketteja -- voit tarkistaa pip list -komennolla täydet tiedot kaikista mukana olevista paketeista. Näiden nimeämiskäytäntö on jax/<version>, tai voit yksinkertaisesti käyttää jax ladataksesi oletusversion (uusimman).

Roihu-GPU:ssa on vain yksi moduuli, jossa on täysi pakettivalikoima. Roihu-GPU:n moduulin nimi on muutettu muotoon python-jax.

Kaikki moduulit perustuvat Apptaineria käyttäviin kontteihin (tunnettiin aiemmin nimellä Singularity). tykky:n kautta on tarjottu wrapper-skriptit, jotta yleiset komennot kuten python, python3, pip ja pip3 sekä asennettujen pakettien tarjoamat komennot toimisivat normaalisti. Muiden komentojen kohdalla ne voi olla tarpeen suorittaa etuliitteellä apptainer_wrapper exec. Lisätietoja on CSC:n yleisissä ohjeissa Apptainer-konttien ajamisesta.

Lisäpaketit

Jos huomaat, että jokin paketti puuttuu, voit usein asentaa sen itse komennolla pip install. On suositeltavaa käyttää Pythonin virtuaaliympäristöjä. Katso Python-dokumentaatiostamme lisätietoja pakettien asentamisesta itse olemassa oleviin moduuleihin. Jos mielestäsi jokin tärkeä paketti pitäisi sisällyttää CSC:n tarjoamaan moduuliin, ota yhteyttä asiakastukeemme.

Jos muutaman paketin lisääminen pip:llä ei riitä, voit myös rakentaa omia konttejasi esimerkiksi laajentamalla CSC:n koneoppimiskontteja.

Lisenssi

JAX on lisensoitu Apache License 2.0 -lisenssillä.

Käyttö

Käyttääksesi JAXin oletusversiota Roihu-GPU:ssa, alusta se komennolla:

module load python-jax

Saadaksesi JAXin käyttöön Puhdissa tai Mahdissa:

module load jax

tai

module load jax/small

Saadaksesi CSC:n asentaman JAXin käyttöön LUMIssa:

module use /appl/local/csc/modulefiles/
module load jax  # jax/small

Note

Suosittelemme kiinnittämään version skripteissäsi, jotta tulevien päivitysten muutokset eivät riko skriptejä, esimerkiksi: module load jax/0.5.0 tai module load jax/0.5.0-small

Huomaa, että JAX-moduulit sisältävät jo vastaavat CUDA- ja cuDNN- tai ROCm-kirjastot, joten mitään cuda-, cudnn- tai rocm-moduuleja ei tarvitse ladata erikseen!

Tämä näyttää kaikki saatavilla olevat JAX-versiot:

module avail python-jax  # on Roihu-GPU
module avail jax         # on other systems

Note

Huomaa, että kirjautumissolmut eivät ole tarkoitettu raskasta laskentaa varten, käytä sen sijaan slurm-eräajoja. Katso ohjeemme eräajojärjestelmän käytöstä.

Note

Älä lue valtavaa määrää tiedostoja jaetusta tiedostojärjestelmästä, vaan käytä nopeaa paikallista levyä tai paketoi datasi suurempiin tiedostoihin! Katso lisätietoja koneoppimisen oppaamme datan tallennusta käsittelevästä osiosta.