patrick-kidger / patrick-kidger/jaxtyping

Incompatibility with latest typeguard version

Open
#80 11 comments 7 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

I was updating my typeguard version to the latest release on Github: https://github.com/agronholm/typeguard (4.0.0rc5), and now get conflicts when trying to shape annotate torch Tensors (the code was working fine with previous typeguard versions). In particular running the following simple snippet:

from typeguard import typechecked
from jaxtyping import jaxtyped, Float
from torch import Tensor

@jaxtyped
@typechecked
def foo(x: Float[Tensor, 'a b c']):
    return 1

Will throw this error:

  File "/nfs/homedirs/fuchsgru/graph-active-learning/foo.py", line 6, in <module>
    @typechecked
     ^^^^^^^^^^^
  File "/nfs/staff-ssd/fuchsgru/miniconda3/envs/graph_active_learning/lib/python3.11/site-packages/typeguard/_decorators.py", line 213, in typechecked
    retval = instrument(target)
             ^^^^^^^^^^^^^^^^^^
  File "/nfs/staff-ssd/fuchsgru/miniconda3/envs/graph_active_learning/lib/python3.11/site-packages/typeguard/_decorators.py", line 54, in instrument
    instrumentor.visit(module_ast)
  File "/nfs/staff-ssd/fuchsgru/miniconda3/envs/graph_active_learning/lib/python3.11/ast.py", line 418, in visit
    return visitor(node)
           ^^^^^^^^^^^^^
  File "/nfs/staff-ssd/fuchsgru/miniconda3/envs/graph_active_learning/lib/python3.11/site-packages/typeguard/_transformer.py", line 561, in visit_Module
    self.generic_visit(node)
  File "/nfs/staff-ssd/fuchsgru/miniconda3/envs/graph_active_learning/lib/python3.11/ast.py", line 494, in generic_visit
    value = self.visit(value)
            ^^^^^^^^^^^^^^^^^
  File "/nfs/staff-ssd/fuchsgru/miniconda3/envs/graph_active_learning/lib/python3.11/ast.py", line 418, in visit
    return visitor(node)
           ^^^^^^^^^^^^^
  File "/nfs/staff-ssd/fuchsgru/miniconda3/envs/graph_active_learning/lib/python3.11/site-packages/typeguard/_transformer.py", line 672, in visit_FunctionDef
    annotation = self._convert_annotation(deepcopy(arg.annotation))
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/nfs/staff-ssd/fuchsgru/miniconda3/envs/graph_active_learning/lib/python3.11/site-packages/typeguard/_transformer.py", line 546, in _convert_annotation
    new_annotation = cast(expr, AnnotationTransformer(self).visit(annotation))
                                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/nfs/staff-ssd/fuchsgru/miniconda3/envs/graph_active_learning/lib/python3.11/site-packages/typeguard/_transformer.py", line 339, in visit
    new_node = super().visit(node)
               ^^^^^^^^^^^^^^^^^^^
  File "/nfs/staff-ssd/fuchsgru/miniconda3/envs/graph_active_learning/lib/python3.11/ast.py", line 418, in visit
    return visitor(node)
           ^^^^^^^^^^^^^
  File "/nfs/staff-ssd/fuchsgru/miniconda3/envs/graph_active_learning/lib/python3.11/site-packages/typeguard/_transformer.py", line 391, in visit_Subscript
    items = [self.visit(item) for item in slice_value.elts]
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/nfs/staff-ssd/fuchsgru/miniconda3/envs/graph_active_learning/lib/python3.11/site-packages/typeguard/_transformer.py", line 391, in <listcomp>
    items = [self.visit(item) for item in slice_value.elts]
             ^^^^^^^^^^^^^^^^
  File "/nfs/staff-ssd/fuchsgru/miniconda3/envs/graph_active_learning/lib/python3.11/site-packages/typeguard/_transformer.py", line 339, in visit
    new_node = super().visit(node)
               ^^^^^^^^^^^^^^^^^^^
  File "/nfs/staff-ssd/fuchsgru/miniconda3/envs/graph_active_learning/lib/python3.11/ast.py", line 418, in visit
    return visitor(node)
           ^^^^^^^^^^^^^
  File "/nfs/staff-ssd/fuchsgru/miniconda3/envs/graph_active_learning/lib/python3.11/site-packages/typeguard/_transformer.py", line 442, in visit_Constant
    expression = ast.parse(node.value, mode="eval")
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/nfs/staff-ssd/fuchsgru/miniconda3/envs/graph_active_learning/lib/python3.11/ast.py", line 50, in parse
    return compile(source, filename, mode, flags,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<unknown>", line 1
    a b c
      ^
SyntaxError: invalid syntax

Seems like typeguard is parsing the annotation "a b c" with ast, which does not like this format. Is there any workaround for this? I need the latest typeguard version as earlier versions do not seem to support annotating @property decorated functions.

Relevant versions of typeguard and `jaxtyping:

typeguard                4.0.0rc5.post1
jaxtyping                0.2.15

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 reproducing the supplied decorator snippet with typeguard 4.0.0rc5.post1 and jaxtyping 0.2.15, then inspect how the two libraries process the Float[Tensor, 'a b c'] annotation. No repository file or test is named; done means the reported incompatibility is resolved or a verified workaround is documented.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
devtools
Issue type
Bug
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.