google / google/ffn

[Compatibility Issue] Missing numpy>=2 constraint — np.typing crashes training entry points on numpy<2

Open Beginner friendly
#127 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
351
Forks
109
Avg merge
1d 8h
Merged PRs (30d)
7

Description

## Affected File

`ffn/training/examples.py:130`

## Current Code

```python
# ffn/training/examples.py
# (no from __future__ import annotations)

def update_seeds(self, batched_seeds: np.typing.ArrayLike):
"""Distributes updated predictions back to the generator buffers.
...
"""
assert self._seeds is not None
batched_seeds = np.asarray(batched_seeds)
```

No `from __future__ import annotations` — the annotation is evaluated when the `BatchExampleIter` class body executes at module import time.

## Root Cause

`numpy.typing` is only accessible as a module attribute (`np.typing.XXX`) in **numpy >= 2.0.0** (via `__getattr__`). On numpy < 2.0, access without a prior `import numpy.typing` triggers:

```
AttributeError: module 'numpy' has no attribute 'typing'. Did you mean: '_typing'?
```

## Dependency

`setup.py` declares `'numpy>=1.11.1'` (2020-era floor, no upper bound). No lock files exist. Fresh pip installs resolve to numpy 2.x (fine), but any environment with a pre-existing numpy<2 pin (e.g. from an older tensorflow install) keeps numpy 1.x and crashes.

## Upstream Dependency Protection — Could Be Avoided, But Not Guaranteed

If any code in the import chain executes `import numpy.typing` before `examples.py` loads, the annotation evaluates fine even on numpy<2 (the import injects `typing` into `numpy.__dict__`). So the crash may be avoided depending on what the user imports first. However, this protection is incidental — it depends on the internal implementation of dependencies, which can change — so it should not be relied upon.

What we have checked so far:

- **pandas**: declared in `setup.py`, but only imported in `ffn/utils/decision_point.py`, which is loaded only by a test. On the runtime training path, pandas is never imported before `examples.py`. (Note: only pandas >= 3.0 triggers `numpy.typing` at import time; pandas 2.x does not.)
- **xarray / zarr**: not dependencies.
- **scipy**: does not trigger `numpy.typing` at `import scipy` time (verified).
- **tensorflow, connectomics (git dep), absl, `ffn.training.augmentation` and the rest of the training import chain**: not yet checked one by one — any of them could potentially do `import numpy.typing` and avoid the crash.

## Impact

`ffn/training/examples.py` is loaded eagerly by both training entry points:
- `python train.py` — `from ffn.training import examples` at line 33
- JAX path — `ffn.jax.main` → `ffn.jax.train` → `ffn.jax.input_pipeline`

`import ffn` itself is safe (zero imports in `ffn/__init__.py`), and the inference path does not load `examples.py`.

On numpy<2 with no prior `import numpy.typing` in the process, the module raises:

```
AttributeError: module 'numpy' has no attribute 'typing'. Did you mean: '_typing'?
```

(the failure mode itself was verified on numpy 1.26.4; the full `python train.py` reproduction with all dependencies installed has not been run).

## Solution

Any of the following in `examples.py`:

1. **`from __future__ import annotations`** — defers annotation evaluation (simplest, also a free import-time win)
2. **`import numpy.typing`** before the class definition
3. **Use `np.ndarray` or a plain type** in the annotation — the current `ArrayLike` annotation is also type-incorrect: the docstring says it accepts "array-like object backed by accelerator memory" (i.e. JAX arrays), which `np.typing.ArrayLike` does not describe.

Optionally bump the declared floor in `setup.py` from `numpy>=1.11.1` to `numpy>=2` (or at least `>=1.20`).

## References

- NumPy 2.0 release notes — `numpy.typing` exposed as module attribute via `__getattr__`
- Similar fix: [numexpr#540](https://github.com/pydata/numexpr/issues/540)

Contributor guide

Open the contributing guide

Research direction

Start with ffn/training/examples.py:130 and inspect the annotation evaluation during module import, then check setup.py for the declared NumPy floor. Reproduce the import on NumPy 1.26.4 and verify that the training entry points load successfully without the reported AttributeError; account for any dependency constraint change in the final checks.

Written by the indexing model from the issue text.

Assessment

Tech stack
numpy, python
Domain
machine-learning
Issue type
Bug
Difficulty
2/5
Estimated time
1-3 hours
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
78/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.