effectful 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

effectful/handlers/torch.py

@@ -1,7 +1,8 @@
 import functools
 import typing
+from collections.abc import Callable, Mapping, Sequence
 from types import EllipsisType
-from typing import Callable, Mapping, Optional, Sequence, Tuple, TypeVar, Union
+from typing import Annotated, Any
 
 try:
     import torch
@@ -9,62 +10,20 @@ except ImportError:
     raise ImportError("PyTorch is required to use effectful.handlers.torch")
 
 import tree
-from typing_extensions import ParamSpec
 
-import effectful.handlers.numbers  # noqa: F401
 from effectful.internals.runtime import interpreter
-from effectful.ops.semantics import apply, evaluate, fvsof, typeof
-from effectful.ops.syntax import defdata, defop
-from effectful.ops.types import Expr, Operation, Term
-
-P = ParamSpec("P")
-Q = ParamSpec("Q")
-S = TypeVar("S")
-T = TypeVar("T")
-V = TypeVar("V")
-
+from effectful.internals.tensor_utils import _desugar_tensor_index
+from effectful.ops.semantics import apply, evaluate, fvsof, handler, typeof
+from effectful.ops.syntax import Scoped, defdata, defop, defterm, syntactic_eq
+from effectful.ops.types import Expr, NotHandled, Operation, Term
 
 # + An element of a tensor index expression.
-IndexElement = Union[None, int, slice, Sequence[int], EllipsisType, torch.Tensor]
-
-
-def _desugar_tensor_index(shape, key):
-    new_shape = []
-    new_key = []
-
-    def extra_dims(key):
-        return sum(1 for k in key if k is None)
-
-    # handle any missing dimensions by adding a trailing Ellipsis
-    if not any(k is Ellipsis for k in key):
-        key = tuple(key) + (...,)
-
-    for i, k in enumerate(key):
-        if k is None:  # add a new singleton dimension
-            new_shape.append(1)
-            new_key.append(slice(None))
-        elif k is Ellipsis:
-            assert not any(
-                k is Ellipsis for k in key[i + 1 :]
-            ), "only one Ellipsis allowed"
-
-            # determine which of the original dimensions this ellipsis refers to
-            pre_dims = i - extra_dims(key[:i])  # dimensions that precede the ellipsis
-            elided_dims = (
-                len(shape) - pre_dims - (len(key) - i - 1 - extra_dims(key[i + 1 :]))
-            )  #
-            new_shape += shape[pre_dims : pre_dims + elided_dims]
-            new_key += [slice(None)] * elided_dims
-        else:
-            new_shape.append(shape[len(new_shape) - extra_dims(key[:i])])
-            new_key.append(k)
-
-    return new_shape, new_key
+IndexElement = None | int | slice | Sequence[int] | EllipsisType | torch.Tensor
 
 
 def _getitem_ellipsis_and_none(
-    x: torch.Tensor, key: Tuple[IndexElement, ...]
-) -> Tuple[torch.Tensor, Tuple[IndexElement, ...]]:
+    x: torch.Tensor, key: tuple[IndexElement, ...]
+) -> tuple[torch.Tensor, tuple[IndexElement, ...]]:
     """Eliminate ellipses and None in an index expression x[key].
 
     Returns x1, key1 such that x1[key1] == x[key] nand key1 does not contain None or Ellipsis.
@@ -75,7 +34,7 @@ def _getitem_ellipsis_and_none(
     return torch.reshape(x, new_shape), new_key
 
 
-def sizesof(value: Expr) -> Mapping[Operation[[], int], int]:
+def sizesof(value) -> Mapping[Operation[[], torch.Tensor], int]:
     """Return the sizes of named dimensions in a tensor expression.
 
     Sizes are inferred from the tensor shape.
@@ -85,19 +44,14 @@ def sizesof(value: Expr) -> Mapping[Operation[[], int], int]:
 
     **Example usage**:
 
-    >>> a, b = defop(int, name='a'), defop(int, name='b')
-    >>> sizesof(Indexable(torch.ones(2, 3))[a(), b()])
-    {a: 2, b: 3}
+    >>> a, b = defop(torch.Tensor, name='a'), defop(torch.Tensor, name='b')
+    >>> sizes = sizesof(torch.ones(2, 3)[a(), b()])
+    >>> assert sizes[a] == 2 and sizes[b] == 3
     """
-    if isinstance(value, torch.distributions.Distribution) and not isinstance(
-        value, Term
-    ):
-        return {v: s for a in value.__dict__.values() for v, s in sizesof(a).items()}
-
-    sizes: dict[Operation[[], int], int] = {}
+    sizes: dict[Operation[[], torch.Tensor], int] = {}
 
     def _torch_getitem_sizeof(
-        x: Expr[torch.Tensor], key: Tuple[Expr[IndexElement], ...]
+        x: Expr[torch.Tensor], key: tuple[Expr[IndexElement], ...]
     ) -> Expr[torch.Tensor]:
         if isinstance(x, torch.Tensor):
             shape, key_ = _desugar_tensor_index(x.shape, key)
@@ -107,7 +61,7 @@ def sizesof(value: Expr) -> Mapping[Operation[[], int], int]:
                     isinstance(k, Term)
                     and len(k.args) == 0
                     and len(k.kwargs) == 0
-                    and issubclass(typeof(k), int)
+                    and issubclass(typeof(k), torch.Tensor)
                 ):
                     if k.op in sizes and sizes[k.op] != shape[i]:
                         raise ValueError(
@@ -117,55 +71,50 @@ def sizesof(value: Expr) -> Mapping[Operation[[], int], int]:
 
         return defdata(torch_getitem, x, key)
 
-    with interpreter(
-        {
-            torch_getitem: _torch_getitem_sizeof,
-            apply: lambda _, op, *a, **k: defdata(op, *a, **k),
-        }
-    ):
-        evaluate(value)
-
-    return sizes
+    def _apply(op, *args, **kwargs):
+        args, kwargs = tree.map_structure(defterm, (args, kwargs))
+        return defdata(op, *args, **kwargs)
 
+    with interpreter({torch_getitem: _torch_getitem_sizeof, apply: _apply}):
+        evaluate(defterm(value))
 
-def _partial_eval(t: T, order: Optional[Sequence[Operation[[], int]]] = None) -> T:
-    """Partially evaluate a term with respect to its sized free variables.
-
-    Variables in `order` are converted to positional dimensions in the result
-    tensor, in the order they appear. All other variables remain free.
+    return sizes
 
-    """
-    from effectful.ops.syntax import deffn
 
-    if order is None:
-        order = []
+def _partial_eval(t: Expr[torch.Tensor]) -> Expr[torch.Tensor]:
+    """Partially evaluate a term with respect to its sized free variables."""
 
     sized_fvs = sizesof(t)
+    if not sized_fvs:
+        return t
 
-    for x in order:
-        if x not in sized_fvs:
-            raise ValueError(
-                f"Tried to partially evaluate nonexistent free variable {x} (free={sized_fvs})"
-            )
-
-    # if there are no sized free variables, then nothing to do
-    if len(sized_fvs) == 0:
+    if not (
+        isinstance(t, Term)
+        and all(
+            isinstance(a, torch.Tensor) or not isinstance(a, Term) or a.op in sized_fvs
+            for a in tree.flatten((t.args, t.kwargs))
+        )
+    ):
         return t
 
-    order_set = set(order)
-    reindex_fvs = [
-        (var, size) for var, size in sized_fvs.items() if var not in order_set
-    ]
-    ordered_sized_fvs = reindex_fvs + [(var, sized_fvs[var]) for var in order]
+    # note: torch.func.vmap will call repr on the callable, so it's important
+    # that we don't pass something with a slow repr (like a large tensor wrapped
+    # in a deffn)
+    def wrapper(*sized_values):
+        with handler(
+            {
+                k: functools.partial(lambda x: x, v)
+                for (k, v) in zip(sized_fvs.keys(), sized_values)
+            }
+        ):
+            return evaluate(t)
 
-    tpe_torch_fn = torch.func.vmap(
-        deffn(t, *[var for (var, _) in ordered_sized_fvs]), randomness="different"
-    )
+    tpe_torch_fn = torch.func.vmap(wrapper, randomness="different")
 
     inds = torch.broadcast_tensors(
         *(
-            torch.arange(size)[(...,) + (None,) * (len(ordered_sized_fvs) - i - 1)]
-            for i, (_, size) in enumerate(ordered_sized_fvs)
+            torch.arange(size)[(...,) + (None,) * (len(sized_fvs) - i - 1)]
+            for i, size in enumerate(sized_fvs.values())
         )
     )
 
@@ -176,38 +125,125 @@ def _partial_eval(t: T, order: Optional[Sequence[Operation[[], int]]] = None) ->
             return t
 
         result = t.reshape(inds[0].shape + t.shape[1:])
-        return torch_getitem(result, tuple(var() for (var, _) in reindex_fvs))
-
-    return tree.map_structure(reindex_flat_tensor, flat_result)
-
-
-def to_tensor(*args, **kwargs) -> torch.Tensor:
+        return torch_getitem(result, tuple(k() for k in sized_fvs.keys()))
+
+    result = tree.map_structure(reindex_flat_tensor, flat_result)
+    return result
+
+
+@defop
+@functools.singledispatch
+def bind_dims[
+    A,
+    B,
+    HasDims: torch.Tensor
+    | torch.distributions.Distribution
+    | tree.Structure[torch.Tensor | torch.distributions.Distribution],
+](
+    value: Annotated[HasDims, Scoped[A | B]],
+    *names: Annotated[Operation[[], torch.Tensor], Scoped[B]],
+) -> Annotated[HasDims, Scoped[A]]:
     """Convert named dimensions to positional dimensions.
 
     :param t: A tensor.
-    :type t: T
-    :param order: A list of named dimensions to convert to positional dimensions.
+    :param args: Named dimensions to convert to positional dimensions.
                   These positional dimensions will appear at the beginning of the
                   shape.
-    :type order: Optional[Sequence[Operation[[], int]]]
-    :return: A tensor with the named dimensions in ``order`` converted to positional dimensions.
+    :return: A tensor with the named dimensions in ``args`` converted to positional dimensions.
 
     **Example usage**:
 
-    >>> a, b = defop(int, name='a'), defop(int, name='b')
+    >>> a, b = defop(torch.Tensor, name='a'), defop(torch.Tensor, name='b')
     >>> t = torch.ones(2, 3)
-    >>> to_tensor(Indexable(t)[a(), b()], [b, a]).shape
+    >>> bind_dims(t[a(), b()], b, a).shape
     torch.Size([3, 2])
     """
-    return _partial_eval(*args, **kwargs)
+    if tree.is_nested(value):
+        return tree.map_structure(lambda v: bind_dims(v, *names), value)
+    raise NotHandled
+
+
+@bind_dims.register  # type: ignore
+def _bind_dims_tensor(
+    value: torch.Tensor, *names: Operation[[], torch.Tensor]
+) -> torch.Tensor:
+    def _evaluate(expr):
+        if isinstance(expr, Term):
+            (args, kwargs) = tree.map_structure(_evaluate, (expr.args, expr.kwargs))
+            return _partial_eval(expr)
+        if tree.is_nested(expr):
+            return tree.map_structure(_evaluate, expr)
+        return expr
+
+    t = value
+    args = names
+
+    if not isinstance(t, Term):
+        return t
+
+    result = _evaluate(t)
+    if not isinstance(result, Term) or not args:
+        return result
+
+    # ensure that the result is a torch_getitem with a tensor as the first argument
+    if not (result.op is torch_getitem and isinstance(result.args[0], torch.Tensor)):
+        raise NotHandled
+
+    tensor = result.args[0]
+    dims = result.args[1]
+    assert isinstance(dims, Sequence)
+
+    # ensure that the order is a subset of the named dimensions
+    order_set = set(args)
+    if not order_set <= set(a.op for a in dims if isinstance(a, Term)):
+        raise NotHandled
+
+    # permute the inner tensor so that the leading dimensions are in the order
+    # specified and the trailing dimensions are the remaining named dimensions
+    # (or slices)
+    reindex_dims = [
+        i
+        for i, o in enumerate(dims)
+        if not isinstance(o, Term) or o.op not in order_set
+    ]
+    dim_ops = [a.op if isinstance(a, Term) else None for a in dims]
+    perm = [dim_ops.index(o) for o in args] + reindex_dims
+    tensor = tensor.permute(perm)
+    return tensor[(slice(None),) * len(args) + tuple(dims[i] for i in reindex_dims)]
+
+
+@defop
+@functools.singledispatch
+def unbind_dims[
+    A,
+    B,
+    HasDims: torch.Tensor
+    | torch.distributions.Distribution
+    | tree.Structure[torch.Tensor | torch.distributions.Distribution],
+](
+    value: Annotated[HasDims, Scoped[A | B]],
+    *names: Annotated[Operation[[], torch.Tensor], Scoped[B]],
+) -> Annotated[HasDims, Scoped[A | B]]:
+    if tree.is_nested(value):
+        return tree.map_structure(lambda v: unbind_dims(v, *names), value)
+    raise NotHandled
+
+
+@unbind_dims.register  # type: ignore
+def _unbind_dims_tensor[A, B](
+    value: torch.Tensor,
+    *names: Annotated[Operation[[], torch.Tensor], Scoped[B]],
+) -> Annotated[torch.Tensor, Scoped[A | B]]:
+    return value[tuple(n() for n in names)]
 
 
 @functools.cache
-def _register_torch_op(torch_fn: Callable[P, T]):
+def _register_torch_op[**P, T](torch_fn: Callable[P, T]):
+    if torch_fn is torch._C.TensorBase.__getitem__:
+        return torch_getitem
 
     @defop
     def _torch_op(*args, **kwargs) -> torch.Tensor:
-
         tm = defdata(_torch_op, *args, **kwargs)
         sized_fvs = sizesof(tm)
 
@@ -218,7 +254,7 @@ def _register_torch_op(torch_fn: Callable[P, T]):
             and args[1]
             and all(isinstance(k, Term) and k.op in sized_fvs for k in args[1])
         ):
-            raise NotImplementedError
+            raise NotHandled
         elif sized_fvs and set(sized_fvs.keys()) == fvsof(tm) - {
             torch_getitem,
             _torch_op,
@@ -234,20 +270,19 @@ def _register_torch_op(torch_fn: Callable[P, T]):
         ):
             return typing.cast(torch.Tensor, torch_fn(*args, **kwargs))
         else:
-            raise NotImplementedError
+            raise NotHandled
 
     functools.update_wrapper(_torch_op, torch_fn)
     return _torch_op
 
 
 @_register_torch_op
-def torch_getitem(x: torch.Tensor, key: Tuple[IndexElement, ...]) -> torch.Tensor:
+def torch_getitem(x: torch.Tensor, key: tuple[IndexElement, ...]) -> torch.Tensor:
     """Operation for indexing a tensor.
 
     .. note::
 
-      This operation is not intended to be called directly. Instead, use
-      :class:`Indexable` to create indexed tensors. :func:`torch_getitem` is
+      This operation is not intended to be called directly. Instead, it is
       exposed so that it can be handled.
 
     """
@@ -288,36 +323,13 @@ def torch_getitem(x: torch.Tensor, key: Tuple[IndexElement, ...]) -> torch.Tenso
                 key_l[i] = flat_arg.reshape((-1,) + (1,) * i)
         elif isinstance(arg, int):
             key_l[i] = torch.tensor(arg, dtype=torch.long, device=x.device)
-        elif isinstance(arg, (list, tuple)):
+        elif isinstance(arg, list | tuple):
             flat_arg = torch.tensor(arg, dtype=torch.long, device=x.device)
             key_l[i] = flat_arg.reshape(flat_arg.shape + (1,) * i)
 
     return torch.ops.aten.index(x, tuple(key_l))
 
 
-class Indexable:
-    """Helper class for constructing indexed tensors.
-
-    **Example usage**:
-
-    >>> width, height = defop(int, name='width'), defop(int, name='height')
-    >>> t = Indexable(torch.ones(2, 3))[width(), height()]
-    >>> t
-    Indexable(tensor([[1., 1., 1.],
-                      [1., 1., 1.]]))[width(), height()]
-    """
-
-    def __init__(self, t: torch.Tensor):
-        if not isinstance(t, torch.Tensor):
-            raise ValueError(f"Expected a torch.Tensor, got {type(t)}")
-        self.t = t
-
-    def __getitem__(self, key) -> torch.Tensor:
-        if not isinstance(key, tuple):
-            key = (key,)
-        return torch_getitem(self.t, key)
-
-
 @defdata.register(torch.Tensor)
 def _embed_tensor(op, *args, **kwargs):
     if (
@@ -325,7 +337,7 @@ def _embed_tensor(op, *args, **kwargs):
         and not isinstance(args[0], Term)
         and len(args[1]) > 0
         and all(
-            typeof(k) is int and not k.args and not k.kwargs
+            typeof(k) is torch.Tensor and not k.args and not k.kwargs
             for k in args[1]
             if isinstance(k, Term)
         )
@@ -356,52 +368,161 @@ class _TensorTerm(Term[torch.Tensor]):
         return self._kwargs
 
     def __getitem__(
-        self, key: Union[Expr[IndexElement], Tuple[Expr[IndexElement], ...]]
+        self, key: Expr[IndexElement] | tuple[Expr[IndexElement], ...]
     ) -> Expr[torch.Tensor]:
         return torch_getitem(self, key if isinstance(key, tuple) else (key,))
 
     @classmethod
-    def __torch_function__(
+    def __torch_function__[T](
         cls, func: Callable[..., T], types, args=(), kwargs=None
     ) -> Expr[T]:
         return _register_torch_op(func)(*args, **({} if kwargs is None else kwargs))
 
+    def __add__(self, other: torch.Tensor) -> torch.Tensor:
+        return torch.add(typing.cast(torch.Tensor, self), other)
+
+    def __radd__(self, other: torch.Tensor) -> torch.Tensor:
+        return torch.add(other, typing.cast(torch.Tensor, self))
+
+    def __neg__(self) -> torch.Tensor:
+        return torch.neg(typing.cast(torch.Tensor, self))
+
+    def __pos__(self) -> torch.Tensor:
+        return torch.positive(typing.cast(torch.Tensor, self))
+
+    def __sub__(self, other: torch.Tensor) -> torch.Tensor:
+        return torch.sub(typing.cast(torch.Tensor, self), other)
+
+    def __rsub__(self, other: torch.Tensor) -> torch.Tensor:
+        return torch.sub(other, typing.cast(torch.Tensor, self))
+
+    def __mul__(self, other: torch.Tensor) -> torch.Tensor:
+        return torch.mul(typing.cast(torch.Tensor, self), other)
+
+    def __rmul__(self, other: torch.Tensor) -> torch.Tensor:
+        return torch.mul(other, typing.cast(torch.Tensor, self))
+
+    def __truediv__(self, other: torch.Tensor) -> torch.Tensor:
+        return torch.div(typing.cast(torch.Tensor, self), other)
+
+    def __rtruediv__(self, other: torch.Tensor) -> torch.Tensor:
+        return torch.div(other, typing.cast(torch.Tensor, self))
+
+    def __pow__(self, other: torch.Tensor) -> torch.Tensor:
+        return torch.pow(typing.cast(torch.Tensor, self), other)
+
+    def __rpow__(self, other: torch.Tensor) -> torch.Tensor:
+        return torch.pow(other, typing.cast(torch.Tensor, self))
+
+    def __abs__(self) -> torch.Tensor:
+        return torch.abs(typing.cast(torch.Tensor, self))
+
+    def __eq__(self, other: Any):
+        return torch.eq(typing.cast(torch.Tensor, self), other)
+
+    def __ne__(self, other: Any):
+        return torch.ne(typing.cast(torch.Tensor, self), other)
+
+    def __floordiv__(self, other: torch.Tensor) -> torch.Tensor:
+        return torch.floor_divide(typing.cast(torch.Tensor, self), other)
+
+    def __rfloordiv__(self, other: torch.Tensor) -> torch.Tensor:
+        return torch.floor_divide(other, typing.cast(torch.Tensor, self))
+
+    def __mod__(self, other: torch.Tensor) -> torch.Tensor:
+        return torch.fmod(typing.cast(torch.Tensor, self), other)
+
+    def __rmod__(self, other: torch.Tensor) -> torch.Tensor:
+        return torch.fmod(other, typing.cast(torch.Tensor, self))
+
+    def __lt__(self, other: torch.Tensor) -> torch.Tensor:
+        return torch.lt(typing.cast(torch.Tensor, self), other)
+
+    def __le__(self, other: torch.Tensor) -> torch.Tensor:
+        return torch.le(typing.cast(torch.Tensor, self), other)
+
+    def __gt__(self, other: torch.Tensor) -> torch.Tensor:
+        return torch.gt(typing.cast(torch.Tensor, self), other)
+
+    def __ge__(self, other: torch.Tensor) -> torch.Tensor:
+        return torch.ge(typing.cast(torch.Tensor, self), other)
+
+    def __lshift__(self, other: torch.Tensor) -> torch.Tensor:
+        return torch.bitwise_left_shift(typing.cast(torch.Tensor, self), other)
+
+    def __rlshift__(self, other: torch.Tensor) -> torch.Tensor:
+        return torch.bitwise_left_shift(other, typing.cast(torch.Tensor, self))
+
+    def __rshift__(self, other: torch.Tensor) -> torch.Tensor:
+        return torch.bitwise_right_shift(typing.cast(torch.Tensor, self), other)
+
+    def __rrshift__(self, other: torch.Tensor) -> torch.Tensor:
+        return torch.bitwise_right_shift(other, typing.cast(torch.Tensor, self))
+
+    def __and__(self, other: torch.Tensor) -> torch.Tensor:
+        return torch.bitwise_and(typing.cast(torch.Tensor, self), other)
+
+    def __rand__(self, other: torch.Tensor) -> torch.Tensor:
+        return torch.bitwise_and(other, typing.cast(torch.Tensor, self))
+
+    def __xor__(self, other: torch.Tensor) -> torch.Tensor:
+        return torch.bitwise_xor(typing.cast(torch.Tensor, self), other)
+
+    def __rxor__(self, other: torch.Tensor) -> torch.Tensor:
+        return torch.bitwise_xor(other, typing.cast(torch.Tensor, self))
+
+    def __or__(self, other: torch.Tensor) -> torch.Tensor:
+        return torch.bitwise_or(typing.cast(torch.Tensor, self), other)
+
+    def __ror__(self, other: torch.Tensor) -> torch.Tensor:
+        return torch.bitwise_or(other, typing.cast(torch.Tensor, self))
+
+    def __invert__(self) -> torch.Tensor:
+        return torch.bitwise_not(typing.cast(torch.Tensor, self))
+
+    def __matmul__(self, other: torch.Tensor) -> torch.Tensor:
+        return torch.matmul(typing.cast(torch.Tensor, self), other)
+
+    def __rmatmul__(self, other: torch.Tensor) -> torch.Tensor:
+        return torch.matmul(other, typing.cast(torch.Tensor, self))
+
+    def __iter__(self):
+        raise TypeError("A free tensor is not iterable.")
+
 
 @Term.register
 class _EagerTensorTerm(torch.Tensor):
-
-    op: Operation[..., torch.Tensor] = torch_getitem
-    args: Tuple[torch.Tensor, Tuple[IndexElement, ...]]
+    args: tuple[torch.Tensor, tuple[IndexElement, ...]]
     kwargs: Mapping[str, object] = {}
 
     __match_args__ = ("op", "args", "kwargs")
 
-    def __new__(cls, x: torch.Tensor, key: Tuple[IndexElement, ...]):
+    def __new__(cls, x: torch.Tensor, key: tuple[IndexElement, ...]):
         assert not isinstance(x, Term)
 
         for k in key:
             if isinstance(k, Term):
-                assert typeof(k) is int and not k.args and not k.kwargs
+                assert typeof(k) is torch.Tensor and not k.args and not k.kwargs
 
         x, key = _getitem_ellipsis_and_none(x, key)
         ret = x.as_subclass(cls)
         ret.args = (x, key)
         return ret
 
-    def __repr__(self):
-        indexed_constr = "Indexable"
-
-        # correct indentation
-        parts = str(self.args[0]).split("\n")
-        tensor_str = "\n".join(
-            [parts[0]] + [(len(indexed_constr) + 1) * " " + p for p in parts[1:]]
-        )
+    @property
+    def op(self) -> Operation[..., torch.Tensor]:
+        return torch_getitem
 
+    def __str__(self):
+        tensor_str = str(self.args[0])
         key_str = ", ".join(str(k) for k in self.args[1])
-        return f"{indexed_constr}({tensor_str})[{key_str}]"
+        return f"{tensor_str}[{key_str}]"
+
+    def __repr__(self):
+        return str(self)
 
     @classmethod
-    def __torch_function__(
+    def __torch_function__[T](
         cls, func: Callable[..., T], types, args=(), kwargs=None
     ) -> Expr[T]:
         return _register_torch_op(func)(*args, **({} if kwargs is None else kwargs))
@@ -422,7 +543,7 @@ class _EagerTensorTerm(torch.Tensor):
         x, key = self.args
         return torch.Size([s for s, k in zip(x.shape, key) if not isinstance(k, Term)])
 
-    def size(self, dim: Optional[int] = None):
+    def size(self, dim: int | None = None):
         if dim is None:
             return self.shape
         return self.shape[dim]
@@ -458,14 +579,17 @@ class _EagerTensorTerm(torch.Tensor):
     def requires_grad(self):
         return self.args[0].requires_grad
 
+    def requires_grad_(self, requires_grad=True):
+        return self.args[0].requires_grad_(requires_grad=requires_grad)
+
     @property
     def grad_fn(self):
         return self.args[0].grad_fn
 
 
-def _indexed_func_wrapper(
-    func: Callable[P, T]
-) -> Tuple[Callable[P, S], Callable[[S], T]]:
+def _indexed_func_wrapper[**P, S, T](
+    func: Callable[P, T],
+) -> tuple[Callable[P, S], Callable[[S], T]]:
     # index expressions for the result of the function
     indexes = None
 
@@ -480,7 +604,7 @@ def _indexed_func_wrapper(
         nonlocal indexes
 
         def deindex_tensor(t, i):
-            t_ = to_tensor(t, i.sizes.keys())
+            t_ = bind_dims(t, *i.sizes.keys())
             assert all(t_.shape[j] == i.sizes[v] for j, v in enumerate(i.sizes))
             return t_
 
@@ -554,7 +678,7 @@ def vjp(func, *indexed_primals, **kwargs):
     unpacked_primals = []
     for t in indexed_primals:
         indices = list(sizesof(t).keys())
-        unpacked = to_tensor(t, indices)
+        unpacked = bind_dims(t, *indices)
         unpacked_primals.append((unpacked, indices))
 
     indexed_result = None
@@ -569,7 +693,7 @@ def vjp(func, *indexed_primals, **kwargs):
         nonlocal indexed_result
         indexed_result = func(*repack_primals(primals))
         return tree.map_structure(
-            lambda t: to_tensor(t, list(sizesof(t).keys())), indexed_result
+            lambda t: bind_dims(t, *list(sizesof(t).keys())), indexed_result
         )
 
     unindexed_primals = [t[0] for t in unpacked_primals]
@@ -577,7 +701,7 @@ def vjp(func, *indexed_primals, **kwargs):
 
     def vjpfunc_wrapper(*tangents):
         unindexed_tangents = tree.map_structure(
-            lambda t: to_tensor(t, list(sizesof(t).keys())), tangents
+            lambda t: bind_dims(t, *list(sizesof(t).keys())), tangents
         )
         grads = vjpfunc(*unindexed_tangents)
         return repack_primals(grads)
@@ -593,3 +717,8 @@ def vmap(func, *args, **kwargs):
     # indexed_dim_n, pos_dim_1, ..., pos_dim_m], so we reapply indexes starting
     # at dim 1
     return lambda *a, **k: reindex(vmap_func(*a, *k), starting_dim=1)
+
+
+@syntactic_eq.register
+def _(x: torch.Tensor, other) -> bool:
+    return isinstance(other, torch.Tensor) and bool((x == other).all())