#include <torch/library.h>

#include "registration.h"
#include "torch_binding.h"

TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.def("apply_token_bitmask_inplace(Tensor(a!) logits, Tensor bitmask) -> ()");
  ops.def("sample(Tensor logits, Tensor temperature, Tensor top_p, Tensor top_k, "
          "Tensor min_p, int seed, int offset) -> Tensor");

#if defined(CUDA_KERNEL) || defined(ROCM_KERNEL)
  ops.impl("apply_token_bitmask_inplace", torch::kCUDA, &apply_token_bitmask_inplace);
  ops.impl("sample", torch::kCUDA, &sample);
#endif
}

REGISTER_EXTENSION(TORCH_EXTENSION_NAME)