add Physarum flow solver (shortest path / network); card leads with it

0a16898 verified 17 days ago

655 Bytes

	#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)