Two small JAX-related cleanups
维护者通常 1 天内回复
还没有人认领这个 Issue。
评估
调研方向
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.
由索引模型根据 Issue 内容生成。
描述
1. Drop the capabilities() shim for jax < 0.6.0?
_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
# 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.
- 主要语言
- Python
- 星标
- 33
- 派生
- 33
- 平均合并
- 1 小时 16 分钟
- 30 天内合并 PR
- 32
环境准备
- 没有 Dockerfile 或 Docker Compose 文件
- 有 Pull Request 模板
- 没有贡献指南
从这里开始
- 先读完整个 Issue，再读项目的贡献指南。
- 在 Issue 下留言说明你要接手 —— 这能避免两个人做同样的事。
- Fork 仓库，在一个分支上完成修改。
- 提交 Pull Request，并在描述里引用这个 Issue 编号。
data-apis/array-api-extra 的其他 Issue
-
难度 2/5 1-3 小时 新手友好度 76/100
data-apis/array-api-extra#1005 · 15 条评论 ·
维护者通常 1 天内回复
-
enhancement new function
难度 2/5 1-3 小时 新手友好度 72/100
data-apis/array-api-extra#992 · 1 条评论 ·
维护者通常 1 天内回复
-
难度 3/5 1-2 天 新手友好度 56/100
data-apis/array-api-extra#1014 · 1 条评论 ·
维护者通常 1 天内回复
-
ENH, TST: Add support for MLX as array backend可能已有人在做 @prady0t 于 4 天前认领。 未关闭
难度 4/5 3-5 天 新手友好度 35/100
data-apis/array-api-extra#1006 ·
维护者通常 1 天内回复
-
难度 5/5 一周以上 新手友好度 35/100
data-apis/array-api-extra#1002 ·
维护者通常 1 天内回复
查看 data-apis/array-api-extra 的全部 Issue
相似的 Issue
-
难度 2/5 1-3 小时 新手友好度 72/100
Graphify-Labs/graphify#4241 · 1 条评论 ·
维护者通常 1 天内回复
-
难度 1/5 1 小时以内 新手友好度 72/100
-
DeviceTracker未关闭
难度 2/5 1-3 小时 新手友好度 63/100
XiaoMi/ha_xiaomi_home#1821 ·
维护者通常 1 天内回复
-
Maven path-index: "Ambiguous or noncanonical artifact path" error does not report the offending path未关闭
难度 2/5 1-3 小时 新手友好度 76/100
pulp/pulp_maven#524 ·
维护者通常 1 天内回复
-
难度 1/5 1-3 小时 新手友好度 82/100
维护者通常 1 天内回复