Hacktoberfest 2026: as issues que os mantenedores marcaram para outubro, abertas e boas para iniciantes. Ver issues do Hacktoberfest

Two small JAX-related cleanups

Aberta
  • Algo errado?
#1,011 1 comentário 0 reações 0 responsáveis Ver no GitHub

Mantenedores costumam responder em até 1 dia

Ninguém assumiu esta issue ainda.

Avaliação

Dificuldade
4/5
Tempo estimado
1-2 dias
Facilidade para iniciantes
52/100
Tipo de issue
Refatoração
Clareza
Razoavelmente clara
Status de atividade
Ativa
Stack de tecnologia
python
Domínio
backend

Direção de pesquisa

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.

Escrita pelo modelo de indexação a partir do texto da issue.

Descrição

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.

Linguagem predominante
Python
Estrelas
33
Forks
33
Merge médio
1h 16min
PRs com merge (30d)
32

Preparar o ambiente

  • Sem Dockerfile nem arquivo Docker Compose
  • Tem um modelo de pull request
  • Sem guia de contribuição

Primeiros passos

  1. Leia a issue inteira e depois o guia de contribuição do projeto.
  2. Comente na issue dizendo que vai assumir — evita que duas pessoas façam o mesmo trabalho.
  3. Faça um fork do repositório e trabalhe em uma branch.
  4. Abra um pull request que referencie o número da issue.

Mais de data-apis/array-api-extra

Todas as issues de data-apis/array-api-extra

Issues semelhantes

Mais issues de Python

Receba novas issues na sua caixa de entrada

Um resumo curto de issues do GitHub para quem está começando.