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);
}