Skip to content

Conversation

@copybara-service
Copy link

Require JAX>=0.7.0

Directly use jax.extend.backend for register_backend_cache.

This change updates pathwaysutils/lru_cache.py to import and use jax.extend.backend.register_backend_cache directly. The re-export of this function from pathwaysutils.jax is removed, along with version-specific compatibility code for older JAX versions.

----
Directly use jax.extend.backend for register_backend_cache.

This change updates pathwaysutils/lru_cache.py to import and use `jax.extend.backend.register_backend_cache` directly. The re-export of this function from `pathwaysutils.jax` is removed, along with version-specific compatibility code for older JAX versions.

PiperOrigin-RevId: 853322003
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant