Skip to content

Commit 0d435df

Browse files
pranavm-nvidiapranavm-nvidia
committed
Updates MLIR-TRT to 0.1.36
1 parent efaf830 commit 0d435df

File tree

9 files changed

+20
-21
lines changed

9 files changed

+20
-21
lines changed

‎tripy/pyproject.toml

+2-2
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,8 @@ requires-python = ">= 3.9"
88
license = {text = "Apache 2.0"}
99
dependencies = [
1010
"tensorrt~=10.0",
11-
"mlir-tensorrt-compiler==0.1.34+cuda12.trt102",
12-
"mlir-tensorrt-runtime==0.1.34+cuda12.trt102",
11+
"mlir-tensorrt-compiler==0.1.36+cuda12.trt102",
12+
"mlir-tensorrt-runtime==0.1.36+cuda12.trt102",
1313
"colored==2.2.3",
1414
]
1515

‎tripy/tests/integration/test_iota.py

+1-3
Original file line numberDiff line numberDiff line change
@@ -91,9 +91,7 @@ def test_negative_no_casting(self, dtype):
9191
a = tp.ones((2, 2))
9292
out = Iota.build([frontend_utils.tensor_from_shape_like(a.shape)], dim=0, output_rank=2, dtype=dtype)
9393

94-
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"
95-
if dtype == tp.bool:
96-
exception_str = "InternalError: failed to run compilation"
94+
exception_str = "InternalError: failed to run compilation"
9795
with helper.raises(
9896
tp.TripyException,
9997
match=exception_str,

‎tripy/tests/integration/test_quantize.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import tripy as tp
2424
from tests.helper import raises, TORCH_DTYPES
2525
from tests.conftest import skip_if_older_than_sm80, skip_if_older_than_sm89
26+
import cupy as cp
2627

2728

2829
class TestQuantize:
@@ -118,4 +119,4 @@ def test_non_constant_scale(self):
118119
scale = tp.ones((4,))
119120
quantized = tp.quantize(input, scale, tp.int8, dim=0)
120121

121-
assert bool(tp.all(quantized == tp.ones((4, 4), dtype=tp.int8)))
122+
assert bool(cp.all(cp.from_dlpack(quantized) == cp.ones((4, 4), dtype=cp.int8)))

‎tripy/tripy/frontend/ops/tensor_initializers.py

+4-4
Original file line numberDiff line numberDiff line change
@@ -165,7 +165,7 @@ def zeros_like(input: "tripy.Tensor", dtype: Optional[datatype.dtype] = None) ->
165165
@constraints.dtypes(
166166
constraints={"tensor": "T1", constraints.RETURN_VALUE: "T1"},
167167
variables={
168-
"T1": ["float32", "float16", "bfloat16", "int32", "bool"],
168+
"T1": ["float32", "float16", "bfloat16", "int32", "int64", "bool"],
169169
},
170170
)
171171
def tril(tensor: "tripy.Tensor", diagonal: int = 0) -> "tripy.Tensor":
@@ -223,7 +223,7 @@ def tril(tensor: "tripy.Tensor", diagonal: int = 0) -> "tripy.Tensor":
223223
@constraints.dtypes(
224224
constraints={"tensor": "T1", constraints.RETURN_VALUE: "T1"},
225225
variables={
226-
"T1": ["float32", "float16", "bfloat16", "int32", "bool"],
226+
"T1": ["float32", "float16", "bfloat16", "int32", "int64", "bool"],
227227
},
228228
)
229229
def triu(tensor: "tripy.Tensor", diagonal: int = 0) -> "tripy.Tensor":
@@ -281,7 +281,7 @@ def triu(tensor: "tripy.Tensor", diagonal: int = 0) -> "tripy.Tensor":
281281
@constraints.dtypes(
282282
constraints={"dtype": "T1", constraints.RETURN_VALUE: "T1"},
283283
variables={
284-
"T1": ["float32", "float16", "bfloat16", "int8", "int32", "bool"],
284+
"T1": ["float32", "float16", "bfloat16", "int8", "int32", "int64", "bool"],
285285
},
286286
)
287287
def arange(
@@ -346,7 +346,7 @@ def arange(
346346
@constraints.dtypes(
347347
constraints={"dtype": "T1", constraints.RETURN_VALUE: "T1"},
348348
variables={
349-
"T1": ["float32", "float16", "bfloat16", "int8", "int32", "bool"],
349+
"T1": ["float32", "float16", "bfloat16", "int8", "int32", "int64", "bool"],
350350
},
351351
)
352352
def arange(

‎tripy/tripy/frontend/trace/ops/flip.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ def to_flat_ir(self, inputs, outputs):
4040
@constraints.dtypes(
4141
constraints={"input": "T1", constraints.RETURN_VALUE: "T1"},
4242
variables={
43-
"T1": ["float32", "float16", "bfloat16", "int32", "bool"],
43+
"T1": ["float32", "float16", "bfloat16", "int32", "int64", "bool"],
4444
},
4545
)
4646
def flip(input: "tripy.Tensor", dims: Optional[Union[int, Sequence[int]]] = None) -> "tripy.Tensor":

‎tripy/tripy/frontend/trace/ops/gather.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,7 @@ def to_flat_ir(self, inputs, outputs):
9191
@constraints.dtypes(
9292
constraints={"input": "T1", "index": "T2", constraints.RETURN_VALUE: "T1"},
9393
variables={
94-
"T1": ["float32", "float16", "bfloat16", "int8", "int32", "bool"],
94+
"T1": ["float8", "float32", "float16", "bfloat16", "int4", "int8", "int32", "int64", "bool"],
9595
"T2": ["int32"],
9696
},
9797
)

‎tripy/tripy/frontend/trace/ops/iota.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,7 @@ def iota_impl(shape: "tripy.Tensor", dim: int, dtype: datatype.dtype, output_ran
6969
@constraints.dtypes(
7070
constraints={"dtype": "T1", constraints.RETURN_VALUE: "T1"},
7171
variables={
72-
"T1": ["float32", "float16", "bfloat16", "float8", "int4", "int8", "int32", "bool"],
72+
"T1": ["float32", "float16", "bfloat16", "float8", "int4", "int8", "int32", "int64", "bool"],
7373
},
7474
)
7575
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
101101
constraints={"input": "T1", "dtype": "T2", constraints.RETURN_VALUE: "T2"},
102102
variables={
103103
"T1": ["float32", "float16", "bfloat16", "float8", "int4", "int8", "int32", "int64", "bool"],
104-
"T2": ["float32", "float16", "bfloat16", "float8", "int4", "int8", "int32", "bool"],
104+
"T2": ["float32", "float16", "bfloat16", "float8", "int4", "int8", "int32", "int64", "bool"],
105105
},
106106
)
107107
def iota_like(input: "tripy.Tensor", dim: int = 0, dtype: Optional[datatype.dtype] = None) -> "tripy.Tensor":

‎tripy/tripy/frontend/trace/ops/reduce.py

+6-6
Original file line numberDiff line numberDiff line change
@@ -135,7 +135,7 @@ def _reduce_impl(input: "tripy.Tensor", kind: Reduce.Kind, dim: Union[int, Seque
135135
@export.public_api(document_under="operations/functions")
136136
@constraints.dtypes(
137137
constraints={"input": "T1", constraints.RETURN_VALUE: "T1"},
138-
variables={"T1": ["float32", "int32", "float16", "bfloat16"]},
138+
variables={"T1": ["float32", "int32", "int64", "float16", "bfloat16"]},
139139
)
140140
def sum(
141141
input: "tripy.Tensor", dim: Optional[Union[int, Sequence[int]]] = None, keepdim: bool = False
@@ -232,7 +232,7 @@ def any(
232232
@export.public_api(document_under="operations/functions")
233233
@constraints.dtypes(
234234
constraints={"input": "T1", constraints.RETURN_VALUE: "T1"},
235-
variables={"T1": ["float32", "int32", "float16", "bfloat16"]},
235+
variables={"T1": ["float32", "int32", "int64", "float16", "bfloat16"]},
236236
)
237237
def max(
238238
input: "tripy.Tensor", dim: Optional[Union[int, Sequence[int]]] = None, keepdim: bool = False
@@ -265,7 +265,7 @@ def max(
265265
@export.public_api(document_under="operations/functions")
266266
@constraints.dtypes(
267267
constraints={"input": "T1", constraints.RETURN_VALUE: "T1"},
268-
variables={"T1": ["float32", "int32", "float16", "bfloat16"]},
268+
variables={"T1": ["float32", "int32", "int64", "float16", "bfloat16"]},
269269
)
270270
def prod(
271271
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:
313313
@export.public_api(document_under="operations/functions")
314314
@constraints.dtypes(
315315
constraints={"input": "T1", constraints.RETURN_VALUE: "T1"},
316-
variables={"T1": ["float32", "int32", "float16", "bfloat16"]},
316+
variables={"T1": ["float32", "int32", "int64", "float16", "bfloat16"]},
317317
)
318318
def mean(
319319
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
413413
@export.public_api(document_under="operations/functions")
414414
@constraints.dtypes(
415415
constraints={"input": "T1", constraints.RETURN_VALUE: "T2"},
416-
variables={"T1": ["float32", "float16", "bfloat16", "int32", "bool", "int8"], "T2": ["int32"]},
416+
variables={"T1": ["float32", "float16", "bfloat16", "int32"], "T2": ["int32"]},
417417
)
418418
def argmax(input: "tripy.Tensor", dim: Optional[int] = None, keepdim: bool = False) -> "tripy.Tensor":
419419
"""
@@ -445,7 +445,7 @@ def argmax(input: "tripy.Tensor", dim: Optional[int] = None, keepdim: bool = Fal
445445
@export.public_api(document_under="operations/functions")
446446
@constraints.dtypes(
447447
constraints={"input": "T1", constraints.RETURN_VALUE: "T2"},
448-
variables={"T1": ["float32", "float16", "bfloat16", "int32", "bool", "int8"], "T2": ["int32"]},
448+
variables={"T1": ["float32", "float16", "bfloat16", "int32"], "T2": ["int32"]},
449449
)
450450
def argmin(input: "tripy.Tensor", dim: Optional[int] = None, keepdim: bool = False) -> "tripy.Tensor":
451451
"""

‎tripy/tripy/frontend/trace/ops/unary_elementwise.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -252,7 +252,7 @@ def log(input: "tripy.Tensor") -> "tripy.Tensor":
252252
@export.public_api(document_under="operations/functions")
253253
@constraints.dtypes(
254254
constraints={"input": "T1", constraints.RETURN_VALUE: "T1"},
255-
variables={"T1": ["float32", "float16", "bfloat16", "int8", "int32"]},
255+
variables={"T1": ["float32", "float16", "bfloat16", "int8", "int32", "int64"]},
256256
)
257257
def abs(input: "tripy.Tensor") -> "tripy.Tensor":
258258
r"""

0 commit comments

Comments
 (0)