From 3f60bc21a7ddaf13aa1125025779c9d0248ac0c4 Mon Sep 17 00:00:00 2001
From: Zachary Charles <zachcharles@google.com>
Date: Tue, 4 Mar 2025 12:05:16 -0800
Subject: [PATCH] Remove RNG usage in impls_test.

PiperOrigin-RevId: 733416782
---
 drjax/_src/impls_test.py | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/drjax/_src/impls_test.py b/drjax/_src/impls_test.py
index cad507d..ff91f6a 100644
--- a/drjax/_src/impls_test.py
+++ b/drjax/_src/impls_test.py
@@ -48,13 +48,10 @@ def temp_sens_example(m, t):
       )
       return comp_factory.mean_from_placement(total_over)
 
-    key = jax.random.PRNGKey(2)
-    random_measurements = jax.random.uniform(
-        key, shape=[self._placements['clients']]
-    )
+    measurements = jnp.arange(self._placements['clients'])
 
     self.assertEqual(
-        temp_sens_example(random_measurements, jnp.array(0.5)), 0.53
+        temp_sens_example(measurements, jnp.median(measurements)), 0.5
     )
 
   def test_runs_fake_training(self):