Added environment config flags.

These flags are `JAXTYPING_DISABLE` and `JAXTYPING_REMOVE_TYPECHECKER_STACK`.

In addition, have now added warnings when using old-style double-decorator syntax, which also serves to guard against the easy mistake of
```python
@jaxtyped(typechecker)
def foo(...)
```
which actually decorates the `typechecker`, not `foo`.
This commit is contained in:
Patrick Kidger
2023-12-05 19:06:00 -08:00
parent 8e47c9081c
commit 125bc89ee9
6 changed files with 160 additions and 46 deletions
+1 -2
View File
@@ -39,8 +39,7 @@ class _ErrorableThread(threading.Thread):
def test_threading_jaxtyped():
@jaxtyped
@typechecked
@jaxtyped(typechecker=typechecked)
def add(x: Float[Array, "a b"], y: Float[Array, "a b"]) -> Float[Array, "a b"]:
return x + y