@@ -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)
140140def 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)
237237def 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)
270270def 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)
318318def 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)
418418def 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)
450450def argmin (input : "tripy.Tensor" , dim : Optional [int ] = None , keepdim : bool = False ) -> "tripy.Tensor" :
451451 """
0 commit comments