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.

Sign up for GitHub

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

Jump to bottom

Fix main and add nnx.fori_loop test #4472

Merged

copybara-service merged 1 commit into main from nnx-fix-fori

Jan 6, 2025

Collaborator

cgarciae commented Jan 6, 2025 •

edited

Loading

What does this PR do?

Creates a custom copy of chex.set_n_cpu_devices that doesn't call jax.lib.xla_bridge._backends (which doesn't exist anymore) to fix jax_utils_test.py.
Adds a simple test to check that the sample specified in Why doesn't nnx.fori_loop work here? #4436 works now.

cgarciae mentioned this pull request

Why doesn't nnx.fori_loop work here? #4436

Open


          [nnx] test fori_loop output

f8164dd

cgarciae force-pushed the nnx-fix-fori branch from d8accc9 to f8164dd Compare

January 6, 2025 21:26

cgarciae changed the title ~~[nnx] test fori_loop output~~ Fix main and add test

cgarciae changed the title ~~Fix main and add test~~ Fix main and add nnx.fori_loop test

levskaya approved these changes

View reviewed changes

levskaya added the pull ready label

copybara-service bot merged commit 595e711 into main

19 checks passed

copybara-service bot deleted the nnx-fix-fori branch

January 6, 2025 22:50

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels