replicate/megablocks

1#include "cumsum.h"2#include "histogram.h"3#include "indices.h"4#include "replicate.h"5#include "sort.h"6 7#include <torch/extension.h>8 9namespace megablocks {10 11PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {12  m.def("exclusive_cumsum", &exclusive_cumsum, "batched exclusive cumsum.");13  m.def("histogram", &histogram, "even width histogram.");14  m.def("inclusive_cumsum", &inclusive_cumsum, "batched inclusive cumsum");15  m.def("indices", &indices, "indices construction for sparse matrix.");16  m.def("replicate_forward", &replicate_forward, "(fwd) replicate a vector dynamically.");17  m.def("replicate_backward", &replicate_backward, "(bwd) replicate a vector dynamically.");18  m.def("sort", &sort, "key/value sort.");19}20 21}  // namespace megablocks22