patrick-kidger / patrick-kidger/jaxtyping

Disable jaxtyping for individual tests

Open
#331 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

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

Description

Hi, I'm using jaxtyping in my test suite have a particular test that is slowed down extremely by jaxtyping and typeguard. I'm talking ~7 minutes with

[tool.pytest.ini_options]
addopts = "--jaxtyping-packages=mypackage,typeguard.typechecked"

and only ~4s without. I tried disabling type checking for this test using

@typing.no_type_check
def test_my_function() -> None:
    ...

but sadly, that doesn't seem to speed things up.

Are there any workarounds for this?

P.S.:

The above example is a simplified version of the following decorator which I am already using successfully on "normal" functions outside of the test files.

W = TypeVar("W")
def no_runtime_typecheck(fun: W) -> W:
    """Applies the ``no_type_check`` decorator during runtime (for typeguard/beartype/
    jaxtyping/...) but not during static analysis (e.g. mypy)."""
    if typing.TYPE_CHECKING:
        return fun
    return typing.no_type_check(fun)  # type: ignore[unreachable]

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 how the pytest option --jaxtyping-packages=mypackage,typeguard.typechecked applies runtime checking to individual tests, and compare that behavior with typing.no_type_check and the shown no_runtime_typecheck decorator. Done means a single test can opt out of jaxtyping and typeguard overhead without disabling checking for the rest of the suite; the issue does not name a file or test to run.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
testing-qa
Issue type
Feature
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.