diff --git a/tripy/pyproject.toml b/tripy/pyproject.toml index 6e6a9c91d..43f6bf6ee 100644 --- a/tripy/pyproject.toml +++ b/tripy/pyproject.toml @@ -8,8 +8,8 @@ requires-python = ">= 3.9" license = {text = "Apache 2.0"} dependencies = [ "tensorrt~=10.0", - "mlir-tensorrt-compiler==0.1.34+cuda12.trt102", - "mlir-tensorrt-runtime==0.1.34+cuda12.trt102", + "mlir-tensorrt-compiler==0.1.36+cuda12.trt102", + "mlir-tensorrt-runtime==0.1.36+cuda12.trt102", "colored==2.2.3", ] diff --git a/tripy/tests/integration/test_iota.py b/tripy/tests/integration/test_iota.py index 39df35787..2df779da2 100644 --- a/tripy/tests/integration/test_iota.py +++ b/tripy/tests/integration/test_iota.py @@ -91,9 +91,7 @@ def test_negative_no_casting(self, dtype): a = tp.ones((2, 2)) out = Iota.build([frontend_utils.tensor_from_shape_like(a.shape)], dim=0, output_rank=2, dtype=dtype) - exception_str = "error: 'tensorrt.linspace' op result #0 must be 0D/1D/2D/3D/4D/5D/6D/7D/8D tensor of 32-bit float or 32-bit signless integer values" - if dtype == tp.bool: - exception_str = "InternalError: failed to run compilation" + exception_str = "InternalError: failed to run compilation" with helper.raises( tp.TripyException, match=exception_str, diff --git a/tripy/tests/integration/test_quantize.py b/tripy/tests/integration/test_quantize.py index b50293869..826db83bf 100644 --- a/tripy/tests/integration/test_quantize.py +++ b/tripy/tests/integration/test_quantize.py @@ -23,6 +23,7 @@ import tripy as tp from tests.helper import raises, TORCH_DTYPES from tests.conftest import skip_if_older_than_sm80, skip_if_older_than_sm89 +import cupy as cp class TestQuantize: @@ -118,4 +119,4 @@ def test_non_constant_scale(self): scale = tp.ones((4,)) quantized = tp.quantize(input, scale, tp.int8, dim=0) - assert bool(tp.all(quantized == tp.ones((4, 4), dtype=tp.int8))) + assert bool(cp.all(cp.from_dlpack(quantized) == cp.ones((4, 4), dtype=cp.int8))) diff --git a/tripy/tripy/frontend/ops/tensor_initializers.py b/tripy/tripy/frontend/ops/tensor_initializers.py index da5b1e221..faa395f7a 100644 --- a/tripy/tripy/frontend/ops/tensor_initializers.py +++ b/tripy/tripy/frontend/ops/tensor_initializers.py @@ -165,7 +165,7 @@ def zeros_like(input: "tripy.Tensor", dtype: Optional[datatype.dtype] = None) -> @constraints.dtypes( constraints={"tensor": "T1", constraints.RETURN_VALUE: "T1"}, variables={ - "T1": ["float32", "float16", "bfloat16", "int32", "bool"], + "T1": ["float32", "float16", "bfloat16", "int32", "int64", "bool"], }, ) def tril(tensor: "tripy.Tensor", diagonal: int = 0) -> "tripy.Tensor": @@ -223,7 +223,7 @@ def tril(tensor: "tripy.Tensor", diagonal: int = 0) -> "tripy.Tensor": @constraints.dtypes( constraints={"tensor": "T1", constraints.RETURN_VALUE: "T1"}, variables={ - "T1": ["float32", "float16", "bfloat16", "int32", "bool"], + "T1": ["float32", "float16", "bfloat16", "int32", "int64", "bool"], }, ) def triu(tensor: "tripy.Tensor", diagonal: int = 0) -> "tripy.Tensor": @@ -281,7 +281,7 @@ def triu(tensor: "tripy.Tensor", diagonal: int = 0) -> "tripy.Tensor": @constraints.dtypes( constraints={"dtype": "T1", constraints.RETURN_VALUE: "T1"}, variables={ - "T1": ["float32", "float16", "bfloat16", "int8", "int32", "bool"], + "T1": ["float32", "float16", "bfloat16", "int8", "int32", "int64", "bool"], }, ) def arange( @@ -346,7 +346,7 @@ def arange( @constraints.dtypes( constraints={"dtype": "T1", constraints.RETURN_VALUE: "T1"}, variables={ - "T1": ["float32", "float16", "bfloat16", "int8", "int32", "bool"], + "T1": ["float32", "float16", "bfloat16", "int8", "int32", "int64", "bool"], }, ) def arange( diff --git a/tripy/tripy/frontend/trace/ops/flip.py b/tripy/tripy/frontend/trace/ops/flip.py index 15e12b2fa..ab9820174 100644 --- a/tripy/tripy/frontend/trace/ops/flip.py +++ b/tripy/tripy/frontend/trace/ops/flip.py @@ -40,7 +40,7 @@ def to_flat_ir(self, inputs, outputs): @constraints.dtypes( constraints={"input": "T1", constraints.RETURN_VALUE: "T1"}, variables={ - "T1": ["float32", "float16", "bfloat16", "int32", "bool"], + "T1": ["float32", "float16", "bfloat16", "int32", "int64", "bool"], }, ) def flip(input: "tripy.Tensor", dims: Optional[Union[int, Sequence[int]]] = None) -> "tripy.Tensor": diff --git a/tripy/tripy/frontend/trace/ops/gather.py b/tripy/tripy/frontend/trace/ops/gather.py index 0df478585..d18294d91 100644 --- a/tripy/tripy/frontend/trace/ops/gather.py +++ b/tripy/tripy/frontend/trace/ops/gather.py @@ -91,7 +91,7 @@ def to_flat_ir(self, inputs, outputs): @constraints.dtypes( constraints={"input": "T1", "index": "T2", constraints.RETURN_VALUE: "T1"}, variables={ - "T1": ["float32", "float16", "bfloat16", "int8", "int32", "bool"], + "T1": ["float32", "float16", "bfloat16", "int4", "int8", "int32", "int64", "bool"], "T2": ["int32"], }, ) diff --git a/tripy/tripy/frontend/trace/ops/iota.py b/tripy/tripy/frontend/trace/ops/iota.py index a252d756c..98933f92a 100644 --- a/tripy/tripy/frontend/trace/ops/iota.py +++ b/tripy/tripy/frontend/trace/ops/iota.py @@ -69,7 +69,7 @@ def iota_impl(shape: "tripy.Tensor", dim: int, dtype: datatype.dtype, output_ran @constraints.dtypes( constraints={"dtype": "T1", constraints.RETURN_VALUE: "T1"}, variables={ - "T1": ["float32", "float16", "bfloat16", "float8", "int4", "int8", "int32", "bool"], + "T1": ["float32", "float16", "bfloat16", "float8", "int4", "int8", "int32", "int64", "bool"], }, ) def iota(shape: ShapeLike, dim: int = 0, dtype: datatype.dtype = datatype.float32) -> "tripy.Tensor": @@ -101,7 +101,7 @@ def iota(shape: ShapeLike, dim: int = 0, dtype: datatype.dtype = datatype.float3 constraints={"input": "T1", "dtype": "T2", constraints.RETURN_VALUE: "T2"}, variables={ "T1": ["float32", "float16", "bfloat16", "float8", "int4", "int8", "int32", "int64", "bool"], - "T2": ["float32", "float16", "bfloat16", "float8", "int4", "int8", "int32", "bool"], + "T2": ["float32", "float16", "bfloat16", "float8", "int4", "int8", "int32", "int64", "bool"], }, ) def iota_like(input: "tripy.Tensor", dim: int = 0, dtype: Optional[datatype.dtype] = None) -> "tripy.Tensor": diff --git a/tripy/tripy/frontend/trace/ops/reduce.py b/tripy/tripy/frontend/trace/ops/reduce.py index e12a8129a..cba001f78 100644 --- a/tripy/tripy/frontend/trace/ops/reduce.py +++ b/tripy/tripy/frontend/trace/ops/reduce.py @@ -135,7 +135,7 @@ def _reduce_impl(input: "tripy.Tensor", kind: Reduce.Kind, dim: Union[int, Seque @export.public_api(document_under="operations/functions") @constraints.dtypes( constraints={"input": "T1", constraints.RETURN_VALUE: "T1"}, - variables={"T1": ["float32", "int32", "float16", "bfloat16"]}, + variables={"T1": ["float32", "int32", "int64", "float16", "bfloat16"]}, ) def sum( input: "tripy.Tensor", dim: Optional[Union[int, Sequence[int]]] = None, keepdim: bool = False @@ -232,7 +232,7 @@ def any( @export.public_api(document_under="operations/functions") @constraints.dtypes( constraints={"input": "T1", constraints.RETURN_VALUE: "T1"}, - variables={"T1": ["float32", "int32", "float16", "bfloat16"]}, + variables={"T1": ["float32", "int32", "int64", "float16", "bfloat16"]}, ) def max( input: "tripy.Tensor", dim: Optional[Union[int, Sequence[int]]] = None, keepdim: bool = False @@ -265,7 +265,7 @@ def max( @export.public_api(document_under="operations/functions") @constraints.dtypes( constraints={"input": "T1", constraints.RETURN_VALUE: "T1"}, - variables={"T1": ["float32", "int32", "float16", "bfloat16"]}, + variables={"T1": ["float32", "int32", "int64", "float16", "bfloat16"]}, ) def prod( input: "tripy.Tensor", dim: Optional[Union[int, Sequence[int]]] = None, keepdim: bool = False @@ -313,7 +313,7 @@ def mean_impl(tensor: "tripy.Tensor", dim: Union[int, Sequence] = None, keepdim: @export.public_api(document_under="operations/functions") @constraints.dtypes( constraints={"input": "T1", constraints.RETURN_VALUE: "T1"}, - variables={"T1": ["float32", "int32", "float16", "bfloat16"]}, + variables={"T1": ["float32", "int32", "int64", "float16", "bfloat16"]}, ) def mean( input: "tripy.Tensor", dim: Optional[Union[int, Sequence[int]]] = None, keepdim: bool = False @@ -413,7 +413,7 @@ def _arg_min_max_impl(tensor: "tripy.Tensor", kind: ArgMinMax.Kind, dim: Optiona @export.public_api(document_under="operations/functions") @constraints.dtypes( constraints={"input": "T1", constraints.RETURN_VALUE: "T2"}, - variables={"T1": ["float32", "float16", "bfloat16", "int32", "bool", "int8"], "T2": ["int32"]}, + variables={"T1": ["float32", "float16", "bfloat16", "int32"], "T2": ["int32"]}, ) def argmax(input: "tripy.Tensor", dim: Optional[int] = None, keepdim: bool = False) -> "tripy.Tensor": """ @@ -445,7 +445,7 @@ def argmax(input: "tripy.Tensor", dim: Optional[int] = None, keepdim: bool = Fal @export.public_api(document_under="operations/functions") @constraints.dtypes( constraints={"input": "T1", constraints.RETURN_VALUE: "T2"}, - variables={"T1": ["float32", "float16", "bfloat16", "int32", "bool", "int8"], "T2": ["int32"]}, + variables={"T1": ["float32", "float16", "bfloat16", "int32"], "T2": ["int32"]}, ) def argmin(input: "tripy.Tensor", dim: Optional[int] = None, keepdim: bool = False) -> "tripy.Tensor": """ diff --git a/tripy/tripy/frontend/trace/ops/unary_elementwise.py b/tripy/tripy/frontend/trace/ops/unary_elementwise.py index 08e630f0c..bb2ad3d00 100644 --- a/tripy/tripy/frontend/trace/ops/unary_elementwise.py +++ b/tripy/tripy/frontend/trace/ops/unary_elementwise.py @@ -252,7 +252,7 @@ def log(input: "tripy.Tensor") -> "tripy.Tensor": @export.public_api(document_under="operations/functions") @constraints.dtypes( constraints={"input": "T1", constraints.RETURN_VALUE: "T1"}, - variables={"T1": ["float32", "float16", "bfloat16", "int8", "int32"]}, + variables={"T1": ["float32", "float16", "bfloat16", "int8", "int32", "int64"]}, ) def abs(input: "tripy.Tensor") -> "tripy.Tensor": r"""