Merge pull request #16 from HMUNACHI/dev

patches in random
HMUNACHI · Mar 12, 2024 · 8508e7a · 8508e7a
2 parents 43ed416 + a1bc92d
commit 8508e7a
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/nanodl/__src/utils/random.py b/nanodl/__src/utils/random.py
@@ -11,7 +11,7 @@ def time_rng_key(seed=None) -> jnp.ndarray:
         jnp.ndarray: A JAX random key.
     """
     key = int(time.time()) if seed is None else seed
-    return random.PRNGKey(seed)
+    return random.PRNGKey(key)
 
 def uniform(shape: Tuple[int, ...], 
             dtype: Any = jnp.float32, 

diff --git a/tests/test_random.py b/tests/test_random.py
@@ -15,7 +15,7 @@ def test_time_rng_key(self):
         self.assertTrue(jnp.array_equal(key1, key2), "Keys should be equal for the same seed")
 
     def test_uniform(self):
-        result = uniform((2, 3), seed=42)
+        result = uniform((2, 3))
         self.assertEqual(result.shape, (2, 3))
         self.assertEqual(result.dtype, jnp.float32)