patrick-kidger / patrick-kidger/jaxtyping

[Feature Request] `+` to match "at least one"/"1 or more" array axes

Open
#378 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
1.9k
Forks
96
PR merge metrics
No merged PRs in 30d

Description

Currently, there is the * modifier to match "0 or more" axes. I think it would be very useful to have the closely related "1 or more" axes modifier, perhaps using +. For example:

Float[Array, "batch *features"] would match arrays with a batch dimension and any number of features, but also matches arrays with no feature axes.

Float[Array, "batch +features"] would only match arrays with a batch dimension and at least one more axis. So [b, a, b, c], [b, a], [b, 1] would all pass, but [b] would not.

Currently, this can be achieved by manually specifying that one differently named axis is present:

Float[Array, "batch feature0 *features"]

But I think it would be both more concise and clear with the + modifier.

Let me know if I'm missing anything with this idea, and thanks again for all your work on this package and the many others.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by tracing the existing * axis-modifier handling and the shape-matching tests that cover it. Use the examples in the issue to define the expected + behavior, with matching for one or more axes and rejection when no additional axis is present.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
tooling
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
42/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.