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

Update pjit_test to skip GDA tests with Array is enabled.

Browse Source 
PiperOrigin-RevId: 475684445
This commit is contained in:
Yash Katariya 2022-09-20 16:38:09 -07:00 committed by jax authors
parent 310bcd57a2
commit 6183727acc
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
tests/pjit_test.py
View File

@ -1090,9 +1090,13 @@ class PJitTest(jtu.BufferDonationTestCase):
pjit_f(jnp.array([1, 2, 3]))
@jax_array(False)
class GDAPjitTest(jtu.JaxTestCase):
def setUp(self):
super().setUp()
if config.jax_array:
self.skipTest('GDA and Array cannot be enabled together.')
def test_pjit_gda_single_output(self):
global_mesh = jtu.create_global_mesh((4, 2), ('x', 'y'))
global_input_shape = (8, 2)