[Hotfix] Restore CI (#1175)

sotetsuk · Mar 13, 2024 · 6c582b6 · 6c582b6
1 parent 262afea
commit 6c582b6
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 33 deletions.
diff --git a/tests/test_baseline.py b/tests/test_baseline.py
diff --git a/tests/test_bridge_bidding.py b/tests/test_bridge_bidding.py
@@ -25,7 +25,7 @@
 )
 
 
-def init(rng: jax.random.KeyArray) -> State:
+def init(rng: jax.Array) -> State:
     rng1, rng2, rng3, rng4, rng5, rng6 = jax.random.split(rng, num=6)
     hand = jnp.arange(0, 52)
     hand = jax.random.permutation(rng2, hand)