JAX

JAX on Autograd ja XLA yhdistettynä suorituskykyiseen koneoppimistutkimukseen.

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 tarkoitetun JAXin, jossa on GPU-tuki CUDA:n/ROCm:n kautta, sekä laajan 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:ää, ja sitä uudemmat moduulit Python 3.12:ta.

LUMIssa "*" merkityt versiot 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 tutkimustuki).

Versiosta 0.4.30 alkaen JAX-moduulista on kaksi versiota:

"small"-versio 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 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, joka sisältää täyden pakettivalikoiman. Roihu-GPU:n moduulin nimi on muutettu muotoon python-jax.

Kaikki moduulit perustuvat kontteihin, jotka käyttävät Apptaineria (aiemmin tunnettu 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 suorittamisesta.

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. Jos mielestäsi jokin tärkeä paketti pitäisi sisällyttää CSC:n tarjoamaan moduuliin, ota yhteyttä asiakastukeemme.

Lisenssi

JAX on lisensoitu Apache License 2.0:lla.

Käyttö

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

module load python-jax

Käyttääksesi JAXia Puhdissa tai Mahdissa:

module load jax

tai

module load jax/small

Käyttääksesi CSC:n asentamaa JAXia LUMIssa:

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

Note

Skripteissä suosittelemme kiinnittämään version, 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  # Roihu-GPU:ssa
module avail jax         # muissa järjestelmissä

Note

Huomaa, että kirjautumissolmut eivät ole tarkoitettu raskasta laskentaa varten, vaan 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 data suurempiin tiedostoihin! Katso lisätietoja koneoppimisen oppaamme datan tallennus -osiosta.