Function run_op_test_main

backends/mlx/test/test_utils.py:1057–1130

Common main() function for op tests. This handles the common argparse setup, rebuild logic, and generate/compare/run action handling that is shared across all op tests. Args: test_factory: A callable that takes parsed args (argparse.Namespace) and returns an Op

(
    test_factory,
    description: str,
    add_args_fn=None,
)

def run_op_test_main(
    test_factory,
    description: str,
    add_args_fn=None,
):
    """
    Common main() function for op tests.

    This handles the common argparse setup, rebuild logic, and generate/compare/run
    action handling that is shared across all op tests.

    Args:
        test_factory: A callable that takes parsed args (argparse.Namespace) and
            returns an OpTestCase instance.
        description: Description for the argparse help message.
        add_args_fn: Optional callable that takes a parser and adds test-specific
            arguments. Signature: add_args_fn(parser) -> None
    """
    import argparse
    import sys

    parser = argparse.ArgumentParser(description=description)
    parser.add_argument(
        "action",
        choices=["generate", "compare", "run"],
        help="Action to perform: generate (create test files), compare (compare outputs), run (full test)",
    )
    parser.add_argument("--verbose", "-v", action="store_true", help="Verbose output")
    parser.add_argument(
        "--rebuild",
        action="store_true",
        help="Rebuild the C++ test runner before running",
    )

    # Add test-specific arguments
    if add_args_fn is not None:
        add_args_fn(parser)

    args = parser.parse_args()

    # Rebuild if requested
    if args.rebuild:
        if not rebuild_op_test_runner(verbose=args.verbose):
            sys.exit(1)

    # Create test case from factory
    test = test_factory(args)

    if args.action == "generate":
        pte_path, input_path, expected_path = test.generate_test_files(
            verbose=args.verbose
        )
        print("\nGenerated files:")
        print(f"  PTE:      {pte_path}")
        print(f"  Input:    {input_path}")
        print(f"  Expected: {expected_path}")
        print_mlx_graph_summary(pte_path)