Ruff Compatibility

[siddk]

Hi! I'm trying to use Ruff (https://github.com/charliermarsh/ruff) as my Python linter, but getting some weird rule errors when plugging in to jaxtyping.

Here's a minimal example:

import numpy as np
from jaxtyping import Float32

def hello(x: Float32[np.ndarray, "batch dim"]) -> Float32[np.ndarray, "batch"]:
    return x[:, 1]

if __name__ == "__main__":
    print(hello(np.random.randn(10, 8)))

Running ruff minimal.py (the above file) results in two errors:

• F722 "Syntax error in forward annotation: batch dim"
• F821 "Undefined name batch"

Not sure if this is a jaxtyping or a ruff problem, but saw that jaxtyping is using Ruff to lint now, so curious if this has come up!

[siddk]

I see that F722 is ignored in the repo's pyproject.toml; however ignoring F821 doesn't quite feel safe.

[patrick-kidger]

Take a look at the FAQ regarding flake8 compatibility. I'll update that now to also mention Ruff, which uses the same ruleset.

In your example, that means ignoring F721, and then switching the return annotation to be " batch" rather than "batch".

[siddk]

Thanks so much, that fixed things! Appreciate the explanation in the FAQ as well.