Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

update Jax-Metal CI runner with nightly jaxlib config #21070

Merged
merged 4 commits into from
May 6, 2024

Conversation

shuhand0
Copy link
Collaborator

@shuhand0 shuhand0 commented May 4, 2024

Add nightly jaxlib config to Jax-Metal CI runner.
Sync lax_metal_test with lax_numpy_test

@shuhand0 shuhand0 requested review from skye and jakevdp May 4, 2024 00:58
.github/workflows/metal_plugin_ci.yml Outdated Show resolved Hide resolved
@@ -302,13 +302,11 @@ def testCountNonzero(self, shape, dtype, axis):
self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False)
self._CompileAndCheck(jnp_fun, args_maker)

@jtu.sample_product(shape=all_shapes, dtype=all_dtypes)
@jtu.sample_product(shape=nonzerodim_shapes, dtype=all_dtypes)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

OOC why did you need to reduce coverage? Were new failures introduced in jaxlib nightly?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

metal_plugin_ci has been failing on tip since two weeks ago. I suspect there is a change in JAX? So the coverage changes here are synced from lax_numpy_test.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah. More reason to sync up the test files!

@@ -302,13 +302,11 @@ def testCountNonzero(self, shape, dtype, axis):
self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False)
self._CompileAndCheck(jnp_fun, args_maker)

@jtu.sample_product(shape=all_shapes, dtype=all_dtypes)
@jtu.sample_product(shape=nonzerodim_shapes, dtype=all_dtypes)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah. More reason to sync up the test files!

@google-ml-butler google-ml-butler bot added kokoro:force-run pull ready Ready for copybara import and testing labels May 6, 2024
@copybara-service copybara-service bot merged commit eee2783 into google:main May 6, 2024
14 of 15 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
pull ready Ready for copybara import and testing
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

3 participants