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	LUMI	Notes
0.10.2	`python-jax/0.10`	X	-	Default on Roihu-GPU
0.4.38	`jax/0.4.38`		default*	all packages
	`jax/0.4.38-small`		X*	framework only
0.4.30	`jax/0.4.30`		X*	all packages
	`jax/0.4.30-small`		X*	framework only
0.4.23	`jax/0.4.23-py3.9`		X*
0.4.20	`jax/0.4.20`		X*
0.4.18	`jax/0.4.18`		X*

Moduulit sisältävät JAXin Pythonille GPU-tuella 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:ää, myöhemmät moduulit käyttävät Python 3.12:ta.

LUMIssa 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 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ämä noudattavat nimeämiskäytäntöä 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ämä noudattavat nimeämiskäytäntöä 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 (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 ajaa 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. 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ö

Jos haluat käyttää JAXin oletusversiota Roihu-GPU:ssa, alusta se komennolla:

module load python-jax

Jos haluat käyttää CSC:n asentamaa JAXia LUMIssa:

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

Note

Suosittelemme määrittämään skripteissäsi version kiinteästi, jotta tulevien päivitysten muutokset eivät riko skriptejä, esimerkiksi: module load jax/0.10 (Roihu-GPU) tai module load jax/0.4.30-small (LUMI)

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 LUMI

Note

Huomaa, että kirjautumissolmuja ei ole tarkoitettu raskaaseen laskentaan, 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 suuremmiksi tiedostoiksi! Katso lisätietoja koneoppimisen oppaamme Datan tallennus -osiosta.