diff --git a/src/argon/base.py b/src/argon/base.py index d0d7734..ed2985d 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( @@ -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) 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/step/errors b/step/errors new file mode 100644 index 0000000..e69de29 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 new file mode 100644 index 0000000..b6eaf28 --- /dev/null +++ b/step/types/buffer.py @@ -0,0 +1,27 @@ +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") + +@dataclass(config=ConfigDict(arbitrary_types_allowed=True)) +class Ndarray[T]: + value: np.ndarray[T] + +BT = TypeVar("BT") +BRK = TypeVar("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 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 new file mode 100644 index 0000000..3756ea4 --- /dev/null +++ b/step/types/rankgen.py @@ -0,0 +1,9 @@ +from typing import Any, Tuple + +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)] \ No newline at end of file diff --git a/step/types/stream.py b/step/types/stream.py new file mode 100644 index 0000000..6263002 --- /dev/null +++ b/step/types/stream.py @@ -0,0 +1,117 @@ +from pydantic.dataclasses import dataclass +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: + 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) + +@dataclass +class Index: + value: int + + def __str__(self) -> str: + return str(self.value) + + +ST = TypeVar("ST") +ST2 = TypeVar("ST2") +SRK = TypeVar("SRK") +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 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 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(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 new file mode 100644 index 0000000..f742baa --- /dev/null +++ b/tests/test_step.py @@ -0,0 +1,234 @@ +from argon.state import State +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_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.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 +def test_int(): + 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)] + ) + 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: + 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: + R1 = RankGen().get_rank(1) + 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,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,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