Digital Drosophila Archive
19 lines
1.2 KiB
C++
19 lines
1.2 KiB
C++
#include <torch/extension.h>
|
|
|
|
void author_state_cuda(torch::Tensor v, torch::Tensor g, torch::Tensor ref,
|
|
torch::Tensor requests, torch::Tensor spikes, torch::Tensor counts,
|
|
double gd, double c0, double cg, double cv, double threshold,
|
|
double reset, double seed_volt, int64_t refractory_steps);
|
|
void author_fill_edges_cuda(torch::Tensor indptr, torch::Tensor sources,
|
|
torch::Tensor offsets, torch::Tensor output);
|
|
void author_arrivals_cuda(torch::Tensor edges, torch::Tensor starts, torch::Tensor targets,
|
|
torch::Tensor weights, torch::Tensor g, torch::Tensor ref);
|
|
void author_external_cuda(torch::Tensor external, torch::Tensor g, torch::Tensor ref);
|
|
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
|
m.def("state_step", &author_state_cuda, "Author NumPy-ordered integration/threshold/reset");
|
|
m.def("fill_edges", &author_fill_edges_cuda, "Expand only active-source canonical CSR edges");
|
|
m.def("apply_arrivals", &author_arrivals_cuda, "Sequential per-target arrival additions");
|
|
m.def("apply_external", &author_external_cuda, "Post-anatomical external arrival additions");
|
|
}
|