Skip to content

Commit 0338fcc

Browse files
authored
Fix recursion error in sizesof (#406)
* fix recursion error in sizesof * format
1 parent 6524e58 commit 0338fcc

1 file changed

Lines changed: 3 additions & 2 deletions

File tree

effectful/handlers/jax/_handlers.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,8 @@
1212

1313
import tree
1414

15-
from effectful.ops.semantics import apply, evaluate, fvsof, handler, typeof
15+
from effectful.internals.runtime import interpreter
16+
from effectful.ops.semantics import apply, evaluate, fvsof, typeof
1617
from effectful.ops.syntax import (
1718
Scoped,
1819
_CustomSingleDispatchCallable,
@@ -74,7 +75,7 @@ def _getitem_sizeof(x: jax.Array, key: tuple[Expr[IndexElement], ...]):
7475
def _apply(op, *args, **kwargs):
7576
return defdata(op, *args, **kwargs)
7677

77-
with handler({jax_getitem: _getitem_sizeof, apply: _apply}):
78+
with interpreter({jax_getitem: _getitem_sizeof, apply: _apply}):
7879
evaluate(value)
7980

8081
return sizes

0 commit comments

Comments
 (0)