Hacktoberfest 2026: le issue che i maintainer hanno segnato per ottobre, aperte e adatte ai principianti. Sfoglia le issue Hacktoberfest

Two small JAX-related cleanups

Aperta
#1,011 1 commento 0 reazioni 0 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 1 giorno

Nessuno ha ancora preso questa issue.

Valutazione

Difficoltà
4/5
Tempo stimato
1-2 giorni
Idoneità per principianti
52/100
Tipo di issue
Refactoring
Chiarezza
Abbastanza chiara
Stato di attività
Attiva
Stack tecnologico
python
Ambito
backend

Direzione di ricerca

Read src/array_api_extra/_lib/_helpers.py around the JAX capabilities() workaround and src/array_api_extra/_agnostic/_set.py around nunique. The issue leaves removal of the workaround contingent on deciding whether to support JAX versions below 0.6.0, so first check project policy and tests. Done means the comments accurately describe JAX's unique_counts behavior, and any shim or minimum-version documentation reflects the project's support decision.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

1. Drop the capabilities() shim for jax < 0.6.0?

https://github.com/data-apis/array-api-extra/blob/04407df9d1f8a06f72faf1d9951c82f81f237d06/src/array_api_extra/_lib/_helpers.py#L315-L320

_helpers.capabilities forces "boolean indexing" to False for JAX. This works around jax-ml/jax#27418, which was fixed in JAX 0.6.0 (March 2025). CI pins jax >= 0.10.2, so this branch never runs there (it is marked # pragma: no cover).

The package doesn't declare a minimum JAX version, so whether to remove it is a policy question: do we still want to support JAX < 0.6.0 at runtime? If not, the branch can be deleted:

out = xp.__array_namespace_info__().capabilities()
if _compat.is_torch_namespace(xp):
    ...

Removing it would also be a good moment to document a minimum supported JAX version.

2. Inaccurate comment in nunique

https://github.com/data-apis/array-api-extra/blob/04407df9d1f8a06f72faf1d9951c82f81f237d06/src/array_api_extra/_agnostic/_set.py#L31-L35

# 3. backend does not have unique_counts; e.g. wrapped JAX

JAX does have unique_counts. It takes the O(n log n) sort-based path because capabilities()["data-dependent shapes"] is False (the output shape of unique_counts can't be known under jax.jit). Also, JAX is no longer "wrapped": jax.numpy is used directly. Suggested wording:

# 3. backend lacks data-dependent shapes, so unique_counts is unusable
#    (e.g. JAX, whose output shape can't be known under jax.jit)

The comment just below it, # xp does not have unique_counts; O(n*logn) complexity, should be updated the same way.

Lingua principale
Python
Stelle
33
Fork
33
Merge medio
1h 16m
PR unite (30g)
32

Preparare l'ambiente

  • Nessun Dockerfile né file Docker Compose
  • Ha un modello di pull request
  • Nessuna guida per i contributori

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di data-apis/array-api-extra

Tutte le issue di data-apis/array-api-extra

Issue simili

Altre issue su Python

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.