mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-26 22:36:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Increase the number of iterations in a test that compares rolled versus unrolled HLO for length.

Browse Source 
A change that avoids duplicating subcomputations in XLA causes this test to fail, but we can make it work again by increasing the number of iterations.

PiperOrigin-RevId: 735875835
This commit is contained in:
Peter Hawkins 2025-03-11 13:44:37 -07:00 committed by jax authors
parent e0545a71eb
commit 67aa997f84
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tests/lax_control_flow_test.py
View File

@ -2445,7 +2445,7 @@ class LaxControlFlowTest(jtu.JaxTestCase):
assert b.shape == () assert b.shape == ()
return c, b return c, b
xs = jnp.ones((5, 3)) xs = jnp.ones((20, 3))
c = jnp.ones(4) c = jnp.ones(4)
scan = lambda c, xs: lax.scan(f, c, xs) scan = lambda c, xs: lax.scan(f, c, xs)