#pragma once

// Local build compatibility shim for kernel-builder's registration.h

#include <torch/extension.h>
#include <torch/library.h>

// TORCH_LIBRARY_EXPAND may not be defined in all PyTorch versions
#ifndef TORCH_LIBRARY_EXPAND
#define TORCH_LIBRARY_EXPAND(ns, m) TORCH_LIBRARY(ns, m)
#endif

// Generate the PyInit_<name> entry point for the shared library
#define REGISTER_EXTENSION(name)                                               \
  PYBIND11_MODULE(name, m) {}