#include <torch/library.h>

#include "registration.h"
#include "torch_binding.h"

TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.def("physarum_step(Tensor! trail, Tensor! tmp, Tensor! ax, Tensor! ay, "
          "Tensor! ah, int seed, float sense_dist, float sense_ang, float turn, "
          "float speed, float deposit, float decay) -> ()");
  ops.def("flow_cg(Tensor cE, Tensor cS, Tensor b, Tensor! p, int ground, "
          "int max_iters, float tol) -> ()");
#if defined(CPU_KERNEL)
  ops.impl("physarum_step", torch::kCPU, &physarum_step);
  ops.impl("flow_cg", torch::kCPU, &flow_cg);
#endif
}

REGISTER_EXTENSION(TORCH_EXTENSION_NAME)