wassname/jaxtyping
1
0
Fork 0
mirror of https://github.com/wassname/jaxtyping.git synced 2026-09-09 11:24:55 +08:00
Code Issues Packages Projects Releases Wiki Activity

The import hook now decorates dataclass __init__ methods (#48)

Browse Source
This commit is contained in:
Patrick Kidger
2022-11-16 13:38:04 -08:00
committed by GitHub
parent 784aa78f7c
commit a220df9964
8 changed files with 110 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
API.md
+2
View File
@@ -168,6 +168,8 @@ The import hook can be applied to multiple packages via
install_import_hook(["foo", "bar.baz"], ...)
```
The import hook will automatically decorate all functions, and the `__init__` method of dataclasses.
**Example: writing an end-user script**
```python
jaxtyping/decorator.py
+26 -3
View File
@@ -17,6 +17,7 @@
# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
import dataclasses
import functools as ft
import inspect
import threading
@@ -46,8 +47,30 @@ class _Jaxtyped:
def jaxtyped(fn):
if inspect.isclass(fn): # allow decorators on class definitions
init = jaxtyped(fn.__init__)
fn.__init__ = init
return fn
if dataclasses.is_dataclass(fn):
init = jaxtyped(fn.__init__)
fn.__init__ = init
return fn
else:
raise ValueError(
"jaxtyped may only be added as a class decorator to dataclasses"
)
else:
return ft.wraps(fn)(_Jaxtyped(fn))
def _jaxtyped_typechecker(typechecker):
# typechecker is expected to probably be either `typeguard.typechecked`, or
# `beartype.beartype`, or `None`.
if typechecker is None:
typechecker = lambda x: x
def _wrapper(kls):
assert inspect.isclass(kls)
if dataclasses.is_dataclass(kls):
init = jaxtyped(typechecker(kls.__init__))
kls.__init__ = init
return kls
return _wrapper
jaxtyping/import_hook.py
+23 -15
View File
@@ -67,7 +67,15 @@ def _call_with_frames_removed(f, *args, **kwargs):
def _optimized_cache_from_source(path, debug_override=None):
# Version 2: change the position of the `@jaxtyped` decorator, so need a
# different name to avoid hitting old __pycache__
return cache_from_source(path, debug_override, optimization="jaxtyping2")
# Version 3: now also annotating classes.
return cache_from_source(path, debug_override, optimization="jaxtyping3")
def _dot_lookup(*elements):
out = ast.Name(id=elements[0], ctx=ast.Load())
for element in elements[1:]:
out = ast.Attribute(out, element, ctx=ast.Load())
return out
class _JaxtypingTransformer(ast.NodeVisitor):
@@ -97,6 +105,18 @@ class _JaxtypingTransformer(ast.NodeVisitor):
self._parents.pop()
return node
def visit_ClassDef(self, node: ast.ClassDef):
func = _dot_lookup("jaxtyping", "decorator", "_jaxtyped_typechecker")
if self._typechecker is None:
args = [ast.Constant(None)]
else:
args = [_dot_lookup(*self._typechecker)]
node.decorator_list.append(ast.Call(func, args, keywords=[]))
self._parents.append(node)
self.generic_visit(node)
self._parents.pop()
return node
def visit_FunctionDef(self, node: ast.FunctionDef):
has_annotated_args = any(arg for arg in node.args.args if arg.annotation)
has_annotated_return = bool(node.returns)
@@ -110,24 +130,12 @@ class _JaxtypingTransformer(ast.NodeVisitor):
# case we're just going to have to need to ask the user to remove their
# typechecking annotation (and let this decorator do it instead).
# It's more important we be compatible with normal JAX code.
node.decorator_list.append(
ast.Attribute(
ast.Name(id="jaxtyping", ctx=ast.Load()), "jaxtyped", ast.Load()
),
)
node.decorator_list.append(_dot_lookup("jaxtyping", "jaxtyped"))
if self._typechecker is not None:
# Place at the end of the decorator list, as decorators
# frequently remove annotations from functions and we'd like to
# use those annotations.
typechecker_module, typechecker_function = self._typechecker
node.decorator_list.append(
ast.Attribute(
ast.Name(id=typechecker_module, ctx=ast.Load()),
typechecker_function,
ast.Load(),
)
)
node.decorator_list.append(_dot_lookup(*self._typechecker))
self._parents.append(node)
self.generic_visit(node)
self._parents.pop()
test/import_hook_tester_beartype.py
+13
View File
@@ -17,6 +17,7 @@
# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
import equinox as eqx
import jax.numpy as jnp
import pytest
@@ -32,3 +33,15 @@ def g(x: Float32[jnp.ndarray, " b"]):
g(jnp.array([1.0]))
with pytest.raises(ParamError):
g(jnp.array(1))
class M(eqx.Module):
foo: int
bar: Float32[jnp.ndarray, " a"]
M(1, jnp.array([1.0]))
with pytest.raises(ParamError):
M(1.0, jnp.array([1.0]))
with pytest.raises(ParamError):
M(1, jnp.array(1.0))
test/import_hook_tester_broken_checker.py
+13
View File
@@ -17,6 +17,7 @@
# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
import equinox as eqx
import jax.numpy as jnp
import pytest
@@ -32,3 +33,15 @@ def g(x: Float32[jnp.ndarray, " b"]):
g(jnp.array([1.0]))
with pytest.raises(ParamError):
g(jnp.array(1))
class M(eqx.Module):
foo: int
bar: Float32[jnp.ndarray, " a"]
M(1, jnp.array([1.0]))
with pytest.raises(ParamError):
M(1.0, jnp.array([1.0]))
with pytest.raises(ParamError):
M(1, jnp.array(1.0))
test/import_hook_tester_transitive/another_file.py
+13
View File
@@ -17,6 +17,7 @@
# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
import equinox as eqx
import jax.numpy as jnp
import pytest
@@ -32,3 +33,15 @@ def g(x: Float32[jnp.ndarray, " b"]):
g(jnp.array([1.0]))
with pytest.raises(ParamError):
g(jnp.array(1))
class M(eqx.Module):
foo: int
bar: Float32[jnp.ndarray, " a"]
M(1, jnp.array([1.0]))
with pytest.raises(ParamError):
M(1.0, jnp.array([1.0]))
with pytest.raises(ParamError):
M(1, jnp.array(1.0))
test/import_hook_tester_typeguard.py
+13
View File
@@ -17,6 +17,7 @@
# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
import equinox as eqx
import jax.numpy as jnp
import pytest
@@ -32,3 +33,15 @@ def g(x: Float32[jnp.ndarray, " b"]):
g(jnp.array([1.0]))
with pytest.raises(ParamError):
g(jnp.array(1))
class M(eqx.Module):
foo: int
bar: Float32[jnp.ndarray, " a"]
M(1, jnp.array([1.0]))
with pytest.raises(ParamError):
M(1.0, jnp.array([1.0]))
with pytest.raises(ParamError):
M(1, jnp.array(1.0))
test/test_import_hook.py
+7 -11
View File
@@ -26,9 +26,8 @@ def test_import_hook_typeguard():
hook = install_import_hook(
"test.import_hook_tester_typeguard", ("typeguard", "typechecked")
)
from . import import_hook_tester_typeguard # noqa: F401
hook.uninstall()
with hook:
from . import import_hook_tester_typeguard # noqa: F401
def test_import_hook_beartype():
@@ -40,24 +39,21 @@ def test_import_hook_beartype():
hook = install_import_hook(
"test.import_hook_tester_beartype", ("beartype", "beartype")
)
from . import import_hook_tester_beartype # noqa: F401
hook.uninstall()
with hook:
from . import import_hook_tester_beartype # noqa: F401
def test_import_hook_transitive():
hook = install_import_hook(
"test.import_hook_tester_transitive", ("typeguard", "typechecked")
)
from . import import_hook_tester_transitive # noqa: F401
hook.uninstall()
with hook:
from . import import_hook_tester_transitive # noqa: F401
def test_import_hook_broken_checker():
hook = install_import_hook(
"test.import_hook_tester_broken_checker", ("jaxtyping", "does_not_exist")
)
with pytest.raises(AttributeError):
with hook, pytest.raises(AttributeError):
from . import import_hook_tester_broken_checker # noqa: F401
hook.uninstall()