tensor-lang/tensorlang.py at main · davro/tensor-lang · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
import os
import sys

from tensorlang.tensor_lang import TensorLang
from tensorlang.compiler import TensorCompiler
from tensorlang.test_runner import TestRunner

def main():
    tensorlang = TensorLang()
    args = tensorlang.parse_arguments()

    try:
        if args.test:
            # Test Mode: Run test suite
            runner = TestRunner(
                parallel=not args.no_parallel,
                jobs=args.jobs,
                verify_tensors=args.verify_tensors,
                debug_mode=args.debug
            )

            # Determine which tests to run
            if args.file:
                if args.file.endswith('.tl'):
                    test_files = [os.path.basename(args.file)]
                else:
                    all_tests = runner.discover_tests()
                    test_files = [t for t in all_tests if args.file in t]
            else:
                test_files = runner.discover_tests()
                if args.filter:
                    test_files = [t for t in test_files if args.filter in t]

            if not test_files:
                print(f"No tests found matching: {args.file or args.filter}")
                sys.exit(1)

            runner.run_test_suite(test_files)

        else:
            # Normal Mode: Compile and execute
            compiler = TensorCompiler(
                debug_mode=args.debug,
                debug_info=args.debug_info,
                debug_ast=args.debug_ast,
                cache_layers=args.cache_layers
            )
            compiler.compile_and_execute(args.file)

    except FileNotFoundError as e:
        print(f"Error: File not found - {args.file}")
        if args.debug:
            print(f"Details: {e}")
        sys.exit(1)

    except KeyboardInterrupt:
        print("\nInterrupted by user")
        sys.exit(0)

    except Exception as e:
        print(f"Unexpected error: {e}")
        if args.debug:
            import traceback
            traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()