Function PYBIND11_MODULE

inference_lib/src/aqlm/inference_kernels/cuda_kernel.cpp:686–699

}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("code1x16_matmat", &code1x16_matmat, "1x16 (2bit) codebook matrix-matrix product through matvec.");
  m.def("code1x16_dequant", &code1x16_dequant, "1x16 (2bit) codebook dequantization.");
  m.def("code1x16_matmat_dequant", &code1x16_matmat_dequant, "1x16 (2bit) codebook matrix-matrix dequantization product.");
  m.def("code1x16_matmat_dequant_transposed", &code1x16_matmat_dequant_transposed, "1x16 (2bit) codebook matrix-matrix dequantization product for backward pass.");
  m.def("code2x8_matmat", &code2x8_matmat, "2x8 (2bit) codebook matrix-matrix product.");
  m.def("code2x8_dequant", &code2x8_dequant, "2x8 (2bit) codebook dequantization.");
  m.def("code2x8_matmat_dequant", &code2x8_matmat_dequant, "2x8 (2bit) codebook matrix-matrix dequantization product.");
  m.def("code2x8_matmat_dequant_transposed", &code2x8_matmat_dequant_transposed, "2x8 (2bit) codebook matrix-matrix dequantization product for backward pass.");
  m.def("code1x8_matmat", &code1x8_matmat, "1x8 (1bit) codebook matrix-matrix product.");
  m.def("code1x8_dequant", &code1x8_dequant, "1x8 (1bit) codebook dequantization.");
  m.def("code1x8_matmat_dequant", &code1x8_matmat_dequant, "1x8 (1bit) codebook matrix-matrix dequantization product.");
  m.def("code1x8_matmat_dequant_transposed", &code1x8_matmat_dequant_transposed, "1x8 (1bit) codebook matrix-matrix dequantization product for backward pass.");
}