Various bug fixes to fix L1 tests

- Adds tests to ensure we only use ASCII characters in our code and markdown files.
     Some non-ASCII characters that look like ASCII characters were breaking tests
     that make assumptions about the encoding.

- Updates dtype constraint tests to properly construct inputs for various APIs.

- Moves dtype constraint tests to L0 since they only add about 20 seconds and were
     breaking very often otherwise.

- Updates binary elementwise ops to only return `ShapeScalar`s if *all* inputs are also
     `ShapeScalar`s.

- Corrects dtype constraints for several APIs.

- Simplifies `mean` implementation.

- Corrects type annotations for some APIs and introduces a `ShapeLike` type.

- Adds support for `Optional` during type checking.

- Updates `convert_inputs_to_tensors` to now potentially convert tensors to
     `Tensor` subclasses. See the note in the code for details on the logic.

Author: pranavm-nvidia

tripy/docs/post0_developer_guides/design-decisions.md

@@ -100,7 +100,7 @@ to MLIR if needed or return the newly created ones.
 
 ## Why Not Build On [JAX](https://github.com/google/jax)?
 
-Tripy’s architecture looks very similar to JAX's, where python code is staged out and
+Tripy's architecture looks very similar to JAX's, where python code is staged out and
 lowered into a custom IR and eventually to MLIR.
 
 Then why not build on top of JAX? There are a couple reasons:

tripy/docs/pre0_user_guides/02-compiler.md

@@ -58,7 +58,7 @@ fast_geglu(inp).eval()
 ### Optimization Profiles
 
 In the example above, we assumed `inp` has a static shape of `(1, 2)`.
-Now, let’s assume that the shape of `inp` can vary from `(1, 2)` to `(16, 2)`, with `(8, 2)`
+Now, let's assume that the shape of `inp` can vary from `(1, 2)` to `(16, 2)`, with `(8, 2)`
 being the shape we'd like to optimize for. To express this constraint to the compiler,
 we can provide the range of shapes to `InputInfo` using `shape=((1, 8, 16), 2)`.
 This indicates to the compiler that the first dimension can vary from 1 to 16,

tripy/tests/flat_ir/ops/test_gather.py

@@ -20,6 +20,7 @@
 
 from tripy.flat_ir.ops import DynamicGatherOp
 from tripy.frontend.trace import Trace
+import re
 
 
 class TestGatherOp:
@@ -38,9 +39,9 @@ def test_gather_str(self, axis):
         reshape = flat_ir.ops[-2]
         print(str(reshape))
         assert isinstance(gather, DynamicGatherOp)
-        assert (
-            str(gather)
-            == f"out: [rank=(3), dtype=(float32), loc=(gpu:0)] = DynamicGatherOp(data, indices, t_inter4, axis={axis})"
+        assert re.match(
+            rf"out: \[rank=\(3\), dtype=\(float32\), loc=\(gpu:0\)\] = DynamicGatherOp\(data, indices, t_inter[0-9]+, axis={axis}\)",
+            str(gather),
         )
 
     @pytest.mark.parametrize("axis", [0, 1])

tripy/tests/helper.py

@@ -43,17 +43,24 @@
 
 ROOT_DIR = os.path.realpath(os.path.join(os.path.dirname(__file__), os.path.pardir))
 
-MARKDOWN_FILES = [
-    path
-    for path in glob.glob(os.path.join(ROOT_DIR, "**", "*.md"), recursive=True)
-    if not path.startswith(
-        (
-            os.path.join(ROOT_DIR, "build"),
-            os.path.join(ROOT_DIR, "mlir-tensorrt"),
-            os.path.join(ROOT_DIR, "stablehlo"),
+
+def get_files_with_extension(ext):
+    return [
+        path
+        for path in glob.glob(os.path.join(ROOT_DIR, "**", f"*{ext}"), recursive=True)
+        if not path.startswith(
+            (
+                os.path.join(ROOT_DIR, "build"),
+                os.path.join(ROOT_DIR, "mlir-tensorrt"),
+                os.path.join(ROOT_DIR, "stablehlo"),
+            )
         )
-    )
-]
+    ]
+
+
+MARKDOWN_FILES = get_files_with_extension(".md")
+
+PYTHON_FILES = get_files_with_extension(".py")
 
 
 @contextlib.contextmanager

tripy/tests/spec_verification/object_builders.py

@@ -60,6 +60,7 @@ def default_builder(init, dtype, namespace):
 
 find_func = {
     "tripy.Tensor": tensor_builder,
+    "tripy.types.TensorLike": tensor_builder,
     "tripy.Shape": tensor_builder,
     "tripy.dtype": dtype_builder,
     datatype.dtype: dtype_builder,
@@ -83,6 +84,12 @@ def default_builder(init, dtype, namespace):
 default_constraints_all = {
     "__getitem__": {"index": 2},
     "__matmul__": {"self": tp.ones((2, 3))},
+    # Force broadcasting for binary ops so the entire broadcasting code path is triggered.
+    "__add__": {"other": 1},
+    "__mul__": {"other": 1},
+    "__pow__": {"other": 1},
+    "__sub__": {"other": 1},
+    "__truediv__": {"other": 1},
     "__radd__": {"self": 1},
     "__rmul__": {"self": 1},
     "__rpow__": {"self": 1},
@@ -105,29 +112,31 @@ def default_builder(init, dtype, namespace):
     },
     "cumsum": {"dim": 0},
     "dequantize": {"scale": tp.Tensor([1, 1, 1]), "dim": 0},
-    "expand": {"sizes": tp.Tensor([3, 4]), "input": tp.ones((3, 1))},
+    "expand": {"sizes": [3, 4], "input": tp.ones((3, 1))},
     "flip": {"dim": 1},
     "full_like": {"value": 1},
-    "full": {"shape": tp.Tensor([3]), "value": 1},
+    "full": {"shape": [3], "value": 1},
     "gather": {"dim": 0, "index": tp.Tensor([1])},
-    "iota": {"shape": tp.Tensor([4])},
+    "iota": {"shape": [4]},
     "masked_fill": {"value": 1},
+    "maxpool": {"input": tp.ones((1, 3, 5, 5)), "kernel_dims": (3, 3)},
     "max": {"dim": 0},
     "mean": {"dim": 0},
-    "ones": {"shape": tp.Tensor([3, 2])},
+    "ones": {"shape": [3, 2]},
+    "outer": {"vec1": tp.Tensor([2, 3, 4, 5]), "vec2": tp.Tensor([1, 2, 3, 4])},
     "permute": {"perm": [1, 0]},
     "prod": {"dim": 0},
     "quantize": {"scale": tp.Tensor([1, 1, 1]), "dim": 0},
     "repeat": {"repeats": 2, "dim": 0},
-    "reshape": {"shape": tp.Tensor([6])},
+    "reshape": {"shape": [6]},
     "softmax": {"dim": 1},
     "split": {"indices_or_sections": 2},
     "squeeze": {"input": tp.ones((3, 1)), "dims": (1)},
     "sum": {"dim": 0},
     "transpose": {"dim0": 0, "dim1": 1},
     "unsqueeze": {"dim": 1},
     "var": {"dim": 0},
-    "zeros": {"shape": tp.Tensor([3, 2])},
+    "zeros": {"shape": [3, 2]},
 }
 
 

tripy/tests/spec_verification/test_dtype_constraints.py

@@ -172,8 +172,6 @@ def _run_dtype_constraints_subtest(test_data):
     return ret_val, namespace
 
 
-# Positive dtype testing is run during L1 testing.
-@pytest.mark.l1
 @pytest.mark.parametrize("test_data", DTYPE_CONSTRAINT_CASES, ids=lambda val: val[-1])
 def test_dtype_constraints(test_data):
     # If data type checking is enabled, negative tests will trivially pass (we will throw an

tripy/tests/test_files.py

@@ -0,0 +1,34 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import pytest
+from tests import helper
+
+
+# Checks all Python and markdown files to ensure there are no non-ASCII characters
+@pytest.mark.parametrize("file", helper.MARKDOWN_FILES + helper.PYTHON_FILES)
+def test_no_non_ascii_characters(file):
+    with open(file, "rb") as f:
+        contents = f.read()
+
+    try:
+        contents.decode("ascii")
+    except UnicodeDecodeError as err:
+        str_contents = contents.decode("utf-8")
+
+        non_ascii = str_contents[err.start : err.end]
+
+        line_num = [line_num for line_num, line in enumerate(str_contents.splitlines()) if non_ascii in line][0] + 1
+
+        assert False, f"Detected non-ASCII character(s) on line {line_num}: {non_ascii}"

tripy/tripy/common/exception.py

@@ -97,9 +97,6 @@ def apply_color(inp, color):
         frame_info += " " * start + apply_color("^" * (size), Fore.red)
         if not is_first_frame:
             frame_info += " --- required from here"
-    else:
-        if not is_first_frame:
-            frame_info = "Required from:\n" + frame_info
     frame_info += "\n\n"
     return frame_info
 

tripy/tripy/export.py

@@ -19,7 +19,7 @@
 from dataclasses import dataclass
 from typing import List, Optional, Any
 from types import ModuleType
-
+from textwrap import dedent
 from tripy.function_registry import FunctionRegistry
 
 
@@ -40,7 +40,11 @@ class PublicAPI:
 
 
 def public_api(
-    document_under: str = "", autodoc_options: Optional[List[str]] = None, module: ModuleType = None, symbol: str = None
+    document_under: str = "",
+    autodoc_options: Optional[List[str]] = None,
+    module: ModuleType = None,
+    symbol: str = None,
+    doc: str = None,
 ):
     """
     Decorator that exports a function/class to the public API under the top-level module and
@@ -71,6 +75,9 @@ def public_api(
         module: The module under which to export this public API. Defaults to the top-level Tripy module.
 
         symbol: The name of the symbol, if different from ``__name__``.
+
+        doc: Optional docstring. This is useful in cases where the docstring cannot be provided as normal.
+            For example, global variables sometimes don't register docstrings correctly.
     """
     assert not autodoc_options or (
         ":no-members:" not in autodoc_options or ":no-special-members:" in autodoc_options
@@ -80,6 +87,9 @@ def export_impl(obj):
         nonlocal module, symbol
         import tripy
 
+        if doc is not None:
+            obj.__doc__ = dedent(doc)
+
         module = module or tripy
 
         symbol = symbol or obj.__name__

tripy/tripy/frontend/module/parameter.py

@@ -17,7 +17,6 @@
 
 from typing import Any, Sequence
 
-import tripy.frontend.utils as frontend_utils
 from tripy import export, utils
 from tripy.frontend.tensor import Tensor
 from tripy.utils import Result
@@ -30,7 +29,6 @@ class Parameter(Tensor):
     constant, enabling additional optimization opportunities.
     """
 
-    @frontend_utils.convert_inputs_to_tensors()
     def __init__(self, tensor: Any) -> None:
         """
         Args: