test: revert getkey removal

Signed-off-by: Nathaniel Starkman <[email protected]>
patrick-kidger · Feb 1, 2025 · 7d65e9b · 7d65e9b
1 parent a576b8d
commit 7d65e9b
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 5 deletions.
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -1,7 +1,7 @@
-import jax.random as jr
+import equinox.internal as eqxi
 import pytest
 
 
 @pytest.fixture()
 def getkey():
-    return lambda: jr.PRNGKey(0)
+    return eqxi.GetKey()
diff --git a/tests/test_lora.py b/tests/test_lora.py
@@ -95,12 +95,16 @@ def run3(mlp, vector):
     run3(mlp, vector)
 
 
-def test_materialise(getkey):
+def test_materialise():
+    key = jr.key(0)
+
+    key, *subkeys = jr.split(key, 3)
     x_false = lora.LoraArray(
-        jr.normal(getkey(), (3, 3)), rank=2, allow_materialise=False, key=getkey()
+        jr.normal(subkeys[0], (3, 3)), rank=2, allow_materialise=False, key=subkeys[1]
     )
+    key, *subkeys = jr.split(key, 3)
     x_true = lora.LoraArray(
-        jr.normal(getkey(), (3, 3)), rank=2, allow_materialise=True, key=getkey()
+        jr.normal(subkeys[0], (3, 3)), rank=2, allow_materialise=True, key=subkeys[1]
     )
 
     _ = quax.quaxify(jax.nn.relu)(x_true)