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

Merge pull request #22840 from jakevdp:fix-backends

Browse Source 
PiperOrigin-RevId: 658817878
This commit is contained in:
jax authors 2024-08-02 09:47:25 -07:00
commit 6c79c10446
3 changed files with 1 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
jax/extend/backend.py
View File

@ -21,7 +21,6 @@ from jax._src.api import (
from jax._src.xla_bridge import (
backends as backends,
backend_xla_version as backend_xla_version,
default_backend as default_backend,
get_backend as get_backend,
register_backend_factory as register_backend_factory,
)

2
jax/lib/xla_bridge.py
View File

@ -35,7 +35,7 @@ _deprecations = {
_deprecated_get_backend
),
"default_backend": (
"jax.lib.xla_bridge.default_backend is deprecated; use jax.extend.backend.default_backend.",
"jax.lib.xla_bridge.default_backend is deprecated; use jax.default_backend.",
_deprecated_default_backend
),
}

1
tests/extend_test.py
View File

@ -50,7 +50,6 @@ class ExtendTest(jtu.JaxTestCase):
# Assume these are tested elsewhere, only check equivalence
self.assertIs(jex.backend.backends, xla_bridge.backends)
self.assertIs(jex.backend.backend_xla_version, xla_bridge.backend_xla_version)
self.assertIs(jex.backend.default_backend, xla_bridge.default_backend)
self.assertIs(jex.backend.clear_backends, api.clear_backends)
self.assertIs(jex.backend.get_backend, xla_bridge.get_backend)
self.assertIs(jex.backend.register_backend_factory, xla_bridge.register_backend_factory)