Updates maximum, minimum, and where to accept non-tripy-tensors

Author: pranavm-nvidia

tripy/nvtripy/frontend/ops/binary/maximum.py

@@ -15,15 +15,17 @@
 from nvtripy import export
 from nvtripy.frontend.ops.binary.create import create_binary_op
 from nvtripy.trace.ops.binary import Max
+from nvtripy.types import TensorLike
 from nvtripy.utils import wrappers
 
 
 @export.public_api(document_under="operations/functions")
 @wrappers.interface(
     dtype_constraints={"lhs": "T1", "rhs": "T1", wrappers.RETURN_VALUE: "T1"},
     dtype_variables={"T1": ["float32", "float16", "bfloat16", "int8", "int32", "int64", "bool"]},
+    convert_to_tensors=True,
 )
-def maximum(lhs: "nvtripy.Tensor", rhs: "nvtripy.Tensor") -> "nvtripy.Tensor":
+def maximum(lhs: TensorLike, rhs: TensorLike) -> "nvtripy.Tensor":
     """
     Performs an elementwise maximum.
 

tripy/nvtripy/frontend/ops/binary/minimum.py

@@ -15,15 +15,17 @@
 from nvtripy import export
 from nvtripy.frontend.ops.binary.create import create_binary_op
 from nvtripy.trace.ops.binary import Min
+from nvtripy.types import TensorLike
 from nvtripy.utils import wrappers
 
 
 @export.public_api(document_under="operations/functions")
 @wrappers.interface(
     dtype_constraints={"lhs": "T1", "rhs": "T1", wrappers.RETURN_VALUE: "T1"},
     dtype_variables={"T1": ["float32", "float16", "bfloat16", "int8", "int32", "int64", "bool"]},
+    convert_to_tensors=True,
 )
-def minimum(lhs: "nvtripy.Tensor", rhs: "nvtripy.Tensor") -> "nvtripy.Tensor":
+def minimum(lhs: TensorLike, rhs: TensorLike) -> "nvtripy.Tensor":
     """
     Performs an elementwise minimum.
 

tripy/nvtripy/frontend/ops/where.py

@@ -19,6 +19,7 @@
 from nvtripy import export
 from nvtripy.frontend.ops import utils as op_utils
 from nvtripy.trace.ops.where import Where
+from nvtripy.types import TensorLike
 from nvtripy.utils import wrappers
 
 
@@ -29,8 +30,9 @@
         "T1": ["float32", "float16", "bfloat16", "int8", "int32", "int64"],
         "T2": ["bool"],
     },
+    convert_to_tensors=True,
 )
-def where(condition: "nvtripy.Tensor", input: "nvtripy.Tensor", other: "nvtripy.Tensor") -> "nvtripy.Tensor":
+def where(condition: "nvtripy.Tensor", input: TensorLike, other: TensorLike) -> "nvtripy.Tensor":
     r"""
     Returns a new tensor of elements selected from either ``input`` or ``other``, depending on ``condition``.