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

Deprecate several APIs in jax.lib.xla_bridge

Browse Source 
PiperOrigin-RevId: 658274719
This commit is contained in:
Jake VanderPlas 2024-07-31 22:59:49 -07:00 committed by jax authors
parent 6c083d78e6
commit 3551fcc077
2 changed files with 34 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
CHANGELOG.md
View File

@ -19,6 +19,10 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
* Deprecations
* Complex inputs to {func}`jax.numpy.clip` and {func}`jax.numpy.hypot` are
no longer allowed, after being deprecated since JAX v0.4.27.
* Deprecated the following APIs:
* `jax.lib.xla_bridge.xla_client`: use {mod}`jax.lib.xla_client` directly.
* `jax.lib.xla_bridge.get_backend`: use {func}`jax.extend.backend.get_backend`.
* `jax.lib.xla_bridge.default_backend`: use {func}`jax.extend.backend.default_backend`.
## jaxlib 0.4.32

33
jax/lib/xla_bridge.py
View File

@ -14,12 +14,39 @@
# ruff: noqa: F401
from jax._src.xla_bridge import (
default_backend as default_backend,
get_backend as get_backend,
xla_client as xla_client,
default_backend as _deprecated_default_backend,
get_backend as _deprecated_get_backend,
xla_client as _deprecated_xla_client,
_backends as _backends,
)
from jax._src.compiler import (
get_compile_options as get_compile_options,
)
_deprecations = {
# Added July 31, 2024
"xla_client": (
"jax.lib.xla_bridge.xla_client is deprecated; use jax.lib.xla_client directly.",
_deprecated_xla_client
),
"get_backend": (
"jax.lib.xla_bridge.get_backend is deprecated; use jax.extend.backend.get_backend.",
_deprecated_get_backend
),
"default_backend": (
"jax.lib.xla_bridge.default_backend is deprecated; use jax.extend.backend.default_backend.",
_deprecated_default_backend
),
}
import typing as _typing
if _typing.TYPE_CHECKING:
from jax._src.xla_bridge import default_backend as default_backend
from jax._src.xla_bridge import get_backend as get_backend
from jax._src.xla_bridge import xla_client as xla_client
else:
from jax._src.deprecations import deprecation_getattr as _deprecation_getattr
__getattr__ = _deprecation_getattr(__name__, _deprecations)
del _deprecation_getattr
del _typing