STeP Stream Type

src/argon/base.py

@@ -9,9 +9,7 @@
 def r_resolve(globalns, localns, rarg):
     match rarg:
         case typing.TypeVar():
-            if isinstance(globalns[rarg.__name__], type) and isinstance(
-                localns[rarg.__name__], type
-            ):
+            if rarg.__name__ in globalns and rarg.__name__ in localns:
                 return localns[rarg.__name__]
             else:
                 raise ArgonError(
@@ -103,6 +101,7 @@ def accessor_parent_tparam(self, arg=arg, param=param):  # type: ignore -- PyRig
                             def accessor_parent_tparam(self, arg=arg):  # type: ignore -- PyRight and other tools falsely report this as conflicting defs
 
                                 aug_ns = {}
+                                print(tparam_set)
                                 for key in tparam_set:
                                     aug_ns[key] = getattr(self, key)
 
@@ -116,6 +115,7 @@ def accessor_parent_tparam(self, arg=arg):  # type: ignore -- PyRight and other
 
                                 # recursively resolve the GenericAlias
                                 arg_list = []
+                                print(typing.get_args(arg))
                                 for arg_i in typing.get_args(arg):
                                     arg_list.append(
                                         r_resolve(temp_globalns, temp_localns, arg_i)

src/argon/errors.py

@@ -3,4 +3,4 @@ class ArgonError(Exception):
 
 
 class StagingError(ArgonError):
-    pass
+    pass

step/ops/accum.py

@@ -0,0 +1,26 @@
+import typing
+import pydantic
+from argon.ref import Exp, Op, Sym
+from typing import Tuple, Callable
+
+
+from pydantic.dataclasses import dataclass
+from ..types.stream import Stream, RStream
+
+
+A = typing.TypeVar("A", bound=Exp[typing.Any, typing.Any], covariant=True)
+B = typing.TypeVar("B", bound=Exp[typing.Any, typing.Any], covariant=True)
+a = typing.TypeVar("a", bound=Exp[typing.Any, typing.Any], covariant=True)
+b = typing.TypeVar("b", bound=Exp[typing.Any, typing.Any], covariant=True)
+d = typing.TypeVar("d", bound=Exp[typing.Any, typing.Any], covariant=True)
+
+
+@dataclass(config=pydantic.ConfigDict(arbitrary_types_allowed=True))
+class Accum[A, B, a, b, d](Op[Stream[B, d]]):
+    stream: RStream[A, B, a, b]
+    func: Callable[[A,B], B]
+
+    @property
+    @typing.override
+    def inputs(self) -> typing.List[Sym[typing.Any]]:
+        return [self.stream, self.func]  # type: ignore

step/ops/bufferize.py

@@ -0,0 +1,25 @@
+import typing
+import pydantic
+from argon.ref import Exp, Op, Sym
+from typing import Tuple
+
+
+from pydantic.dataclasses import dataclass
+from ..types.stream import Stream
+from ..types.buffer import Buffer
+
+
+A = typing.TypeVar("A", bound=Exp[typing.Any, typing.Any], covariant=True)
+a = typing.TypeVar("a", bound=Exp[typing.Any, typing.Any], covariant=True)
+b = typing.TypeVar("b", bound=Exp[typing.Any, typing.Any], covariant=True)
+d = typing.TypeVar("d", bound=Exp[typing.Any, typing.Any], covariant=True)
+
+
+@dataclass(config=pydantic.ConfigDict(arbitrary_types_allowed=True))
+class Bufferize[A, a, b, d](Op[Stream[Buffer[A, a], d]]):
+    stream: Stream[A, b]
+
+    @property
+    @typing.override
+    def inputs(self) -> typing.List[Sym[typing.Any]]:
+        return [self.stream]  # type: ignore

step/ops/enumerate.py

@@ -0,0 +1,23 @@
+import typing
+import pydantic
+from argon.ref import Exp, Op, Sym
+from typing import Tuple
+
+
+from pydantic.dataclasses import dataclass
+from ..types.stream import Stream, Index
+
+
+A = typing.TypeVar("A", bound=Exp[typing.Any, typing.Any], covariant=True)
+a = typing.TypeVar("a", bound=Exp[typing.Any, typing.Any], covariant=True)
+
+
+@dataclass(config=pydantic.ConfigDict(arbitrary_types_allowed=True))
+class Enumerate[A, a](Op[Stream[Tuple[A, Index], a]]):
+    stream: Stream[A, a]
+    level: int
+
+    @property
+    @typing.override
+    def inputs(self) -> typing.List[Sym[typing.Any]]:
+        return [self.stream, self.level]  # type: ignore

step/ops/flatmap.py

@@ -0,0 +1,43 @@
+import typing
+import pydantic
+from argon.ref import Exp, Op, Sym
+from typing import Tuple, Callable
+
+
+from pydantic.dataclasses import dataclass
+from ..types.stream import Stream
+
+
+A = typing.TypeVar("A", bound=Exp[typing.Any, typing.Any], covariant=True)
+B = typing.TypeVar("B", bound=Exp[typing.Any, typing.Any], covariant=True)
+a = typing.TypeVar("a", bound=Exp[typing.Any, typing.Any], covariant=True)
+b = typing.TypeVar("b", bound=Exp[typing.Any, typing.Any], covariant=True)
+
+
+@dataclass(config=pydantic.ConfigDict(arbitrary_types_allowed=True))
+class Flatmap[a, b, A, B](Op[Stream[B, a]]):
+    func: Callable[[A], Stream[B, b]]
+    stream: Stream[A, a]
+
+    @property
+    @typing.override
+    def inputs(self) -> typing.List[Sym[typing.Any]]:
+        return [self.func, self.stream]  # type: ignore
+
+# R = typing.TypeVar("R", bound=Exp[typing.Any, typing.Any], covariant=True)
+
+# class Range[R](Op[Stream[B, a]]):
+#     func: R
+
+#     @property
+#     @typing.override
+#     def inputs(self) -> typing.List[Sym[typing.Any]]:
+#         return [self.func, self.stream]  # type: ignore
+
+# class Streamify[A, a]:
+#     func: R
+
+#     @property
+#     @typing.override
+#     def inputs(self) -> typing.List[Sym[typing.Any]]:
+#         return [self.func, self.stream]  # type: ignore

step/ops/flatten.py

@@ -0,0 +1,24 @@
+import typing
+import pydantic
+from argon.ref import Exp, Op, Sym
+from typing import Tuple
+
+
+from pydantic.dataclasses import dataclass
+from ..types.stream import Stream, Index
+
+
+A = typing.TypeVar("A", bound=Exp[typing.Any, typing.Any], covariant=True)
+a = typing.TypeVar("a", bound=Exp[typing.Any, typing.Any], covariant=True)
+b = typing.TypeVar("b", bound=Exp[typing.Any, typing.Any], covariant=True)
+
+
+@dataclass(config=pydantic.ConfigDict(arbitrary_types_allowed=True))
+class Flatten[a, A, b](Op[Stream[A, b]]):
+    stream: Stream[A, a]
+    dims: Tuple[Index, ...]
+
+    @property
+    @typing.override
+    def inputs(self) -> typing.List[Sym[typing.Any]]:
+        return [self.stream, self.dims]  # type: ignore

step/ops/map.py

@@ -0,0 +1,37 @@
+import typing
+import pydantic
+from argon.ref import Exp, Op, Sym
+from typing import Tuple, Callable
+
+from types import FunctionType
+
+
+from pydantic.dataclasses import dataclass
+from ..types.stream import Stream, HStream
+
+
+A = typing.TypeVar("A", bound=Exp[typing.Any, typing.Any], covariant=True)
+B = typing.TypeVar("B", bound=Exp[typing.Any, typing.Any], covariant=True)
+a = typing.TypeVar("a", bound=Exp[typing.Any, typing.Any], covariant=True)
+
+
+@dataclass(config=pydantic.ConfigDict(arbitrary_types_allowed=True))
+class Map[A, B, a](Op[Stream[B, a]]):
+    stream: HStream[A, B, a]
+    func: Callable[[A], B]
+
+    @property
+    @typing.override
+    def inputs(self) -> typing.List[Sym[typing.Any]]:
+        return [self.stream, self.func]  # type: ignore
+
+# P = typing.TypeVar("P", bound=Exp[typing.Any, typing.Any], covariant=True)
+
+# class Permute[P]:
+#     stream: Stream[A, a]
+#     func: Callable[[A], B]
+
+#     @property
+#     @typing.override
+#     def inputs(self) -> typing.List[Sym[typing.Any]]:
+#         return [self.stream, self.func]  # type: ignore

step/ops/partition.py

@@ -0,0 +1,26 @@
+import typing
+import pydantic
+from argon.ref import Exp, Op, Sym
+from typing import Tuple
+
+
+from pydantic.dataclasses import dataclass
+from ..types.stream import Stream, RStream
+
+A = typing.TypeVar("A", bound=Exp[typing.Any, typing.Any], covariant=True)
+B = typing.TypeVar("B", bound=Exp[typing.Any, typing.Any], covariant=True)
+a = typing.TypeVar("a", bound=Exp[typing.Any, typing.Any], covariant=True)
+b = typing.TypeVar("b", bound=Exp[typing.Any, typing.Any], covariant=True)
+d = typing.TypeVar("d", bound=Exp[typing.Any, typing.Any], covariant=True)
+
+
+@dataclass(config=pydantic.ConfigDict(arbitrary_types_allowed=True))
+class Partition[A,B,a,b,d](Op[Stream[A, d]]):
+    stream1: RStream[A, B, a, b]
+    N: int
+    stream2: RStream[A, B, a, b]
+
+    @property
+    @typing.override
+    def inputs(self) -> typing.List[Sym[typing.Any]]:
+        return [self.stream1, self.N, self.stream2]  # type: ignore

step/ops/promote.py

@@ -0,0 +1,24 @@
+import typing
+import pydantic
+from argon.ref import Exp, Op, Sym
+from typing import Tuple
+
+
+from pydantic.dataclasses import dataclass
+from ..types.stream import Stream
+
+
+A = typing.TypeVar("A", bound=Exp[typing.Any, typing.Any], covariant=True)
+a = typing.TypeVar("a", bound=Exp[typing.Any, typing.Any], covariant=True)
+b = typing.TypeVar("b", bound=Exp[typing.Any, typing.Any], covariant=True)
+
+
+@dataclass(config=pydantic.ConfigDict(arbitrary_types_allowed=True))
+class Promote[A, a, b](Op[Stream[A, b]]):
+    stream: Stream[A, a]
+    level: int
+
+    @property
+    @typing.override
+    def inputs(self) -> typing.List[Sym[typing.Any]]:
+        return [self.stream, self.level]  # type: ignore

step/ops/repeat.py

@@ -0,0 +1,25 @@
+import typing
+import pydantic
+from argon.ref import Exp, Op, Sym
+from typing import Tuple, Union, Any
+
+
+from pydantic.dataclasses import dataclass
+from ..types.stream import Stream, RStream
+from ..types.rankgen import RankGen
+
+
+A = typing.TypeVar("A", bound=Exp[typing.Any, typing.Any], covariant=True)
+b = typing.TypeVar("b", bound=Exp[typing.Any, typing.Any], covariant=True)
+c = typing.TypeVar("c", bound=Exp[typing.Any, typing.Any], covariant=True)
+
+
+@dataclass(config=pydantic.ConfigDict(arbitrary_types_allowed=True))
+class Repeat[A, b, c](Op[Stream[A, c]]):
+    stream1: RStream[A, Any, b, c]
+    stream2: RStream[A, Any, b, c]
+
+    @property
+    @typing.override
+    def inputs(self) -> typing.List[Sym[typing.Any]]:
+        return [self.stream1, self.stream2]  # type: ignore

step/ops/reshape.py

@@ -0,0 +1,24 @@
+import typing
+import pydantic
+from argon.ref import Exp, Op, Sym
+from typing import Tuple
+
+
+from pydantic.dataclasses import dataclass
+from ..types.stream import Stream, Index
+
+
+A = typing.TypeVar("A", bound=Exp[typing.Any, typing.Any], covariant=True)
+a = typing.TypeVar("a", bound=Exp[typing.Any, typing.Any], covariant=True)
+b = typing.TypeVar("b", bound=Exp[typing.Any, typing.Any], covariant=True)
+
+@dataclass(config=pydantic.ConfigDict(arbitrary_types_allowed=True))
+class Reshape[A, a, b](Op[Stream[A, b]]):
+    stream: Stream[A, a]
+    dims: Tuple[Index, ...]
+    size: Tuple[int, ...]
+
+    @property
+    @typing.override
+    def inputs(self) -> typing.List[Sym[typing.Any]]:
+        return [self.stream, self.dims, self.size]  # type: ignore

step/ops/window.py

@@ -0,0 +1,24 @@
+# import typing
+# import pydantic
+# from argon.ref import Exp, Op, Sym
+# from typing import Tuple
+
+
+# from pydantic.dataclasses import dataclass
+# from ..types.stream import Stream
+
+# A = typing.TypeVar("A", bound=Exp[typing.Any, typing.Any], covariant=True)
+# N = typing.TypeVar("N", bound=Exp[typing.Any, typing.Any], covariant=True)
+# S = typing.TypeVar("S", bound=Exp[typing.Any, typing.Any], covariant=True)
+# a = typing.TypeVar("a", bound=Exp[typing.Any, typing.Any], covariant=True)
+# b = typing.TypeVar("b", bound=Exp[typing.Any, typing.Any], covariant=True)
+
+
+# @dataclass(config=pydantic.ConfigDict(arbitrary_types_allowed=True))
+# class Window[a, b, A, N, S](Op[Stream[A, a+b]]):
+#     stream: Stream[a, a]
+
+#     @property
+#     @typing.override
+#     def inputs(self) -> typing.List[Sym[typing.Any]]:
+#         return [self.stream]  # type: ignore

step/ops/zip.py

@@ -0,0 +1,24 @@
+import typing
+import pydantic
+from argon.ref import Exp, Op, Sym
+from typing import Tuple
+
+
+from pydantic.dataclasses import dataclass
+from ..types.stream import Stream, HStream
+
+
+A = typing.TypeVar("A", bound=Exp[typing.Any, typing.Any], covariant=True)
+B = typing.TypeVar("B", bound=Exp[typing.Any, typing.Any], covariant=True)
+b = typing.TypeVar("b", bound=Exp[typing.Any, typing.Any], covariant=True)
+
+
+@dataclass(config=pydantic.ConfigDict(arbitrary_types_allowed=True))
+class Zip[A, B, b](Op[Stream[Tuple[A,B], b]]):
+    stream1: HStream[A, B, b]
+    stream2: HStream[A, B, b]
+
+    @property
+    @typing.override
+    def inputs(self) -> typing.List[Sym[typing.Any]]:
+        return [self.stream1, self.stream2]  # type: ignore