google-deepmind / google-deepmind/PGMax

Bug Report: Incompatibility with JAX 0.7.0+ due to deprecated API usage

Open
#10 1 comment 1 reaction 0 assignees View on GitHub
Dominant language
Jupyter Notebook
Stars
166
Forks
12
PR merge metrics
No merged PRs in 30d

Description

### **Description**
PGMax is currently incompatible with JAX version 0.7.0 and later due to the use of a deprecated API. The deprecation warning was introduced in JAX 0.7.0, with the API slated for removal in JAX 0.8.0.

### **Error**
```python
AttributeError: jax.lib.xla_bridge.get_backend is deprecated and will be removed in JAX v0.8.0; use jax.extend.backend.get_backend.
```

Full traceback shows the error originates from:
```python
File "/usr/local/lib/python3.12/site-packages/pgmax/infer/inferer.py", line 66, in __post_init__
if jax.lib.xla_bridge.get_backend().platform == "tpu": # pragma: no cover
```

### **Location**
**File:** `pgmax/infer/inferer.py`
**Line:** 66
**Class:** `InfererContext.__post_init__()`

### **Current Code**
```python
def __post_init__(self):
if jax.lib.xla_bridge.get_backend().platform == "tpu": # pragma: no cover
warnings.warn(
"PGMax is not optimized for the TPU backend. Please consider using"
" GPUs!"
)
```

### **Proposed Fix**
Replace the deprecated API with the new `jax.extend.backend.get_backend()`:

```python
def __post_init__(self):
import jax.extend # Explicit import required
if jax.extend.backend.get_backend().platform == "tpu": # pragma: no cover
warnings.warn(
"PGMax is not optimized for the TPU backend. Please consider using"
" GPUs!"
)
```

### **JAX Changelog Reference**
From the [JAX v0.7.0 changelog](https://docs.jax.dev/en/latest/changelog.html):
> `jax.lib.xla_bridge.get_backend` is deprecated in JAX v0.7.0 and will be removed in JAX v0.8.0; use `jax.extend.backend.get_backend`

### **Impact**
- Users cannot use PGMax with JAX 0.7.0 or later without encountering this error
- This blocks users from receiving JAX security updates and new features
- Forces dependency pinning to JAX <0.7.0 in downstream projects

### **Environment**
- **PGMax version:** 0.6.1 (from main branch)
- **JAX version:** 0.7.0+
- **Python version:** 3.12

### **Workaround**
Currently, users must pin JAX to versions below 0.7.0:
```
jax<0.7.0
jaxlib<0.7.0
```

---

**Note:** The fix requires adding `import jax.extend` as per JAX's deprecation notice: *"please note that you must `import jax.extend` explicitly."*

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.