add stubs for jax.lax element-wise operations (#4756)

Summary:
Update `tensor-shapes/pyrefly-jax-stubs` with shape annotations for `jax.lax` element-wise APIs. Followup to https://github.com/facebook/pyrefly/issues/4708

The tests required updating some `jax.numpy` stubs to ensure that an empty tuple leads to a concrete 0-D array.

Pull Request resolved: https://github.com/facebook/pyrefly/pull/4756

Test Plan:
Unit test file added
```
$ uv tool run --from ruff==0.16.5 ruff format tensor-shapes
259 files left unchanged
$ uv run python test.py --no-test --no-conformance --no-jsonschema
Running formatting...
Finished in 2.60 seconds.
Running linting...
    Finished `dev` profile [unoptimized + debuginfo] target(s) in 1.06s
Finished in 1.75 seconds.
$ uv run python tensor-shapes/pyrefly-jax-stubs/run_pyrefly.py --release
+ cargo build -p pyrefly --release
    Finished `release` profile [optimized + debuginfo] target(s) in 0.58s
PASS stubs (7 files)
PASS arithmetic (1 files)
PASS creation (1 files)
PASS fft (1 files)
PASS lax (1 files)
PASS matmul (1 files)
PASS nn (1 files)
PASS reductions (1 files)
PASS reshape (1 files)
```

# AI usage

Changes in this PR were generated with a Gemini coding agent.

Reviewed By: stroxler

Differential Revision: D118146402

fbshipit-source-id: a350fb4d8213b396434d22aa30d9ebb95696ee5c
5 files changed
tree: 5e1c0b2e322aa6fa082bfd534d63aab2b14c2d5d
  1. .cargo/
  2. .claude/
  3. .github/
  4. .llms/
  5. .vscode/
  6. conformance/
  7. crates/
  8. lsp/
  9. pyrefly/
  10. pyrefly_wasm/
  11. release_notes/
  12. schemas/
  13. scripts/
  14. tensor-shapes/
  15. test/
  16. website/
  17. .gitattributes
  18. .gitignore
  19. action.yml
  20. AGENTS.md
  21. ARCHITECTURE.md
  22. Cargo.lock
  23. Cargo.toml
  24. CODE_OF_CONDUCT.md
  25. CONTRIBUTING.md
  26. debug.html
  27. empty.py
  28. LICENSE
  29. package-lock.json
  30. README.md
  31. REVIEW.md
  32. ruff.toml
  33. rust-toolchain.toml
  34. rustfmt.toml
  35. TENSOR_SHAPES_CONTRIBUTING.md
  36. test.py
  37. version.bzl
  38. yarn.lock
README.md

Pyrefly: A fast type checker and language server for Python with powerful IDE features

pyrefly PyPI VS Code Open VSX Discord License: MIT

Pyrefly is a type checker and language server for Python, which provides lightning-fast type checking along with IDE features such as code navigation, semantic highlighting, and code completion. It is available as a command-line tool and an extension for popular IDEs and editors such as VSCode, Neovim, Zed, and more.

See the Pyrefly website for full documentation and how to add Pyrefly to your editor of choice.

Pyrefly's current development status is stable.

Key Features

  • Fast. Pyrefly checks over 1.85 million lines of code per second, type checking projects like PyTorch 15x faster than Mypy and Pyright. In the IDE, rechecks typically complete in under 10 milliseconds after saving a file.
  • Production-proven at scale. Pyrefly is the default type checker for Instagram's 20-million-line Python codebase at Meta, and has been adopted by large open source projects including PyTorch and JAX.
  • Full-featured language server. Code navigation, autocomplete, hover information, inlay hints, semantic highlighting, and more, with consistent results across the CLI and your editor of choice.
  • Understands real-world Python. Built-in support for frameworks and tools like Pydantic, Django, and pytest, with model validation, field types, fixture navigation, and autocomplete that work out of the box.
  • Adoption-ready. Migrate from Mypy or Pyright with pyrefly init, silence existing errors with pyrefly suppress, and generate type annotations with pyrefly infer. Start with one file and expand at your own pace.

Getting Started

Version Policy

Pyrefly releases new minor versions (1.x.0) monthly and patch versions in between as-needed for critical fixes. Pyrefly does not follow strict semantic versioning: minor versions contain more significant changes than patch versions, but any version may introduce new type errors and other breaking changes. The pyrefly suppress command can be used to easily silence errors when upgrading to a new version.

Getting Involved

If you have questions or would like to report a bug, please create an issue.

See our contributing guide and architecture overview for information on how to contribute to Pyrefly.

Join our Discord to chat about Pyrefly and types. This is also where we hold biweekly office hours.