jax-ml/jax
improve pjit pytree error involving Nones (and PytreeLeafs)
オープン
#13,074 opened on 2022/11/02
better_errorsgood first issue
Repository metrics
- Stars
- (36,195 個のスター)
- PR merge metrics
- (平均マージ 2d 19h) (30d で 439 merged PRs)
説明
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)