from setuptools import setup
from torch.utils.cpp_extension import CppExtension, BuildExtension

setup(
    name='sparse_linear',
    ext_modules=[
        CppExtension(
            name='sparse_linear',
            sources=['sparse_linear_ops.mm'],
            extra_compile_args=['-ObjC++'],
            extra_link_args=['-framework', 'Metal', '-framework', 'Foundation']
        )
    ],
    cmdclass={'build_ext': BuildExtension}
)