Merge pull request jax-md#304 from muhrin/fix-key-array

Fixes jax-md#298
abhijeetgangan · Apr 17, 2024 · c451353 · c451353
2 parents e703937 + 38a1c8b
commit c451353
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/jax_md/rigid_body.py b/jax_md/rigid_body.py
@@ -73,7 +73,7 @@
 PyTree = Any
 f64 = util.f64
 f32 = util.f32
-KeyArray = random.KeyArray
+KeyArray = jax.Array
 NeighborListFns = partition.NeighborListFns
 ShiftFn = space.ShiftFn