File size: 6,844 Bytes
1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c 1bfc3f0 7c09a0c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | // pybind11 entry point for the ActionCodec2 native kernels.
//
// Every function here has an exactly equivalent Python implementation in the
// package. The native module is a pure accelerator: if it is not built, the
// package still works, and the test suite asserts the two agree bit for bit.
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <pybind11/stl.h>
#include <cstdint>
#include <vector>
namespace py = pybind11;
void register_resample(py::module_&);
void register_setbpe(py::module_&);
void register_inference(py::module_&);
void register_preprocessing(py::module_&);
void register_grammar(py::module_&);
void register_sparse(py::module_&);
py::array_t<int32_t> so3_second_order_encode_batch(
const py::array_t<double, py::array::c_style | py::array::forcecast>&,
const py::array_t<int32_t, py::array::c_style | py::array::forcecast>&,
const py::array_t<double, py::array::c_style | py::array::forcecast>&,
const py::array_t<double, py::array::c_style | py::array::forcecast>&,
int32_t, bool, double, double, bool, int);
py::array_t<int32_t> additive_second_order_encode_batch(
const py::array_t<float, py::array::c_style | py::array::forcecast>&,
const py::array_t<int32_t, py::array::c_style | py::array::forcecast>&,
const py::array_t<double, py::array::c_style | py::array::forcecast>&,
const py::array_t<double, py::array::c_style | py::array::forcecast>&,
const py::array_t<float, py::array::c_style | py::array::forcecast>&, bool, int);
py::tuple train_spatial_native(py::array_t<int32_t, py::array::c_style | py::array::forcecast>,
py::array_t<int64_t, py::array::c_style | py::array::forcecast>,
const std::vector<int32_t>&, int64_t, int);
py::tuple train_temporal_native(py::iterable, int64_t,
py::array_t<uint64_t, py::array::c_style | py::array::forcecast>,
py::array_t<int32_t, py::array::c_style | py::array::forcecast>,
int32_t, int32_t, int64_t, int64_t, int);
py::array_t<float> physical_decode_batch(
const py::array_t<float>&, const py::array_t<int32_t>&, const py::array_t<uint8_t>&,
const py::array_t<float>&, const py::array_t<float>&,
const py::array_t<int32_t, py::array::c_style | py::array::forcecast>&,
const py::array_t<float, py::array::c_style | py::array::forcecast>&,
const py::array_t<float, py::array::c_style | py::array::forcecast>&, int32_t, int);
py::array_t<int32_t> physical_encode_batch(
const py::array_t<float>&, const py::array_t<int32_t>&, const py::array_t<uint8_t>&,
const py::array_t<float>&, const py::array_t<float>&,
const py::array_t<float, py::array::c_style | py::array::forcecast>&,
const py::array_t<float, py::array::c_style | py::array::forcecast>&,
const py::array_t<float, py::array::c_style | py::array::forcecast>&, int32_t, int,
const py::array_t<float>&, bool);
py::array_t<int32_t> so3_encode_batch(
const py::array_t<double, py::array::c_style | py::array::forcecast>&,
const py::array_t<int32_t, py::array::c_style | py::array::forcecast>&,
const py::array_t<float, py::array::c_style | py::array::forcecast>&, int32_t, bool, bool,
double, int32_t, int);
py::tuple train_bpe_native(py::iterable, const std::vector<int32_t>&, int64_t, int64_t, int);
py::tuple spatial_encode_instances(
const py::array_t<int32_t, py::array::c_style | py::array::forcecast>&,
const py::array_t<int64_t, py::array::c_style | py::array::forcecast>&,
const std::vector<int32_t>&,
const py::array_t<int32_t, py::array::c_style | py::array::forcecast>&, int);
PYBIND11_MODULE(_native, module) {
register_resample(module);
register_setbpe(module);
register_inference(module);
register_preprocessing(module);
register_grammar(module);
register_sparse(module);
module.doc() = "ActionCodec2 native kernels (Set-BPE training and batched physical stage)";
module.def("so3_second_order_encode_batch", &so3_second_order_encode_batch,
py::arg("table"), py::arg("bins"), py::arg("quaternions"),
py::arg("previous_first_order"), py::arg("input_order"), py::arg("left"),
py::arg("limit"), py::arg("radius"), py::arg("smooth"), py::arg("threads") = 0);
module.def("additive_second_order_encode_batch", &additive_second_order_encode_batch,
py::arg("table"), py::arg("bins"), py::arg("positions"),
py::arg("previous_first_order"), py::arg("epsilon"), py::arg("smooth"),
py::arg("threads") = 0);
module.def("train_spatial", &train_spatial_native, py::arg("codes"), py::arg("weights"),
py::arg("bins"), py::arg("vocab_size"), py::arg("threads") = 0,
"Exact greedy spatial Set-BPE; returns (merges, report).");
module.def("train_temporal", &train_temporal_native, py::arg("instance_tokens"),
py::arg("total_anchors"), py::arg("spatial_support"),
py::arg("spatial_anchor"), py::arg("dimension"), py::arg("lookahead"),
py::arg("vocab_size"), py::arg("action_count"), py::arg("threads") = 0,
"Exact greedy rooted temporal Set-BPE; returns (merges, report).");
module.def("train_bpe", &train_bpe_native, py::arg("cell_chunks"), py::arg("bins"),
py::arg("vocab_size"), py::arg("max_token_cells"), py::arg("threads") = 0,
"Exact greedy BPE over the row-major flattened cell stream; returns (merges, report).");
module.def("physical_decode_batch", &physical_decode_batch, py::arg("table"),
py::arg("bins"), py::arg("binary"), py::arg("lower"), py::arg("upper"),
py::arg("cells"), py::arg("initial"), py::arg("initial_velocity"),
py::arg("order"), py::arg("threads") = 0);
module.def("physical_encode_batch", &physical_encode_batch, py::arg("table"), py::arg("bins"),
py::arg("binary"), py::arg("lower"), py::arg("upper"), py::arg("positions"),
py::arg("initial"), py::arg("initial_velocity"), py::arg("order"),
py::arg("threads") = 0, py::arg("epsilon") = py::array_t<float>(),
py::arg("smooth") = false);
module.def("so3_encode_batch", &so3_encode_batch, py::arg("table"), py::arg("bins"),
py::arg("actions"), py::arg("input_order"), py::arg("rotation_vector"),
py::arg("left"), py::arg("limit"), py::arg("selection"), py::arg("threads") = 0,
"Batched exact SO(3) error-feedback encoding; selection 0 nearest, 1 smooth, "
"2 momentum_lazy.");
module.def("spatial_encode_instances", &spatial_encode_instances, py::arg("codes"),
py::arg("row_anchors"), py::arg("bins"), py::arg("merges"),
py::arg("threads") = 0);
}
|