refactor: convert jax verification script to pytest and update docs

tests/test_jax_init.py

@@ -1,22 +1,7 @@
 import jax
-import sys
 
-def verify_jax():
-    print(f"Python version: {sys.version}")
-    print(f"JAX version: {jax.__version__}")
-    
+
+def test_jax_initializes():
+    """Verify that JAX initializes and exposes at least one device."""
     devices = jax.devices()
-    print(f"Available devices: {devices}")
-    
-    gpu_found = any(d.device_kind == 'gpu' for d in devices)
-    if gpu_found:
-        print("SUCCESS: GPU detected!")
-    else:
-        print("INFO: Only CPU detected (expected if not in GPU-enabled environment/container).")
-
-if __name__ == "__main__":
-    try:
-        verify_jax()
-    except Exception as e:
-        print(f"ERROR during JAX initialization: {e}")
-        sys.exit(1)
+    assert len(devices) > 0, "JAX should expose at least one device"