jax-ml/jax

improve pjit pytree error involving Nones (and PytreeLeafs)

Open

#13,074 opened on Nov 2, 2022

View on GitHubContributor guide
 (6 comments) (3 reactions) (1 assignee)Python (3,738 forks)github user discovery
better_errorsgood first issue

Repository metrics

Stars
 (36,174 stars)
PR merge metrics
 (Avg merge 2d 19h) (439 merged PRs in 30d)

Description

When pjit's in_axis_resoures is not a pytree prefix of the corresponding arguments' pytrees, we attempt to raise a good error. But at least two things can be improved:

  • don't mention PytreeLeaf instances in the error message
  • handle None specially (since it's surprising for the error message to refer to it as a container)

Contributor guide

Research direction
Investigate the pjit function's argument validation logic, locate where the error message for mismatched pytree structure is raised, and modify the error handling to avoid mentioning PytreeLeaf instances and to treat None as a leaf rather than a container.
Tech stack
python
Domain
machine learningai
Issue type
Refactor
DifficultyEstimated implementation difficulty for a new contributor, from 1 for very small changes to 5 for expert-level work.
3
Estimated timeA rough time range for an experienced contributor to investigate, implement, test, and prepare a pull request.
1-3 hours
Activity statusHow available the issue appears right now: fresh, active, stale, blocked, or waiting on maintainer input.
Active
ClarityHow clearly the issue explains the expected change, acceptance criteria, and next step.
Clear
Prerequisites
GitPython
Newbie friendlinessA 1-100 score estimating how approachable this issue is for first-time contributors.
65