File size: 659 Bytes
bfdbd33 0a16898 bfdbd33 0a16898 bfdbd33 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 | // JIT binding for load_local.py; the published build registers via torch_binding.cpp.
#include <torch/extension.h>
void physarum_step(torch::Tensor trail, torch::Tensor tmp,
torch::Tensor ax, torch::Tensor ay, torch::Tensor ah,
int64_t seed, double sense_dist, double sense_ang,
double turn, double speed, double deposit, double decay);
void flow_cg(torch::Tensor cE, torch::Tensor cS, torch::Tensor b,
torch::Tensor p, int64_t ground, int64_t max_iters, double tol);
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("physarum_step", &physarum_step);
m.def("flow_cg", &flow_cg);
}
|