Remove "torchscript" association from the e2e framework.

We use it for more than TorchScript testing now. This is a purely
mechanical change to adjust some file paths to remove "torchscript".

The most perceptible change here is that now e2e tests are run with

```
./tools/e2e_test.sh
instead of:
./tools/torchscript_e2e_test.sh
```

.github/workflows/buildAndTest.yml

@@ -173,29 +173,29 @@ jobs:
       if: ${{ matrix.os-arch == 'ubuntu-x86_64' && matrix.llvm-build == 'in-tree' }}
       run: |
         export PYTHONPATH="$GITHUB_WORKSPACE/build/tools/torch-mlir/python_packages/torch_mlir"
-        python -m e2e_testing.torchscript.main --config=refbackend -v
+        python -m e2e_testing.main --config=refbackend -v
 
     - name: Run eager_mode e2e integration tests
       if: ${{ matrix.os-arch == 'ubuntu-x86_64' && matrix.llvm-build == 'in-tree' }}
       run: |
         export PYTHONPATH="$GITHUB_WORKSPACE/build/tools/torch-mlir/python_packages/torch_mlir"
-        python -m e2e_testing.torchscript.main --config=eager_mode -v
+        python -m e2e_testing.main --config=eager_mode -v
 
     - name: Run mhlo e2e integration tests
       if: ${{ matrix.os-arch == 'ubuntu-x86_64' && matrix.llvm-build == 'in-tree' }}
       run: |
         export PYTHONPATH="$GITHUB_WORKSPACE/build/tools/torch-mlir/python_packages/torch_mlir"
-        python -m e2e_testing.torchscript.main --config=mhlo -v
+        python -m e2e_testing.main --config=mhlo -v
 
     - name: Run tosa e2e integration tests
       if: ${{ matrix.os-arch == 'ubuntu-x86_64' && matrix.llvm-build == 'in-tree' }}
       run: |
         export PYTHONPATH="$GITHUB_WORKSPACE/build/tools/torch-mlir/python_packages/torch_mlir"
-        python -m e2e_testing.torchscript.main --config=tosa -v
+        python -m e2e_testing.main --config=tosa -v
 
     - name: Run lazy_tensor_core e2e integration tests
       if: ${{ matrix.os-arch == 'ubuntu-x86_64' && matrix.llvm-build == 'in-tree' }}
       run: |
         export PYTHONPATH="$GITHUB_WORKSPACE/build/tools/torch-mlir/python_packages/torch_mlir"
         echo "LTC tests disabled temporarily. https://github.com/llvm/torch-mlir/pull/1292"
-        # python -m e2e_testing.torchscript.main --config=lazy_tensor_core -v
+        # python -m e2e_testing.main --config=lazy_tensor_core -v

build_tools/torchscript_e2e_heavydep_tests/README.md

@@ -13,19 +13,19 @@ self-contained virtual environment. It can be used like so:
 # serialized test artifacts in the other specified directory.
 # This command is safe to re-run if you have already built the virtual
 # environment and just changed the tests.
-build_tools/torchscript_e2e_heavydep_tests/generate_serialized_tests.sh \
+build_tools/e2e_heavydep_tests/generate_serialized_tests.sh \
   path/to/heavydep_venv \
   path/to/heavydep_serialized_tests
 
 # Add the --serialized-test-dir flag to point at the directory containing the
 # serialized tests. All other functionality is the same as the normal invocation
-# of torchscript_e2e_test.sh, but the serialized tests will be available.
-tools/torchscript_e2e_test.sh --serialized-test-dir=path/to/heavydep_serialized_tests
+# of e2e_test.sh, but the serialized tests will be available.
+tools/e2e_test.sh --serialized-test-dir=path/to/heavydep_serialized_tests
 ```
 
 The tests use the same (pure-Python) test framework as the normal
-torchscript_e2e_test.sh, but the tests are added in
-`build_tools/torchscript_e2e_heavydep_tests` instead of
+e2e_test.sh, but the tests are added in
+`build_tools/e2e_heavydep_tests` instead of
 `frontends/pytorch/e2e_testing/torchscript`.
 
 We rely critically on serialized TorchScript compatibility across PyTorch

build_tools/torchscript_e2e_heavydep_tests/generate_serialized_tests.sh

@@ -38,4 +38,4 @@ cd "$torch_mlir_src_root"
 export PYTHONPATH=${PYTHONPATH-}
 source "$torch_mlir_src_root/.env"
 
-python3 -m build_tools.torchscript_e2e_heavydep_tests.main --output_dir="$serialized_test_dir"
+python3 -m build_tools.e2e_heavydep_tests.main --output_dir="$serialized_test_dir"

build_tools/torchscript_e2e_heavydep_tests/hf_sequence_classification.py

@@ -6,9 +6,9 @@
 import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 torch.manual_seed(0)
 

build_tools/torchscript_e2e_heavydep_tests/main.py

@@ -5,7 +5,7 @@
 
 import argparse
 
-from torch_mlir_e2e_test.torchscript.serialization import serialize_all_tests_to
+from torch_mlir_e2e_test.serialization import serialize_all_tests_to
 
 from . import hf_sequence_classification
 from . import vision_models

build_tools/torchscript_e2e_heavydep_tests/train_models.py

@@ -9,9 +9,9 @@ from functorch import make_fx
 from torch.nn.utils import _stateless
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 from torch import fx
 import copy

build_tools/torchscript_e2e_heavydep_tests/vision_models.py

@@ -6,9 +6,9 @@
 import torch
 import torchvision.models as models
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 import timm
 
 torch.manual_seed(0)

docs/development.md

@@ -205,11 +205,11 @@ Torch-MLIR has two types of tests:
 
 ```shell
 # Run all tests on the reference backend
-./tools/torchscript_e2e_test.sh
+./tools/e2e_test.sh
 # Run tests that match the regex `Conv2d`, with verbose errors.
-./tools/torchscript_e2e_test.sh --filter Conv2d --verbose
+./tools/e2e_test.sh --filter Conv2d --verbose
 # Run tests on the TOSA backend.
-./tools/torchscript_e2e_test.sh --config tosa
+./tools/e2e_test.sh --config tosa
 ```
 
 ## Running unit tests.
@@ -310,10 +310,10 @@ Here are some examples of PRs updating the LLVM and MLIR-HLO submodules:
 
 To enable ASAN, pass `-DLLVM_USE_SANITIZER=Address` to CMake. This should "just
 work" with all C++ tools like `torch-mlir-opt`. When running a Python script
-such as through `./tools/torchscript_e2e_test.sh`, you will need to do:
+such as through `./tools/e2e_test.sh`, you will need to do:
 
 ```
-LD_PRELOAD="$(clang -print-file-name=libclang_rt.asan-x86_64.so)" ./tools/torchscript_e2e_test.sh -s
+LD_PRELOAD="$(clang -print-file-name=libclang_rt.asan-x86_64.so)" ./tools/e2e_test.sh -s
 # See instructions here for how to get the libasan path for GCC:
 # https://stackoverflow.com/questions/48833176/get-location-of-libasan-from-gcc-clang
 ```

e2e_testing/main.py

@@ -7,14 +7,14 @@ import argparse
 import re
 import sys
 
-from torch_mlir_e2e_test.torchscript.framework import run_tests
-from torch_mlir_e2e_test.torchscript.reporting import report_results
-from torch_mlir_e2e_test.torchscript.registry import GLOBAL_TEST_REGISTRY
-from torch_mlir_e2e_test.torchscript.serialization import deserialize_all_tests_from
+from torch_mlir_e2e_test.framework import run_tests
+from torch_mlir_e2e_test.reporting import report_results
+from torch_mlir_e2e_test.registry import GLOBAL_TEST_REGISTRY
+from torch_mlir_e2e_test.serialization import deserialize_all_tests_from
 
 
 # Available test configs.
-from torch_mlir_e2e_test.torchscript.configs import (
+from torch_mlir_e2e_test.configs import (
     LazyTensorCoreTestConfig,
     LinalgOnTensorsBackendTestConfig,
     MhloBackendTestConfig,
@@ -62,7 +62,7 @@ The directory containing serialized pre-built tests.
 Right now, these are additional tests which require heavy Python dependencies
 to generate (or cannot even be generated with the version of PyTorch used by
 torch-mlir).
-See `build_tools/torchscript_e2e_heavydep_tests/generate_serialized_tests.sh`
+See `build_tools/e2e_heavydep_tests/generate_serialized_tests.sh`
 for more information on building these artifacts.
 ''')
     parser.add_argument('-s', '--sequential',

python/test/annotations-sugar.py

@@ -7,7 +7,7 @@
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.annotations import annotate_args, export
 from torch_mlir.dialects.torch.importer.jit_ir import ClassAnnotator
 from torch_mlir.dialects.torch.importer.jit_ir.torchscript_annotations import extract_annotations
 

python/test/torchscript_e2e_test/basic.py

@@ -7,10 +7,10 @@
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import run_tests, TestUtils
-from torch_mlir_e2e_test.torchscript.reporting import report_results
-from torch_mlir_e2e_test.torchscript.registry import register_test_case, GLOBAL_TEST_REGISTRY
-from torch_mlir_e2e_test.torchscript.configs import TorchScriptTestConfig
+from torch_mlir_e2e_test.framework import run_tests, TestUtils
+from torch_mlir_e2e_test.reporting import report_results
+from torch_mlir_e2e_test.registry import register_test_case, GLOBAL_TEST_REGISTRY
+from torch_mlir_e2e_test.configs import TorchScriptTestConfig
 
 
 class MmModule(torch.nn.Module):

python/test/torchscript_e2e_test/compilation_failure.py

@@ -7,10 +7,10 @@
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import run_tests, TestUtils
-from torch_mlir_e2e_test.torchscript.reporting import report_results
-from torch_mlir_e2e_test.torchscript.registry import register_test_case, GLOBAL_TEST_REGISTRY
-from torch_mlir_e2e_test.torchscript.configs import TorchScriptTestConfig
+from torch_mlir_e2e_test.framework import run_tests, TestUtils
+from torch_mlir_e2e_test.reporting import report_results
+from torch_mlir_e2e_test.registry import register_test_case, GLOBAL_TEST_REGISTRY
+from torch_mlir_e2e_test.configs import TorchScriptTestConfig
 
 
 class MmModule(torch.nn.Module):

python/test/torchscript_e2e_test/error_reports.py

@@ -9,10 +9,10 @@ from typing import List, Tuple, Dict
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import run_tests, TestUtils
-from torch_mlir_e2e_test.torchscript.reporting import report_results
-from torch_mlir_e2e_test.torchscript.registry import register_test_case, GLOBAL_TEST_REGISTRY
-from torch_mlir_e2e_test.torchscript.configs import TorchScriptTestConfig
+from torch_mlir_e2e_test.framework import run_tests, TestUtils
+from torch_mlir_e2e_test.reporting import report_results
+from torch_mlir_e2e_test.registry import register_test_case, GLOBAL_TEST_REGISTRY
+from torch_mlir_e2e_test.configs import TorchScriptTestConfig
 
 # CHECK: Unexpected outcome summary:
 # CHECK: FAIL - "ErroneousModule_basic"

python/test/torchscript_e2e_test/non_tensor_values.py

@@ -9,10 +9,10 @@ from typing import List, Tuple, Dict
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import run_tests, TestUtils
-from torch_mlir_e2e_test.torchscript.reporting import report_results
-from torch_mlir_e2e_test.torchscript.registry import register_test_case, GLOBAL_TEST_REGISTRY
-from torch_mlir_e2e_test.torchscript.configs import TorchScriptTestConfig
+from torch_mlir_e2e_test.framework import run_tests, TestUtils
+from torch_mlir_e2e_test.reporting import report_results
+from torch_mlir_e2e_test.registry import register_test_case, GLOBAL_TEST_REGISTRY
+from torch_mlir_e2e_test.configs import TorchScriptTestConfig
 
 
 class NonTensorValuesModule(torch.nn.Module):

python/test/torchscript_e2e_test/runtime_failure.py

@@ -7,10 +7,10 @@
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import run_tests, TestUtils
-from torch_mlir_e2e_test.torchscript.reporting import report_results
-from torch_mlir_e2e_test.torchscript.registry import register_test_case, GLOBAL_TEST_REGISTRY
-from torch_mlir_e2e_test.torchscript.configs import TorchScriptTestConfig
+from torch_mlir_e2e_test.framework import run_tests, TestUtils
+from torch_mlir_e2e_test.reporting import report_results
+from torch_mlir_e2e_test.registry import register_test_case, GLOBAL_TEST_REGISTRY
+from torch_mlir_e2e_test.configs import TorchScriptTestConfig
 
 
 class MmModule(torch.nn.Module):

python/test/torchscript_e2e_test/submodule.py

@@ -7,10 +7,10 @@
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import run_tests, TestUtils
-from torch_mlir_e2e_test.torchscript.reporting import report_results
-from torch_mlir_e2e_test.torchscript.registry import register_test_case, GLOBAL_TEST_REGISTRY
-from torch_mlir_e2e_test.torchscript.configs import TorchScriptTestConfig
+from torch_mlir_e2e_test.framework import run_tests, TestUtils
+from torch_mlir_e2e_test.reporting import report_results
+from torch_mlir_e2e_test.registry import register_test_case, GLOBAL_TEST_REGISTRY
+from torch_mlir_e2e_test.configs import TorchScriptTestConfig
 
 class Submodule2(torch.nn.Module):
     def __init__(self):

python/torch_mlir/dialects/torch/importer/jit_ir/torchscript_annotations.py

@@ -18,7 +18,7 @@ from torch_mlir.dialects.torch.importer.jit_ir import ClassAnnotator
 # to be expressed conveniently and gives clearer error reports when
 # the annotations aren't acceptable.
 
-# This module is kept separate from torch_mlir_e2e_test.torchscript.annotations so that
+# This module is kept separate from torch_mlir_e2e_test.annotations so that
 # we can use that module from code without C++ dependencies, which prevent us
 # from interfacing the test framework across environments.
 

python/torch_mlir_e2e_test/configs/eager_mode.py

@@ -7,7 +7,7 @@ import torch
 from torch.utils._pytree import tree_map
 
 from torch_mlir.eager_mode.torch_mlir_tensor import TorchMLIRTensor
-from torch_mlir_e2e_test.torchscript.framework import TestConfig, Trace, TraceItem
+from torch_mlir_e2e_test.framework import TestConfig, Trace, TraceItem
 
 
 def wrap(e):

python/torch_mlir_e2e_test/configs/lazy_tensor_core.py

@@ -7,7 +7,7 @@ import torch_mlir._mlir_libs._REFERENCE_LAZY_BACKEND as lazy_backend
 import torch
 from torch.utils._pytree import tree_map
 
-from torch_mlir_e2e_test.torchscript.framework import TestConfig, Trace, TraceItem
+from torch_mlir_e2e_test.framework import TestConfig, Trace, TraceItem
 
 
 def to_device(device):

python/torch_mlir_e2e_test/configs/linalg_on_tensors_backend.py

@@ -9,7 +9,7 @@ import torch
 import torch_mlir
 
 from torch_mlir_e2e_test.linalg_on_tensors_backends.abc import LinalgOnTensorsBackend
-from torch_mlir_e2e_test.torchscript.framework import TestConfig, Trace, TraceItem
+from torch_mlir_e2e_test.framework import TestConfig, Trace, TraceItem
 from torch_mlir_e2e_test.utils import convert_annotations_to_placeholders
 
 from .utils import (

python/torch_mlir_e2e_test/configs/mhlo_backend.py

@@ -9,7 +9,7 @@ import torch
 import torch_mlir
 
 from torch_mlir_e2e_test.mhlo_backends.abc import MhloBackend
-from torch_mlir_e2e_test.torchscript.framework import (
+from torch_mlir_e2e_test.framework import (
     TestConfig,
     Trace,
     TraceItem

python/torch_mlir_e2e_test/configs/native_torch.py

@@ -8,7 +8,7 @@ from typing import Any
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import TestConfig, Trace, TraceItem
+from torch_mlir_e2e_test.framework import TestConfig, Trace, TraceItem
 
 
 class NativeTorchTestConfig(TestConfig):

python/torch_mlir_e2e_test/configs/torchscript.py

@@ -8,7 +8,7 @@ from typing import Any
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import TestConfig, Trace, TraceItem
+from torch_mlir_e2e_test.framework import TestConfig, Trace, TraceItem
 
 
 class TorchScriptTestConfig(TestConfig):

python/torch_mlir_e2e_test/configs/tosa_backend.py

@@ -9,7 +9,7 @@ import torch
 import torch_mlir
 
 from torch_mlir_e2e_test.tosa_backends.abc import TosaBackend
-from torch_mlir_e2e_test.torchscript.framework import TestConfig, Trace, TraceItem
+from torch_mlir_e2e_test.framework import TestConfig, Trace, TraceItem
 from torch_mlir_e2e_test.utils import convert_annotations_to_placeholders
 from .utils import (
     recursively_convert_to_numpy,

python/torch_mlir_e2e_test/configs/utils.py

@@ -3,17 +3,11 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 # Also available under a BSD-style license. See LICENSE.
 
-import sys
 from typing import Any
-from io import StringIO
 
 import numpy as np
 import torch
 
-from torch_mlir.dialects.torch.importer.jit_ir import ClassAnnotator, ModuleBuilder
-from torch_mlir.dialects.torch.importer.jit_ir.torchscript_annotations import extract_annotations
-from torch_mlir.compiler_utils import run_pipeline_with_repro_report
-
 
 def recursively_convert_to_numpy(o: Any):
     if isinstance(o, torch.Tensor):

python/torch_mlir_e2e_test/test_suite/arange.py

@@ -5,9 +5,9 @@
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 # ==============================================================================
 

python/torch_mlir_e2e_test/test_suite/argmax.py

@@ -4,9 +4,9 @@
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 # ==============================================================================
 
@@ -63,4 +63,3 @@ class ArgmaxKeepDimsModule(torch.nn.Module):
 @register_test_case(module_factory=lambda: ArgmaxKeepDimsModule())
 def ArgmaxModule_keepDim(module, tu: TestUtils):
     module.forward(tu.rand(4, 6))
-

python/torch_mlir_e2e_test/test_suite/backprop.py

@@ -5,9 +5,9 @@
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 # ==============================================================================
 

python/torch_mlir_e2e_test/test_suite/basic.py

@@ -5,9 +5,9 @@
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 # ==============================================================================
 

python/torch_mlir_e2e_test/test_suite/cast.py

@@ -5,9 +5,9 @@
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 # ==============================================================================
 
@@ -122,4 +122,3 @@ class TensorToBool(torch.nn.Module):
 @register_test_case(module_factory=lambda: TensorToBool())
 def TensorToBool_basic(module, tu: TestUtils):
     module.forward(torch.tensor([[1]], dtype=torch.bool))
-

python/torch_mlir_e2e_test/test_suite/constant_alloc.py

@@ -5,9 +5,9 @@
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 # ==============================================================================
 

python/torch_mlir_e2e_test/test_suite/control_flow.py

@@ -7,9 +7,9 @@ from numpy import int64
 import torch
 import random
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 # ==============================================================================
 

python/torch_mlir_e2e_test/test_suite/conv.py

@@ -4,9 +4,9 @@
 # Also available under a BSD-style license. See LICENSE.
 
 import torch
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 # ==============================================================================
 
@@ -704,4 +704,4 @@ class Conv_Transpose3dModule(torch.nn.Module):
 
 @register_test_case(module_factory=lambda: Conv_Transpose3dModule())
 def Conv_Transpose3dModule_basic(module, tu: TestUtils):
-    module.forward(torch.randn(5, 2, 5, 6, 4), torch.randn(2, 5, 2, 2, 2))
+    module.forward(torch.randn(5, 2, 5, 6, 4), torch.randn(2, 5, 2, 2, 2))

python/torch_mlir_e2e_test/test_suite/custom_op_example.py

@@ -5,9 +5,9 @@
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 # ==============================================================================
 
@@ -33,4 +33,3 @@ class CustomOpExampleModule(torch.nn.Module):
 @register_test_case(module_factory=lambda: CustomOpExampleModule())
 def CustomOpExampleModule_basic(module, tu: TestUtils):
     module.forward(tu.rand(3, 4))
-

python/torch_mlir_e2e_test/test_suite/elementwise.py

@@ -5,9 +5,9 @@
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 # TODO: Support scalar !torch.int/!torch.float variants. Add support to
 # ReduceOpVariants to implement them in terms of the tensor-only variants +

python/torch_mlir_e2e_test/test_suite/elementwise_comparison.py

@@ -5,9 +5,9 @@
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 # ==============================================================================
 
@@ -571,4 +571,3 @@ class AllBoolFalseModule(torch.nn.Module):
 @register_test_case(module_factory=lambda: AllBoolFalseModule())
 def AllBoolFalseModule_basic(module, tu: TestUtils):
     module.forward()
-

python/torch_mlir_e2e_test/test_suite/histogram_binning_calibration.py

@@ -5,9 +5,9 @@
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 
 # ==============================================================================

python/torch_mlir_e2e_test/test_suite/index_put.py

@@ -5,9 +5,9 @@
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 # ==============================================================================
 

python/torch_mlir_e2e_test/test_suite/index_select.py

@@ -5,9 +5,9 @@
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 # ==============================================================================
 

python/torch_mlir_e2e_test/test_suite/matmul.py

@@ -5,9 +5,9 @@
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 # ==============================================================================
 

python/torch_mlir_e2e_test/test_suite/mlp.py

@@ -6,9 +6,9 @@
 import torch
 import torch.nn as nn
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 # ==============================================================================
 

python/torch_mlir_e2e_test/test_suite/nll_loss.py

@@ -5,9 +5,9 @@
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 # ==============================================================================
 

python/torch_mlir_e2e_test/test_suite/norm_like.py

@@ -5,9 +5,9 @@
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 # ==============================================================================
 

python/torch_mlir_e2e_test/test_suite/pooling.py

@@ -5,9 +5,9 @@
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 # ==============================================================================
 

python/torch_mlir_e2e_test/test_suite/quantized_models.py

@@ -6,9 +6,9 @@
 import torch
 from torch import nn
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 # ==============================================================================
 

python/torch_mlir_e2e_test/test_suite/reduction.py

@@ -5,9 +5,9 @@
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 # ==============================================================================
 

python/torch_mlir_e2e_test/test_suite/reshape_like.py

@@ -4,9 +4,9 @@
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 # ==============================================================================
 
@@ -561,4 +561,3 @@ class ReshapeAliasCollapseModule(torch.nn.Module):
 @register_test_case(module_factory=lambda: ReshapeAliasCollapseModule())
 def ReshapeAliasCollapseModule_basic(module, tu: TestUtils):
     module.forward(tu.rand(2, 4))
-

python/torch_mlir_e2e_test/test_suite/return_types.py

@@ -5,9 +5,9 @@
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 # ==============================================================================
 

python/torch_mlir_e2e_test/test_suite/rng.py

@@ -1,8 +1,8 @@
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 # ==============================================================================
 

python/torch_mlir_e2e_test/test_suite/scalar.py

@@ -5,9 +5,9 @@
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 # ==============================================================================
 

python/torch_mlir_e2e_test/test_suite/scalar_comparison.py

@@ -5,9 +5,9 @@
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 # ==============================================================================
 

python/torch_mlir_e2e_test/test_suite/slice_like.py

@@ -5,9 +5,9 @@
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 # ==============================================================================
 

python/torch_mlir_e2e_test/test_suite/squeeze.py

@@ -5,9 +5,9 @@
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 # ==============================================================================
 

python/torch_mlir_e2e_test/test_suite/stats.py

@@ -5,9 +5,9 @@
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 # ==============================================================================
 

python/torch_mlir_e2e_test/test_suite/table_batch_embedding.py

@@ -5,9 +5,9 @@
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 # ==============================================================================
 
@@ -58,4 +58,3 @@ def TableBatchEmbeddingModule_basic(module, tu: TestUtils):
     offsets = torch.cumsum(
         torch.tensor([0] + [BAG_SIZE for _ in range(BATCH_SIZE - 1)], dtype=torch.int64), 0)
     module.forward(indices, offsets)
-

python/torch_mlir_e2e_test/test_suite/threshold.py

@@ -5,9 +5,9 @@
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 # ==============================================================================
 

python/torch_mlir_e2e_test/test_suite/type_conversion.py

@@ -5,9 +5,9 @@
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 # ==============================================================================
 

python/torch_mlir_e2e_test/test_suite/type_promotion.py

@@ -5,9 +5,9 @@
 
 import torch
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 # ==============================================================================
 
@@ -111,5 +111,3 @@ class TypePromotionAlphaWiderModule(torch.nn.Module):
 @register_test_case(module_factory=lambda: TypePromotionAlphaWiderModule())
 def TypePromotionAlphaWiderModule_basic(module, tu: TestUtils):
     module.forward(tu.rand(4), tu.rand())
-
-

python/torch_mlir_e2e_test/test_suite/vision_models.py

@@ -6,9 +6,9 @@
 import torch
 import torchvision.models as models
 
-from torch_mlir_e2e_test.torchscript.framework import TestUtils
-from torch_mlir_e2e_test.torchscript.registry import register_test_case
-from torch_mlir_e2e_test.torchscript.annotations import annotate_args, export
+from torch_mlir_e2e_test.framework import TestUtils
+from torch_mlir_e2e_test.registry import register_test_case
+from torch_mlir_e2e_test.annotations import annotate_args, export
 
 # ==============================================================================
 

python/torch_mlir_e2e_test/torchscript/CMakeLists.txt

@@ -1,8 +0,0 @@
-## Declare the sources of the Python module.
-
-declare_mlir_python_sources(TorchMLIRPythonSources.TorchScriptE2ETest
-  ROOT_DIR "${TORCH_MLIR_PYTHON_ROOT_DIR}"
-  ADD_TO_PARENT TorchMLIRPythonSources
-  SOURCES_GLOB
-    dialects/torch/e2e_test/torchscript/*.py
-)

python/torch_mlir_e2e_test/utils.py

@@ -4,7 +4,7 @@
 # Also available under a BSD-style license. See LICENSE.
 
 from torch_mlir import TensorPlaceholder
-from torch_mlir_e2e_test.torchscript.annotations import TORCH_MLIR_ARG_ANNOTATIONS_ATTR_NAME
+from torch_mlir_e2e_test.annotations import TORCH_MLIR_ARG_ANNOTATIONS_ATTR_NAME
 
 def convert_annotations_to_placeholders(forward_method):
     """Converts the annotations on a forward method into tensor placeholders.

tools/e2e_test.sh

@@ -9,4 +9,4 @@ cd "$src_dir"
 export PYTHONPATH=${PYTHONPATH-}
 source .env
 
-python -m e2e_testing.torchscript.main "$@"
+python -m e2e_testing.main "$@"