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

effectful/internals/runtime.py

@@ -1,20 +1,14 @@
 import contextlib
 import dataclasses
 import functools
-from typing import Callable, Generic, Mapping, Tuple, TypeVar
-
-from typing_extensions import ParamSpec
+from collections.abc import Callable, Mapping
 
 from effectful.ops.syntax import defop
 from effectful.ops.types import Interpretation, Operation
 
-P = ParamSpec("P")
-S = TypeVar("S")
-T = TypeVar("T")
-
 
 @dataclasses.dataclass
-class Runtime(Generic[S, T]):
+class Runtime[S, T]:
     interpretation: "Interpretation[S, T]"
 
 
@@ -29,7 +23,6 @@ def get_interpretation():
 
 @contextlib.contextmanager
 def interpreter(intp: "Interpretation"):
-
     r = get_runtime()
     old_intp = r.interpretation
     try:
@@ -40,11 +33,11 @@ def interpreter(intp: "Interpretation"):
 
 
 @defop
-def _get_args() -> Tuple[Tuple, Mapping]:
+def _get_args() -> tuple[tuple, Mapping]:
     return ((), {})
 
 
-def _restore_args(fn: Callable[P, T]) -> Callable[P, T]:
+def _restore_args[**P, T](fn: Callable[P, T]) -> Callable[P, T]:
     @functools.wraps(fn)
     def _cont_wrapper(*a: P.args, **k: P.kwargs) -> T:
         a, k = (a, k) if a or k else _get_args()  # type: ignore
@@ -53,7 +46,7 @@ def _restore_args(fn: Callable[P, T]) -> Callable[P, T]:
     return _cont_wrapper
 
 
-def _save_args(fn: Callable[P, T]) -> Callable[P, T]:
+def _save_args[**P, T](fn: Callable[P, T]) -> Callable[P, T]:
     from effectful.ops.semantics import handler
 
     @functools.wraps(fn)
@@ -64,7 +57,7 @@ def _save_args(fn: Callable[P, T]) -> Callable[P, T]:
     return _cont_wrapper
 
 
-def _set_prompt(
+def _set_prompt[**P, T](
     prompt: Operation[P, T], cont: Callable[P, T], body: Callable[P, T]
 ) -> Callable[P, T]:
     from effectful.ops.semantics import handler

effectful/internals/tensor_utils.py

@@ -0,0 +1,32 @@
+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