Fix failing device_put rules tests

PiperOrigin-RevId: 525179481

tests/rules_test.py

@@ -482,6 +482,7 @@ def _compare_jacobians(self, j_fwd, j_rev, j_rule, primitive):
     else:
       if primitive == lax.reshape_p:
         # Reshape Jacobian is special-case defined as identity.
+        j_rule: np.ndarray
         j_rule = j_rule.reshape(j_fwd.shape)
 
       self.assertAllClose(j_fwd, j_rev)
@@ -581,10 +582,12 @@ def _test_primitive(
   def test_unary(self, primitive: Optional[Primitive], shape, dtype, params):
     if primitive == lax.device_put_p:
       # Can't instantiate devices at test generation time; using subtests.
-      for device in [None] + jax.devices() + jax.devices('cpu'):
-        with self.subTest(device=device):
-          params = {'device': device}
-          self._test_primitive(primitive, [shape], dtype, params)
+      devices = [None] + jax.devices() + jax.devices('cpu')
+      for device in devices:
+        for src in devices:
+          with self.subTest(device=device, src=src):
+            params = {'device': device, 'src': src}
+            self._test_primitive(primitive, [shape], dtype, params)
 
     else:
       self._test_primitive(primitive, [shape], dtype, params)