Function generate_init_py

export_hf.py:450–485

Generate the Python __init__.py for the HF Kernels module.

(
    name: str,
    functions: List[Dict[str, str]],
    repo_id: str,
    backend: str = "cuda",
)

# ---------------------------------------------------------------------------

def generate_init_py(
    name: str,
    functions: List[Dict[str, str]],
    repo_id: str,
    backend: str = "cuda",
) -> str:
    """Generate the Python __init__.py for the HF Kernels module."""
    first_func = functions[0]["name"] if functions else "kernel_fn"

    if backend == "cuda":
        return textwrap.dedent(f'''\
            """
            {name} - Optimized GPU kernel exported from AutoKernel
            https://github.com/RightNow-AI/autokernel

            Usage:
                from kernels import get_kernel
                module = get_kernel("{repo_id}")
                result = module.{first_func}(input)
            """
            from ._C import *  # noqa: F401,F403
        ''')
    else:
        # Triton kernel: import the Python module directly
        return textwrap.dedent(f'''\
            """
            {name} - Optimized Triton GPU kernel exported from AutoKernel
            https://github.com/RightNow-AI/autokernel

            Usage:
                from kernels import get_kernel
                module = get_kernel("{repo_id}")
                result = module.kernel_fn(input)
            """
            from .kernel import kernel_fn  # noqa: F401
        ''')


# ---------------------------------------------------------------------------