From 4386393938cc8a77186cd514da92797397686f58 Mon Sep 17 00:00:00 2001 From: charlieli13 Date: Tue, 28 May 2024 11:53:00 -0700 Subject: [PATCH 1/7] implemented preliminary stream and buffer types --- src/argon/types/buffer.py | 20 ++++++++++++ src/argon/types/stream.py | 26 +++++++++++++++ tests/test_step.py | 67 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 113 insertions(+) create mode 100644 src/argon/types/buffer.py create mode 100644 src/argon/types/stream.py create mode 100644 tests/test_step.py diff --git a/src/argon/types/buffer.py b/src/argon/types/buffer.py new file mode 100644 index 0000000..2ad33ec --- /dev/null +++ b/src/argon/types/buffer.py @@ -0,0 +1,20 @@ +from pydantic.dataclasses import dataclass +from typing import Union, Tuple, override, List, TypeVar +from argon.ref import Ref +import numpy as np + +T = TypeVar("T", int, float, Tuple[int], Tuple[float]) +@dataclass +class Element: + value: T + + def __str__(self) -> str: + return str(self.value) + +BT = TypeVar("BT") + +class Buffer[BT](Ref[np.ndarray, "Buffer[BT]"]): + + @override + def fresh(self) -> "Buffer[BT]": + return Buffer[self.BT]() \ No newline at end of file diff --git a/src/argon/types/stream.py b/src/argon/types/stream.py new file mode 100644 index 0000000..696b834 --- /dev/null +++ b/src/argon/types/stream.py @@ -0,0 +1,26 @@ +from pydantic.dataclasses import dataclass +from typing import Union, Tuple, override, List, TypeVar +from argon.ref import Ref +from argon.types.buffer import Element, Buffer + +@dataclass +class Stop: + level: int + + def __str__(self) -> str: + return f"S{self.level}" + +T = TypeVar("T", Element, Buffer, Tuple[Buffer]) +@dataclass +class Val: + value: T + + def __str__(self) -> str: + return str(self.value) + +VT = TypeVar("VT") +class Stream[VT](Ref[List[Union[Val, Stop]], "Stream[VT]"]): + + @override + def fresh(self) -> "Stream[VT]": + return Stream[self.VT]() \ No newline at end of file diff --git a/tests/test_step.py b/tests/test_step.py new file mode 100644 index 0000000..0a10cb5 --- /dev/null +++ b/tests/test_step.py @@ -0,0 +1,67 @@ +from argon.state import State +from argon.types.stream import Stop, Val, Stream +from argon.types.buffer import Buffer +from typing import Union, Tuple, override, List, TypeVar +import numpy as np +import typing + +# Buffer tests +def test_buffer(): + state = State() + with state: + a = Buffer[int]().const(np.ndarray([2])) + assert a.C == np.ndarray + assert a.A is Buffer[int] + assert a.BT is int + assert a.A().BT is int + +# Stream tests +def test_int(): + state = State() + with state: + a = Stream[int]().const( + [Val(1.0), Val(2.0), Stop(1), Val(3.0), Val(4.0), Stop(2)] + ) + assert a.C == typing.List[typing.Union[Val, Stop]] + assert a.A is Stream[int] + assert a.T is int + assert a.A().T is int + print(state) + +def test_float(): + state = State() + with state: + a = Stream[float]().const( + [Val(1.5), Val(2.5), Stop(1), Val(3.5), Val(4.5), Stop(2)] + ) + assert a.C == typing.List[typing.Union[Val, Stop]] + assert a.A is Stream[float] + assert a.T is float + assert a.A().T is float + print(state) + +def test_tuple(): + state = State() + with state: + a = Stream[Tuple[int]]().const( + [Tuple[Val(1.0), Val(2.0)], Stop(1), Tuple[Val(3.0), Val(4.0)], Stop(2)] + ) + assert a.C == typing.List[typing.Union[Val, Stop]] + assert a.A is Stream[Tuple[int]] + assert a.T is Tuple[int] + assert a.A().T is Tuple[int] + print(state) + +def test_buffer_in_stream(): + state = State() + with state: + a = Buffer[int]().const(np.ndarray([1])) + b = Buffer[int]().const(np.ndarray([2])) + c = Stream[Buffer[int]]().const( + [Val(a), Stop(1), Val(b), Stop(2)] + ) + assert c.C == typing.List[typing.Union[Val, Stop]] + assert c.A is Stream[Buffer[int]] + assert c.VT is Buffer[int] + assert c.A().VT is Buffer[int] + print(state) \ No newline at end of file From 036470511bb64bd8c7c976dbdc7b698d37188df7 Mon Sep 17 00:00:00 2001 From: charlieli13 Date: Tue, 28 May 2024 15:01:46 -0700 Subject: [PATCH 2/7] fixed bugs in stream and buffer types --- src/argon/errors.py | 2 +- src/argon/types/buffer.py | 28 ++++++---- src/argon/types/stream.py | 16 +++--- step/errors | 0 tests/test_step.py | 104 ++++++++++++++++++++------------------ 5 files changed, 85 insertions(+), 65 deletions(-) create mode 100644 step/errors diff --git a/src/argon/errors.py b/src/argon/errors.py index 2398206..c8f0996 100644 --- a/src/argon/errors.py +++ b/src/argon/errors.py @@ -3,4 +3,4 @@ class ArgonError(Exception): class StagingError(ArgonError): - pass + pass \ No newline at end of file diff --git a/src/argon/types/buffer.py b/src/argon/types/buffer.py index 2ad33ec..9cddbba 100644 --- a/src/argon/types/buffer.py +++ b/src/argon/types/buffer.py @@ -1,20 +1,28 @@ from pydantic.dataclasses import dataclass +from pydantic import ConfigDict from typing import Union, Tuple, override, List, TypeVar from argon.ref import Ref +from argon.types.stream import gen_rank import numpy as np +import numpy.typing as npt -T = TypeVar("T", int, float, Tuple[int], Tuple[float]) -@dataclass -class Element: - value: T - - def __str__(self) -> str: - return str(self.value) +T = TypeVar("T") +@dataclass(config=ConfigDict(arbitrary_types_allowed=True)) +class Ndarray[T]: + value: np.ndarray[T] + + BT = TypeVar("BT") +BRK = TypeVar("BRK") -class Buffer[BT](Ref[np.ndarray, "Buffer[BT]"]): +class Buffer[BT,BRK](Ref[Ndarray[BT], "Buffer[BT,BRK]"]): @override - def fresh(self) -> "Buffer[BT]": - return Buffer[self.BT]() \ No newline at end of file + def fresh(self) -> "Buffer[BT,BRK]": + return Buffer[self.BT, self.BRK]() + + @override + def const(self, c: Ndarray[BT]) -> "Buffer[BT,BRK]": + assert type(gen_rank(c.value.shape)) == type(self.BRK) + return super().const(c) \ No newline at end of file diff --git a/src/argon/types/stream.py b/src/argon/types/stream.py index 696b834..4b5b1a3 100644 --- a/src/argon/types/stream.py +++ b/src/argon/types/stream.py @@ -1,7 +1,6 @@ from pydantic.dataclasses import dataclass from typing import Union, Tuple, override, List, TypeVar from argon.ref import Ref -from argon.types.buffer import Element, Buffer @dataclass class Stop: @@ -10,17 +9,22 @@ class Stop: def __str__(self) -> str: return f"S{self.level}" -T = TypeVar("T", Element, Buffer, Tuple[Buffer]) +T = TypeVar("T") @dataclass -class Val: +class Val[T]: value: T def __str__(self) -> str: return str(self.value) VT = TypeVar("VT") -class Stream[VT](Ref[List[Union[Val, Stop]], "Stream[VT]"]): +RK = TypeVar("RK") +class Stream[VT,RK](Ref[List[Union[Val[VT], Stop]], "Stream[VT,RK]"]): @override - def fresh(self) -> "Stream[VT]": - return Stream[self.VT]() \ No newline at end of file + def fresh(self) -> "Stream[VT,RK]": + return Stream[self.VT, self.RK]() + +def gen_rank(rank: int): + rank_name = 'R'+str(rank) + return type(rank_name, (), dict(rank=rank)) \ No newline at end of file diff --git a/step/errors b/step/errors new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_step.py b/tests/test_step.py index 0a10cb5..0814111 100644 --- a/tests/test_step.py +++ b/tests/test_step.py @@ -1,67 +1,75 @@ from argon.state import State -from argon.types.stream import Stop, Val, Stream -from argon.types.buffer import Buffer +from argon.types.stream import Stop, Val, Stream, gen_rank +from argon.types.buffer import Ndarray, Buffer from typing import Union, Tuple, override, List, TypeVar import numpy as np import typing + # Buffer tests def test_buffer(): state = State() with state: - a = Buffer[int]().const(np.ndarray([2])) - assert a.C == np.ndarray - assert a.A is Buffer[int] + R1 = gen_rank(1) + a = Buffer[int,R1]().const(Ndarray[int](np.ndarray([2]))) + assert a.C == Ndarray[int] + assert a.A is Buffer[int, R1] assert a.BT is int assert a.A().BT is int - + + # Stream tests def test_int(): state = State() with state: - a = Stream[int]().const( + R1 = gen_rank(1) + a = Stream[float, R1]().const( [Val(1.0), Val(2.0), Stop(1), Val(3.0), Val(4.0), Stop(2)] ) - assert a.C == typing.List[typing.Union[Val, Stop]] - assert a.A is Stream[int] - assert a.T is int - assert a.A().T is int + print(a.C) + assert a.C == typing.List[typing.Union[Val[float], Stop]] + assert a.A is Stream[float, R1] + assert a.VT is float + assert a.RK is R1 + assert a.A().VT is float + assert a.A().RK is R1 print(state) -def test_float(): - state = State() - with state: - a = Stream[float]().const( - [Val(1.5), Val(2.5), Stop(1), Val(3.5), Val(4.5), Stop(2)] - ) - assert a.C == typing.List[typing.Union[Val, Stop]] - assert a.A is Stream[float] - assert a.T is float - assert a.A().T is float - print(state) - -def test_tuple(): - state = State() - with state: - a = Stream[Tuple[int]]().const( - [Tuple[Val(1.0), Val(2.0)], Stop(1), Tuple[Val(3.0), Val(4.0)], Stop(2)] - ) - assert a.C == typing.List[typing.Union[Val, Stop]] - assert a.A is Stream[Tuple[int]] - assert a.T is Tuple[int] - assert a.A().T is Tuple[int] - print(state) - -def test_buffer_in_stream(): - state = State() - with state: - a = Buffer[int]().const(np.ndarray([1])) - b = Buffer[int]().const(np.ndarray([2])) - c = Stream[Buffer[int]]().const( - [Val(a), Stop(1), Val(b), Stop(2)] - ) - assert c.C == typing.List[typing.Union[Val, Stop]] - assert c.A is Stream[Buffer[int]] - assert c.VT is Buffer[int] - assert c.A().VT is Buffer[int] - print(state) \ No newline at end of file + +# def test_float(): +# state = State() +# with state: +# a = Stream[float]().const( +# [Val(1.5), Val(2.5), Stop(1), Val(3.5), Val(4.5), Stop(2)] +# ) +# assert a.C == typing.List[typing.Union[Val, Stop]] +# assert a.A is Stream[float] +# assert a.T is float +# assert a.A().T is float +# print(state) + +# def test_tuple(): +# state = State() +# with state: +# a = Stream[Tuple[int]]().const( +# [Tuple[Val(1.0), Val(2.0)], Stop(1), Tuple[Val(3.0), Val(4.0)], Stop(2)] +# ) +# assert a.C == typing.List[typing.Union[Val, Stop]] +# assert a.A is Stream[Tuple[int]] +# assert a.T is Tuple[int] +# assert a.A().T is Tuple[int] +# print(state) + +# def test_buffer_in_stream(): +# state = State() +# with state: +# a = Buffer[int]().const(np.ndarray([1])) +# b = Buffer[int]().const(np.ndarray([2])) +# c = Stream[Buffer[int]]().const( +# [Val(a), Stop(1), Val(b), Stop(2)] +# ) +# assert c.C == typing.List[typing.Union[Val, Stop]] +# assert c.A is Stream[Buffer[int]] +# assert c.VT is Buffer[int] +# assert c.A().VT is Buffer[int] +# print(state) From 4327772963227c0ea57f1d55da30e04d246b3261 Mon Sep 17 00:00:00 2001 From: charlieli13 Date: Tue, 4 Jun 2024 13:41:56 -0700 Subject: [PATCH 3/7] Support cases where generic aliases are type variables --- src/argon/base.py | 63 ++++++++++++++++++----------------------------- 1 file changed, 24 insertions(+), 39 deletions(-) diff --git a/src/argon/base.py b/src/argon/base.py index 158186c..9cf3615 100644 --- a/src/argon/base.py +++ b/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( @@ -48,24 +46,6 @@ def __init_subclass__(cls) -> None: # TODO: Make sure that this is actually correct super_init = super().__init_subclass__() - # To handle generic type parameters, we should look them up dynamically at runtime - for ind, tparam in enumerate(cls.__type_params__): - param_name = tparam.__name__ - tparam_set.add(param_name) - - # print(f"setting accessor_tparam for {param_name}") - - def accessor_tparam(self, ind=ind): - # breakpoint() - if not hasattr(self, "__orig_class__"): - raise TypeError( - f"Cannot access type parameter {param_name} of {self.__class__}." - ) - return self.__orig_class__.__args__[ind] - - accessor_tparam.__name__ = param_name - setattr(cls, param_name, property(fget=accessor_tparam)) - # However, if the type parameter hole is filled, we should not use the old accessor anymore. # For example: # class Parent[T]: pass @@ -92,17 +72,19 @@ def accessor_parent_tparam(self, arg=arg): # type: ignore -- PyRight and other if isinstance(retval, typing._GenericAlias): # type: ignore -- We don't have a great alternative way for checking if an object is a GenericAlias aug_ns = {} for key in tparam_set: - if isinstance( - globalns[key], typing.TypeVar - ) and isinstance(localns[key], typing.TypeVar): - # Resolve the type parameters in this class that hasn't been resolved yet - aug_ns[key] = getattr(self, key) + aug_ns[key] = getattr(self, key) # augment the namespace - globalns.update(aug_ns) - localns.update(aug_ns) + temp_globalns = {} + temp_localns = {} + temp_globalns.update(globalns) + temp_globalns.update(aug_ns) + temp_localns.update(localns) + temp_localns.update(aug_ns) - return arg._evaluate(globalns, localns, frozenset()) + return arg._evaluate( + temp_globalns, temp_localns, frozenset() + ) if isinstance(retval, typing.TypeVar): return getattr(self, retval.__name__) @@ -119,21 +101,26 @@ 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: - if isinstance( - globalns[key], typing.TypeVar - ) and isinstance(localns[key], typing.TypeVar): - # Resolve the type parameters in this class that hasn't been resolved yet - aug_ns[key] = getattr(self, key) + aug_ns[key] = getattr(self, key) + print(aug_ns[key]) # augment the namespace - globalns.update(aug_ns) - localns.update(aug_ns) + temp_globalns = {} + temp_localns = {} + temp_globalns.update(globalns) + temp_globalns.update(aug_ns) + temp_localns.update(localns) + temp_localns.update(aug_ns) # recursively resolve the GenericAlias arg_list = [] + print(typing.get_args(arg)) for arg_i in typing.get_args(arg): - arg_list.append(r_resolve(globalns, localns, arg_i)) + arg_list.append( + r_resolve(temp_globalns, temp_localns, arg_i) + ) return typing._GenericAlias( # type: ignore -- We don't have a great alternative way for generating an object that is a GenericAlias typing.get_origin(arg), tuple(arg_list) ) @@ -150,8 +137,6 @@ def accessor_parent_tparam(self, arg=arg): # type: ignore -- PyRight and other param_name = tparam.__name__ tparam_set.add(param_name) - # print(f"setting accessor_tparam for {param_name}") - def accessor_override(self, ind=ind): if not hasattr(self, "__orig_class__"): raise TypeError( From e39db11e86d2d6de661da0d05d0166a5e1484ae7 Mon Sep 17 00:00:00 2001 From: charlieli13 Date: Tue, 4 Jun 2024 13:42:35 -0700 Subject: [PATCH 4/7] implemented buffer and stream types and refactored step types --- src/argon/types/stream.py | 30 ------- {src/argon => step}/types/buffer.py | 16 ++-- step/types/stream.py | 37 +++++++++ tests/test_step.py | 124 +++++++++++++++++----------- 4 files changed, 121 insertions(+), 86 deletions(-) delete mode 100644 src/argon/types/stream.py rename {src/argon => step}/types/buffer.py (64%) create mode 100644 step/types/stream.py diff --git a/src/argon/types/stream.py b/src/argon/types/stream.py deleted file mode 100644 index 4b5b1a3..0000000 --- a/src/argon/types/stream.py +++ /dev/null @@ -1,30 +0,0 @@ -from pydantic.dataclasses import dataclass -from typing import Union, Tuple, override, List, TypeVar -from argon.ref import Ref - -@dataclass -class Stop: - level: int - - def __str__(self) -> str: - return f"S{self.level}" - -T = TypeVar("T") -@dataclass -class Val[T]: - value: T - - def __str__(self) -> str: - return str(self.value) - -VT = TypeVar("VT") -RK = TypeVar("RK") -class Stream[VT,RK](Ref[List[Union[Val[VT], Stop]], "Stream[VT,RK]"]): - - @override - def fresh(self) -> "Stream[VT,RK]": - return Stream[self.VT, self.RK]() - -def gen_rank(rank: int): - rank_name = 'R'+str(rank) - return type(rank_name, (), dict(rank=rank)) \ No newline at end of file diff --git a/src/argon/types/buffer.py b/step/types/buffer.py similarity index 64% rename from src/argon/types/buffer.py rename to step/types/buffer.py index 9cddbba..fc7a5c2 100644 --- a/src/argon/types/buffer.py +++ b/step/types/buffer.py @@ -1,17 +1,21 @@ from pydantic.dataclasses import dataclass from pydantic import ConfigDict -from typing import Union, Tuple, override, List, TypeVar +from typing import Union, Tuple, override, List, TypeVar, Any from argon.ref import Ref -from argon.types.stream import gen_rank import numpy as np -import numpy.typing as npt T = TypeVar("T") @dataclass(config=ConfigDict(arbitrary_types_allowed=True)) class Ndarray[T]: value: np.ndarray[T] - + + +class RankGen: + generated_ranks = {"1": type("R1", (), {}), "2": type("R2", (), {})} + + def get_rank(self, c: int) -> Any: + return self.generated_ranks[str(c)] BT = TypeVar("BT") BRK = TypeVar("BRK") @@ -21,8 +25,8 @@ class Buffer[BT,BRK](Ref[Ndarray[BT], "Buffer[BT,BRK]"]): @override def fresh(self) -> "Buffer[BT,BRK]": return Buffer[self.BT, self.BRK]() - + @override def const(self, c: Ndarray[BT]) -> "Buffer[BT,BRK]": - assert type(gen_rank(c.value.shape)) == type(self.BRK) + assert RankGen().get_rank(c.value.ndim) == self.BRK return super().const(c) \ No newline at end of file diff --git a/step/types/stream.py b/step/types/stream.py new file mode 100644 index 0000000..7e87075 --- /dev/null +++ b/step/types/stream.py @@ -0,0 +1,37 @@ +from pydantic.dataclasses import dataclass +from typing import Union, Tuple, override, List, TypeVar +from argon.ref import Ref +from argon.state import stage +from argon.srcctx import SrcCtx + + +@dataclass +class Stop: + level: int + + def __str__(self) -> str: + return f"S{self.level}" + +VT = TypeVar("VT") + +@dataclass +class Val[VT]: + value: VT + + def __str__(self) -> str: + return str(self.value) + +ST = TypeVar("ST") +SRK = TypeVar("SRK") +B = TypeVar("B") + +class Stream[ST,SRK](Ref[List[Union[Val[ST], Stop]], "Stream[ST,SRK]"]): + + @override + def fresh(self) -> "Stream[ST,SRK]": + return Stream[self.ST, self.SRK]() + + def zip(self, other: "Stream[B,SRK]") -> "Stream[(ST,B),SRK]": + import step.ops.zip as zip + + return stage(zip.Zip[ST,B,SRK](self, other), ctx=SrcCtx.new(2)) \ No newline at end of file diff --git a/tests/test_step.py b/tests/test_step.py index 0814111..2aaebbb 100644 --- a/tests/test_step.py +++ b/tests/test_step.py @@ -1,16 +1,31 @@ from argon.state import State -from argon.types.stream import Stop, Val, Stream, gen_rank -from argon.types.buffer import Ndarray, Buffer -from typing import Union, Tuple, override, List, TypeVar +from step.types.stream import Stop, Val, Stream +from step.types.buffer import RankGen, Ndarray, Buffer +from typing import Tuple, List, Union import numpy as np import typing - +# Op tests +# def test_zip(): +# state = State() +# with state: +# R1 = RankGen().get_rank(1) +# a = Stream[int,R1]().const( +# [Val(1), Val(2), Stop(1), Val(3), Val(4), Stop(2)] +# ) +# b = Stream[int,R1]().const( +# [Val(1), Val(2), Stop(1), Val(3), Val(4), Stop(2)] +# ) +# c = a.zip(b) +# print(c) + + + # Buffer tests def test_buffer(): state = State() with state: - R1 = gen_rank(1) + R1 = RankGen().get_rank(1) a = Buffer[int,R1]().const(Ndarray[int](np.ndarray([2]))) assert a.C == Ndarray[int] assert a.A is Buffer[int, R1] @@ -22,54 +37,63 @@ def test_buffer(): def test_int(): state = State() with state: - R1 = gen_rank(1) - a = Stream[float, R1]().const( - [Val(1.0), Val(2.0), Stop(1), Val(3.0), Val(4.0), Stop(2)] + R1 = RankGen().get_rank(1) + a = Stream[int,R1]().const( + [Val(1), Val(2), Stop(1), Val(3), Val(4), Stop(2)] ) - print(a.C) - assert a.C == typing.List[typing.Union[Val[float], Stop]] - assert a.A is Stream[float, R1] - assert a.VT is float - assert a.RK is R1 - assert a.A().VT is float - assert a.A().RK is R1 + assert a.C == typing.List[typing.Union[Val[int], Stop]] + assert a.A is Stream[int,R1] + assert a.ST is int + assert a.SRK is R1 + assert a.A().ST is int + assert a.A().SRK is R1 print(state) -# def test_float(): -# state = State() -# with state: -# a = Stream[float]().const( -# [Val(1.5), Val(2.5), Stop(1), Val(3.5), Val(4.5), Stop(2)] -# ) -# assert a.C == typing.List[typing.Union[Val, Stop]] -# assert a.A is Stream[float] -# assert a.T is float -# assert a.A().T is float -# print(state) +def test_float(): + state = State() + with state: + R1 = RankGen().get_rank(1) + a = Stream[float,R1]().const( + [Val(1.5), Val(2.5), Stop(1), Val(3.5), Val(4.5), Stop(2)] + ) + assert a.C == typing.List[typing.Union[Val[float], Stop]] + assert a.A is Stream[float,R1] + assert a.ST is float + assert a.SRK is R1 + assert a.A().ST is float + assert a.A().SRK is R1 + print(state) -# def test_tuple(): -# state = State() -# with state: -# a = Stream[Tuple[int]]().const( -# [Tuple[Val(1.0), Val(2.0)], Stop(1), Tuple[Val(3.0), Val(4.0)], Stop(2)] -# ) -# assert a.C == typing.List[typing.Union[Val, Stop]] -# assert a.A is Stream[Tuple[int]] -# assert a.T is Tuple[int] -# assert a.A().T is Tuple[int] -# print(state) +def test_tuple(): + state = State() + with state: + R1 = RankGen().get_rank(1) + a = Stream[Tuple[int],R1]().const( + [Val[Tuple[1,2]], Stop(1), Val[Tuple[3,4]], Stop(2)] + ) + assert a.C == typing.List[typing.Union[Val[Tuple[int]], Stop]] + assert a.A is Stream[Tuple[int],R1] + assert a.ST is Tuple[int] + assert a.SRK is R1 + assert a.A().ST is Tuple[int] + assert a.A().SRK is R1 + print(state) -# def test_buffer_in_stream(): -# state = State() -# with state: -# a = Buffer[int]().const(np.ndarray([1])) -# b = Buffer[int]().const(np.ndarray([2])) -# c = Stream[Buffer[int]]().const( -# [Val(a), Stop(1), Val(b), Stop(2)] -# ) -# assert c.C == typing.List[typing.Union[Val, Stop]] -# assert c.A is Stream[Buffer[int]] -# assert c.VT is Buffer[int] -# assert c.A().VT is Buffer[int] -# print(state) +def test_buffer_in_stream(): + state = State() + with state: + R1 = RankGen().get_rank(1) + a = Buffer[int,R1]().const(Ndarray[int](np.ndarray([1]))) + b = Buffer[int,R1]().const(Ndarray[int](np.ndarray([2]))) + R2 = RankGen().get_rank(2) + c = Stream[Buffer[int,R1],R2]().const( + [Val(a), Stop(1), Val(b), Stop(2)] + ) + assert c.C == typing.List[typing.Union[Val[Buffer[int,R1]], Stop]] + assert c.A is Stream[Buffer[int,R1],R2] + assert c.ST is Buffer[int,R1] + assert c.SRK is R2 + assert c.A().ST is Buffer[int,R1] + assert c.A().SRK is R2 + print(state) \ No newline at end of file From 6936f26dac861745158c9675fe45f650e599b40e Mon Sep 17 00:00:00 2001 From: charlieli13 Date: Tue, 4 Jun 2024 14:43:15 -0700 Subject: [PATCH 5/7] implemented RankGen class and refactored into separate file --- step/types/rankgen.py | 10 ++++++++++ 1 file changed, 10 insertions(+) create mode 100644 step/types/rankgen.py diff --git a/step/types/rankgen.py b/step/types/rankgen.py new file mode 100644 index 0000000..96205fe --- /dev/null +++ b/step/types/rankgen.py @@ -0,0 +1,10 @@ +from typing import Any + + +class RankGen: + generated_ranks = {} + + def get_rank(self, c: int) -> Any: + if str(c) not in self.generated_ranks: + self.generated_ranks[str(c)] = type("R" + str(c), (), {}) + return self.generated_ranks[str(c)] From ac0a44115826c36f3c3fecd890199a21ad98b7f0 Mon Sep 17 00:00:00 2001 From: charlieli13 Date: Tue, 4 Jun 2024 14:44:23 -0700 Subject: [PATCH 6/7] commenting out code not being used --- step/types/buffer.py | 51 +++++++++++++++------------------ step/types/stream.py | 6 ++-- tests/test_step.py | 67 ++++++++++++++++++++++---------------------- 3 files changed, 60 insertions(+), 64 deletions(-) diff --git a/step/types/buffer.py b/step/types/buffer.py index fc7a5c2..b30f3d9 100644 --- a/step/types/buffer.py +++ b/step/types/buffer.py @@ -1,32 +1,27 @@ -from pydantic.dataclasses import dataclass -from pydantic import ConfigDict -from typing import Union, Tuple, override, List, TypeVar, Any -from argon.ref import Ref -import numpy as np - -T = TypeVar("T") - -@dataclass(config=ConfigDict(arbitrary_types_allowed=True)) -class Ndarray[T]: - value: np.ndarray[T] - - -class RankGen: - generated_ranks = {"1": type("R1", (), {}), "2": type("R2", (), {})} - - def get_rank(self, c: int) -> Any: - return self.generated_ranks[str(c)] +# from pydantic.dataclasses import dataclass +# from pydantic import ConfigDict +# from typing import Union, Tuple, override, List, TypeVar, Any +# from argon.ref import Ref +# from rankgen import RankGen +# import numpy as np +# import numpy.typing as npt + +# T = TypeVar("T", bound=np.generic, covariant=True) + +# @dataclass(config=ConfigDict(arbitrary_types_allowed=True)) +# class Ndarray[T]: +# value: npt.NDArray[T] -BT = TypeVar("BT") -BRK = TypeVar("BRK") +# BT = TypeVar("BT") +# BRK = TypeVar("BRK") -class Buffer[BT,BRK](Ref[Ndarray[BT], "Buffer[BT,BRK]"]): +# class Buffer[BT,BRK](Ref[Ndarray[BT], "Buffer[BT,BRK]"]): - @override - def fresh(self) -> "Buffer[BT,BRK]": - return Buffer[self.BT, self.BRK]() +# @override +# def fresh(self) -> "Buffer[BT,BRK]": +# return Buffer[self.BT, self.BRK]() - @override - def const(self, c: Ndarray[BT]) -> "Buffer[BT,BRK]": - assert RankGen().get_rank(c.value.ndim) == self.BRK - return super().const(c) \ No newline at end of file +# @override +# def const(self, c: Ndarray[BT]) -> "Buffer[BT,BRK]": +# assert RankGen().get_rank(c.value.ndim) == self.BRK +# return super().const(c) \ No newline at end of file diff --git a/step/types/stream.py b/step/types/stream.py index 7e87075..0b5674e 100644 --- a/step/types/stream.py +++ b/step/types/stream.py @@ -31,7 +31,7 @@ class Stream[ST,SRK](Ref[List[Union[Val[ST], Stop]], "Stream[ST,SRK]"]): def fresh(self) -> "Stream[ST,SRK]": return Stream[self.ST, self.SRK]() - def zip(self, other: "Stream[B,SRK]") -> "Stream[(ST,B),SRK]": - import step.ops.zip as zip + # def zip(self, other: "Stream[B,SRK]") -> "Stream[(ST,B),SRK]": + # import step.ops.zip as zip - return stage(zip.Zip[ST,B,SRK](self, other), ctx=SrcCtx.new(2)) \ No newline at end of file + # return stage(zip.Zip[ST,B,SRK](self, other), ctx=SrcCtx.new(2)) \ No newline at end of file diff --git a/tests/test_step.py b/tests/test_step.py index 2aaebbb..0dd8c3c 100644 --- a/tests/test_step.py +++ b/tests/test_step.py @@ -1,6 +1,7 @@ from argon.state import State from step.types.stream import Stop, Val, Stream -from step.types.buffer import RankGen, Ndarray, Buffer +# from step.types.buffer import Ndarray, Buffer +from step.types.rankgen import RankGen from typing import Tuple, List, Union import numpy as np import typing @@ -22,15 +23,15 @@ # Buffer tests -def test_buffer(): - state = State() - with state: - R1 = RankGen().get_rank(1) - a = Buffer[int,R1]().const(Ndarray[int](np.ndarray([2]))) - assert a.C == Ndarray[int] - assert a.A is Buffer[int, R1] - assert a.BT is int - assert a.A().BT is int +# def test_buffer(): +# state = State() +# with state: +# R1 = RankGen().get_rank(1) +# a = Buffer[int,R1]().const(Ndarray[int](np.ndarray([2]))) +# assert a.C == Ndarray[int] +# assert a.A is Buffer[int, R1] +# assert a.BT is int +# assert a.A().BT is int # Stream tests @@ -69,31 +70,31 @@ def test_tuple(): state = State() with state: R1 = RankGen().get_rank(1) - a = Stream[Tuple[int],R1]().const( - [Val[Tuple[1,2]], Stop(1), Val[Tuple[3,4]], Stop(2)] + a = Stream[Tuple[int,int],R1]().const( + [Val((1,2)), Stop(1), Val((3,4)), Stop(2)] ) - assert a.C == typing.List[typing.Union[Val[Tuple[int]], Stop]] - assert a.A is Stream[Tuple[int],R1] - assert a.ST is Tuple[int] + assert a.C == typing.List[typing.Union[Val[Tuple[int,int]], Stop]] + assert a.A is Stream[Tuple[int,int],R1] + assert a.ST is Tuple[int,int] assert a.SRK is R1 - assert a.A().ST is Tuple[int] + assert a.A().ST is Tuple[int,int] assert a.A().SRK is R1 print(state) -def test_buffer_in_stream(): - state = State() - with state: - R1 = RankGen().get_rank(1) - a = Buffer[int,R1]().const(Ndarray[int](np.ndarray([1]))) - b = Buffer[int,R1]().const(Ndarray[int](np.ndarray([2]))) - R2 = RankGen().get_rank(2) - c = Stream[Buffer[int,R1],R2]().const( - [Val(a), Stop(1), Val(b), Stop(2)] - ) - assert c.C == typing.List[typing.Union[Val[Buffer[int,R1]], Stop]] - assert c.A is Stream[Buffer[int,R1],R2] - assert c.ST is Buffer[int,R1] - assert c.SRK is R2 - assert c.A().ST is Buffer[int,R1] - assert c.A().SRK is R2 - print(state) \ No newline at end of file +# def test_buffer_in_stream(): +# state = State() +# with state: +# R1 = RankGen().get_rank(1) +# a = Buffer[int,R1]().const(Ndarray[int](np.ndarray([1]))) +# b = Buffer[int,R1]().const(Ndarray[int](np.ndarray([2]))) +# R2 = RankGen().get_rank(2) +# c = Stream[Buffer[int,R1],R2]().const( +# [Val(a), Stop(1), Val(b), Stop(2)] +# ) +# assert c.C == typing.List[typing.Union[Val[Buffer[int,R1]], Stop]] +# assert c.A is Stream[Buffer[int,R1],R2] +# assert c.ST is Buffer[int,R1] +# assert c.SRK is R2 +# assert c.A().ST is Buffer[int,R1] +# assert c.A().SRK is R2 +# print(state) \ No newline at end of file From 03e92210ba4c108966d8b9605ef9a4e667b238f2 Mon Sep 17 00:00:00 2001 From: charlieli13 Date: Mon, 15 Jul 2024 17:13:16 -0700 Subject: [PATCH 7/7] added op types and corresponding stream functions --- step/ops/accum.py | 26 +++++ step/ops/bufferize.py | 25 +++++ step/ops/enumerate.py | 23 +++++ step/ops/flatmap.py | 43 +++++++++ step/ops/flatten.py | 24 +++++ step/ops/map.py | 37 +++++++ step/ops/partition.py | 26 +++++ step/ops/promote.py | 24 +++++ step/ops/repeat.py | 25 +++++ step/ops/reshape.py | 24 +++++ step/ops/window.py | 24 +++++ step/ops/zip.py | 24 +++++ step/types/buffer.py | 44 ++++----- step/types/rankgen.py | 5 +- step/types/stream.py | 94 ++++++++++++++++-- tests/test_step.py | 220 +++++++++++++++++++++++++++++++++--------- 16 files changed, 613 insertions(+), 75 deletions(-) create mode 100644 step/ops/accum.py create mode 100644 step/ops/bufferize.py create mode 100644 step/ops/enumerate.py create mode 100644 step/ops/flatmap.py create mode 100644 step/ops/flatten.py create mode 100644 step/ops/map.py create mode 100644 step/ops/partition.py create mode 100644 step/ops/promote.py create mode 100644 step/ops/repeat.py create mode 100644 step/ops/reshape.py create mode 100644 step/ops/window.py create mode 100644 step/ops/zip.py diff --git a/step/ops/accum.py b/step/ops/accum.py new file mode 100644 index 0000000..305d0bf --- /dev/null +++ b/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 \ No newline at end of file diff --git a/step/ops/bufferize.py b/step/ops/bufferize.py new file mode 100644 index 0000000..36774fe --- /dev/null +++ b/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 \ No newline at end of file diff --git a/step/ops/enumerate.py b/step/ops/enumerate.py new file mode 100644 index 0000000..63d5093 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/step/ops/flatmap.py b/step/ops/flatmap.py new file mode 100644 index 0000000..2dd0520 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/step/ops/flatten.py b/step/ops/flatten.py new file mode 100644 index 0000000..919bafa --- /dev/null +++ b/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 \ No newline at end of file diff --git a/step/ops/map.py b/step/ops/map.py new file mode 100644 index 0000000..4f438b3 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/step/ops/partition.py b/step/ops/partition.py new file mode 100644 index 0000000..b49b98b --- /dev/null +++ b/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 \ No newline at end of file diff --git a/step/ops/promote.py b/step/ops/promote.py new file mode 100644 index 0000000..ddd9299 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/step/ops/repeat.py b/step/ops/repeat.py new file mode 100644 index 0000000..4f27bf5 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/step/ops/reshape.py b/step/ops/reshape.py new file mode 100644 index 0000000..c4af540 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/step/ops/window.py b/step/ops/window.py new file mode 100644 index 0000000..095e4f6 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/step/ops/zip.py b/step/ops/zip.py new file mode 100644 index 0000000..0f1cac8 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/step/types/buffer.py b/step/types/buffer.py index b30f3d9..b6eaf28 100644 --- a/step/types/buffer.py +++ b/step/types/buffer.py @@ -1,27 +1,27 @@ -# from pydantic.dataclasses import dataclass -# from pydantic import ConfigDict -# from typing import Union, Tuple, override, List, TypeVar, Any -# from argon.ref import Ref -# from rankgen import RankGen -# import numpy as np -# import numpy.typing as npt +from pydantic.dataclasses import dataclass +from pydantic import ConfigDict +from typing import Union, Tuple, override, List, TypeVar, Any, Generic +from argon.ref import Ref +from .rankgen import RankGen +import numpy as np +from numpy.typing import NDArray -# T = TypeVar("T", bound=np.generic, covariant=True) +T = TypeVar("T") -# @dataclass(config=ConfigDict(arbitrary_types_allowed=True)) -# class Ndarray[T]: -# value: npt.NDArray[T] - -# BT = TypeVar("BT") -# BRK = TypeVar("BRK") +@dataclass(config=ConfigDict(arbitrary_types_allowed=True)) +class Ndarray[T]: + value: np.ndarray[T] -# class Buffer[BT,BRK](Ref[Ndarray[BT], "Buffer[BT,BRK]"]): +BT = TypeVar("BT") +BRK = TypeVar("BRK") -# @override -# def fresh(self) -> "Buffer[BT,BRK]": -# return Buffer[self.BT, self.BRK]() +class Buffer[BT,BRK](Ref[Ndarray[BT], "Buffer[BT,BRK]"]): -# @override -# def const(self, c: Ndarray[BT]) -> "Buffer[BT,BRK]": -# assert RankGen().get_rank(c.value.ndim) == self.BRK -# return super().const(c) \ No newline at end of file + @override + def fresh(self) -> "Buffer[BT,BRK]": + return Buffer[self.BT, self.BRK]() + + @override + def const(self, c: Ndarray[BT]) -> "Buffer[BT,BRK]": + assert RankGen().get_rank(c.value.ndim) == self.BRK + return super().const(c) \ No newline at end of file diff --git a/step/types/rankgen.py b/step/types/rankgen.py index 96205fe..3756ea4 100644 --- a/step/types/rankgen.py +++ b/step/types/rankgen.py @@ -1,5 +1,4 @@ -from typing import Any - +from typing import Any, Tuple class RankGen: generated_ranks = {} @@ -7,4 +6,4 @@ class RankGen: def get_rank(self, c: int) -> Any: if str(c) not in self.generated_ranks: self.generated_ranks[str(c)] = type("R" + str(c), (), {}) - return self.generated_ranks[str(c)] + return self.generated_ranks[str(c)] \ No newline at end of file diff --git a/step/types/stream.py b/step/types/stream.py index 0b5674e..6263002 100644 --- a/step/types/stream.py +++ b/step/types/stream.py @@ -1,9 +1,9 @@ from pydantic.dataclasses import dataclass -from typing import Union, Tuple, override, List, TypeVar +from typing import Union, Tuple, override, List, TypeVar, Callable, Any from argon.ref import Ref from argon.state import stage from argon.srcctx import SrcCtx - +from .rankgen import RankGen @dataclass class Stop: @@ -20,18 +20,98 @@ class Val[VT]: def __str__(self) -> str: return str(self.value) + +@dataclass +class Index: + value: int + + def __str__(self) -> str: + return str(self.value) + ST = TypeVar("ST") +ST2 = TypeVar("ST2") SRK = TypeVar("SRK") -B = TypeVar("B") +SRK2 = TypeVar("SRK2") + + +class RStream[ST,ST2,SRK,SRK2](Ref[List[Union[Val[ST], Stop]], "RStream[ST,ST2,SRK,SRK2]"]): + @override + def fresh(self) -> "RStream[ST,ST2,SRK,SRK2]": + return RStream[self.ST, self.ST2, self.SRK, self.SRK2]() + + def accum(self, func: Callable[[ST,ST2], SRK]) -> "Stream[ST2,SRK-SRK2]": + from ..ops.accum import Accum + + d = int(self.SRK.__name__[1:]) - int(self.SRK2.__name__[1:]) + assert d >= 0 + return stage(Accum[self.ST, self.ST2, self.SRK, self.SRK2, RankGen().get_rank(d)](self, func), ctx=SrcCtx.new(2)) + + def repeat(self, other: "RStream[ST,ST2,SRK,SRK2]") -> "Stream[ST,SRK+1]": + from ..ops.repeat import Repeat + + return stage(Repeat[self.ST, self.SRK, RankGen().get_rank(int(self.SRK.__name__[1:])+1)](self, other), ctx=SrcCtx.new(2)) + + # def flatmap(self, func: Callable[[A], Stream[B, b]]) -> "Stream[ST,SRK]": + # from ..ops.flatmap import Flatmap + + # return stage(Flatmap[self.SRK, b, self.ST, B](self, func), ctx=SrcCtx.new(2)) + + def partition(self, N: int, other: "RStream[ST,ST2,SRK,SRK2]") -> "Stream[ST,SRK-SRK2+1]": + from ..ops.partition import Partition + + d = RankGen().get_rank(int(self.SRK.__name__[1:]) - int(self.SRK2.__name__[1:]) + 1) + return stage(Partition[self.ST, self.ST2, self.SRK, self.SRK2, d](self, N, other), ctx=SrcCtx.new(2)) -class Stream[ST,SRK](Ref[List[Union[Val[ST], Stop]], "Stream[ST,SRK]"]): +class HStream[ST,ST2,SRK](Ref[List[Union[Val[ST], Stop]], "HStream[ST,ST2,SRK]"]): + @override + def fresh(self) -> "Stream[ST,ST2,SRK]": + return HStream[self.ST, self.ST2, self.SRK]() + + def map(self, func: Callable[[ST], ST2]) -> "HStream[ST2,SRK]": + from ..ops.map import Map + + return stage(Map[self.ST, self.ST2, self.SRK](self, func), ctx=SrcCtx.new(2)) + + def zip(self, other: "HStream[ST,ST2,SRK]") -> "Stream[Tuple[ST, ST2], SRK]": + from ..ops.zip import Zip + + return stage(Zip[self.ST, other.ST2, self.SRK](self, other), ctx=SrcCtx.new(2)) + + +class Stream[ST,SRK](Ref[List[Union[Val[ST], Stop]], "Stream[ST,SRK]"]): @override def fresh(self) -> "Stream[ST,SRK]": return Stream[self.ST, self.SRK]() - # def zip(self, other: "Stream[B,SRK]") -> "Stream[(ST,B),SRK]": - # import step.ops.zip as zip + def bufferize(self, a: int) -> "Stream[Buffer[ST, a], SRK+1]": + from ..ops.bufferize import Bufferize + + return stage(Bufferize[self.ST, RankGen().get_rank(a), self.SRK, RankGen().get_rank(int(self.SRK.__name__[1:])-a)](self), ctx=SrcCtx.new(2)) + + def promote(self, b: int) -> "Stream[ST,SRK+1]": + from ..ops.promote import Promote + + return stage(Promote[self.ST, self.SRK, RankGen().get_rank(int(self.SRK.__name__[1:])+1)](self, b), ctx=SrcCtx.new(2)) + + def reshape(self, L: Tuple[Index, ...], S: Tuple[int, ...]) -> "Stream[ST,SRK+L]": + from ..ops.reshape import Reshape + + return stage(Reshape[self.ST, self.SRK, RankGen().get_rank(int(self.SRK.__name__[1:])+len(L))](self, L, S), ctx=SrcCtx.new(2)) + + def flatten(self, L: Tuple[Index, ...]) -> "Stream[ST,SRK-L]": + from ..ops.flatten import Flatten + + assert len(L) < int(self.SRK.__name__[1:]) - 1 + return stage(Flatten[self.SRK, self.ST, RankGen().get_rank(int(self.SRK.__name__[1:])-len(L))](self, L), ctx=SrcCtx.new(2)) + + def enumerate(self, b: int) -> "Stream[Tuple[ST,Index],SRK]": + from ..ops.enumerate import Enumerate + + return stage(Enumerate[self.ST, self.SRK](self, b), ctx=SrcCtx.new(2)) + + # def window(self, N: Tuple[int], S: Tuple[int]) -> "Stream[ST,SRK]": + # from ..ops.window import Window - # return stage(zip.Zip[ST,B,SRK](self, other), ctx=SrcCtx.new(2)) \ No newline at end of file + # return stage(Window[self.ST, self.SRK, N, S](self, N, S), ctx=SrcCtx.new(2)) \ No newline at end of file diff --git a/tests/test_step.py b/tests/test_step.py index 0dd8c3c..f742baa 100644 --- a/tests/test_step.py +++ b/tests/test_step.py @@ -1,37 +1,171 @@ from argon.state import State -from step.types.stream import Stop, Val, Stream -# from step.types.buffer import Ndarray, Buffer +from step.types.stream import Stop, Val, Index, Stream, RStream, HStream +from step.types.buffer import Ndarray, Buffer from step.types.rankgen import RankGen +# from step.ops.zip import Zip from typing import Tuple, List, Union import numpy as np import typing +# Composite tests + # Op tests -# def test_zip(): -# state = State() -# with state: -# R1 = RankGen().get_rank(1) -# a = Stream[int,R1]().const( -# [Val(1), Val(2), Stop(1), Val(3), Val(4), Stop(2)] -# ) -# b = Stream[int,R1]().const( -# [Val(1), Val(2), Stop(1), Val(3), Val(4), Stop(2)] -# ) -# c = a.zip(b) -# print(c) - - - +def test_map(): + def func(x: int, y: int) -> int: + return x + y + + def func2(x: int, y: float) -> float: + return x + y + + state = State() + with state: + R2 = RankGen().get_rank(2) + a = HStream[int,int,R2]().const( + [Val(1), Val(2), Stop(1), Val(3), Val(4), Stop(2)] + ) + b = a.map(lambda val: func(val, 1)) + print(b.A) + c = HStream[int,float,R2]().const( + [Val(1), Val(2), Stop(1), Val(3), Val(4), Stop(2)] + ) + d = c.map(lambda val: func2(val, 1.0)) + print(d.A) + print(state) + + +def test_accum(): + def func(x: int, y: int) -> int: + return x + y + + state = State() + with state: + R1 = RankGen().get_rank(2) + R2 = RankGen().get_rank(1) + a = RStream[int,int,R1,R2]().const( + [Val(1), Val(2), Stop(1), Val(3), Val(4), Stop(2)] + ) + b = a.accum(lambda val: func(val, 1)) + print(b.A) + print(state) + + +def test_zip(): + state = State() + with state: + R1 = RankGen().get_rank(1) + a = HStream[int,int,R1]().const( + [Val(1), Val(2), Stop(1), Val(3), Val(4), Stop(2)] + ) + b = HStream[int,int,R1]().const( + [Val(1), Val(2), Stop(1), Val(3), Val(4), Stop(2)] + ) + c = a.zip(b) + print(c.A) + d = HStream[int,float,R1]().const( + [Val(1.0), Val(2.0), Stop(1), Val(3.0), Val(4.0), Stop(2)] + ) + e = a.zip(d) + print(e.A) + print(state) + + +def test_repeat(): + state = State() + with state: + R0 = RankGen().get_rank(0) + R1 = RankGen().get_rank(1) + a = RStream[int,int,R1,R0]().const( + [Val(1), Val(2), Stop(1), Val(3), Val(4), Stop(2)] + ) + b = a.repeat(RStream[int,int,R1,R0]()) + print(state) + + +def test_bufferize(): + state = State() + with state: + R1 = RankGen().get_rank(1) + a = Stream[int,R1]().const( + [Val(1), Val(2), Stop(1), Val(3), Val(4), Stop(2)] + ) + b = a.bufferize(1) + print(state) + + +def test_promote(): + state = State() + with state: + R1 = RankGen().get_rank(1) + a = Stream[int,R1]().const( + [Val(1), Val(2), Stop(1), Val(3), Val(4), Stop(2)] + ) + b = a.promote(1) + print(state) + + +def test_reshape(): + state = State() + with state: + R1 = RankGen().get_rank(1) + a = Stream[int,R1]().const( + [Val(1), Val(2), Stop(1), Val(3), Val(4), Stop(2)] + ) + b = a.reshape((Index(1), Index(2)), (1,1)) + print(b.A) + print(b.SRK) + print(state) + + +def test_flatten(): + state = State() + with state: + R3 = RankGen().get_rank(3) + a = Stream[int,R3]().const( + [Val(1), Val(2), Stop(1), Val(3), Val(4), Stop(2)] + ) + b = a.flatten((Index(1),)) + print(b.A) + print(state) + + +def test_enumerate(): + state = State() + with state: + R1 = RankGen().get_rank(1) + a = Stream[int,R1]().const( + [Val(1), Val(2), Stop(1), Val(3), Val(4), Stop(2)] + ) + b = a.enumerate(1) + print(b.A) + print(state) + + +def test_partition(): + state = State() + with state: + R1 = RankGen().get_rank(1) + a = RStream[int,int,R1,R1]().const( + [Val(1), Val(2), Stop(1), Val(3), Val(4), Stop(2)] + ) + b = RStream[int,int,R1,R1]().const( + [Val(1), Val(2), Stop(1), Val(3), Val(4), Stop(2)] + ) + c = a.partition(3, b) + print(c.A) + print(state) + + # Buffer tests -# def test_buffer(): -# state = State() -# with state: -# R1 = RankGen().get_rank(1) -# a = Buffer[int,R1]().const(Ndarray[int](np.ndarray([2]))) -# assert a.C == Ndarray[int] -# assert a.A is Buffer[int, R1] -# assert a.BT is int -# assert a.A().BT is int +def test_buffer(): + state = State() + with state: + R1 = RankGen().get_rank(1) + a = Buffer[int,R1]().const(Ndarray[int](np.array([2]))) + assert a.C == Ndarray[int] + assert a.A is Buffer[int, R1] + assert a.BT is int + assert a.A().BT is int + print(state) # Stream tests @@ -81,20 +215,20 @@ def test_tuple(): assert a.A().SRK is R1 print(state) -# def test_buffer_in_stream(): -# state = State() -# with state: -# R1 = RankGen().get_rank(1) -# a = Buffer[int,R1]().const(Ndarray[int](np.ndarray([1]))) -# b = Buffer[int,R1]().const(Ndarray[int](np.ndarray([2]))) -# R2 = RankGen().get_rank(2) -# c = Stream[Buffer[int,R1],R2]().const( -# [Val(a), Stop(1), Val(b), Stop(2)] -# ) -# assert c.C == typing.List[typing.Union[Val[Buffer[int,R1]], Stop]] -# assert c.A is Stream[Buffer[int,R1],R2] -# assert c.ST is Buffer[int,R1] -# assert c.SRK is R2 -# assert c.A().ST is Buffer[int,R1] -# assert c.A().SRK is R2 -# print(state) \ No newline at end of file +def test_buffer_in_stream(): + state = State() + with state: + R1 = RankGen().get_rank(1) + a = Buffer[int,R1]().const(Ndarray[int](np.ndarray([1]))) + b = Buffer[int,R1]().const(Ndarray[int](np.ndarray([2]))) + R2 = RankGen().get_rank(2) + c = Stream[Buffer[int,R1],R2]().const( + [Val(a), Stop(1), Val(b), Stop(2)] + ) + #assert c.C == typing.List[typing.Union[Val[Buffer[int,R1]], Stop]] + assert c.A is Stream[Buffer[int,R1],R2] + assert c.ST is Buffer[int,R1] + assert c.SRK is R2 + assert c.A().ST is Buffer[int,R1] + assert c.A().SRK is R2 + print(state) \ No newline at end of file