static void fmac_apply(const linop_data_t* _data, complex float* dst, const complex float* src)
{
auto data = CAST_DOWN(fmac_data, _data);
#ifdef USE_CUDA
const complex float* tensor = get_tensor(data, cuda_ondevice(src));
#else
const complex float* tensor = data->tensor;
#endif
md_clear2(data->N, data->odims, data->ostrs, dst, CFL_SIZE);
md_zfmac2(data->N, data->dims, data->ostrs, dst, data->istrs, src, data->tstrs, tensor);
}
PYBIND11_MODULE(ead, m)
{
m.doc() = "ead variables";
// ==== node ====
py::object node = (py::object)
py::module::import("ead.age").attr("NodeptrT<PybindT>");
((py::class_<ead::iNode<PybindT>,ead::NodeptrT<PybindT>>) node)
.def("__str__",
[](py::object self)
{
auto dnode = self.cast<ead::iNode<PybindT>*>();
return dnode->get_tensor()->to_string();
},
"Return string representation of this tensor instance")
.def("shape",
[](py::object self)
{
auto dnode = self.cast<ead::iNode<PybindT>*>();
ade::Shape shape = dnode->get_tensor()->shape();
auto pshape = pyead::c2pshape(shape);
std::vector<int> ipshape(pshape.begin(), pshape.end());
return py::array(ipshape.size(), ipshape.data());
},
"Return this instance's shape")
.def("children",
[](py::object self)
{
auto dnode = self.cast<ead::iNode<PybindT>*>();
std::vector<ade::TensptrT> tens;
if (auto f = dynamic_cast<ade::iFunctor*>(
dnode->get_tensor().get()))
{
auto args = f->get_children();
std::transform(args.begin(), args.end(),
std::back_inserter(tens),
[](ade::FuncArg& mten)
{
return mten.get_tensor();
});
}
return tens;
})
.def("as_tens",
[](py::object self)
{
auto dnode = self.cast<ead::iNode<PybindT>*>();
return dnode->get_tensor();
})
.def("get",
[](py::object self)
{
auto dnode = self.cast<ead::iNode<PybindT>*>();
return pyead::typedata_to_array<PybindT>(dnode,
py::dtype::of<PybindT>());
});
// ==== session ====
py::class_<ead::Session<PybindT>> session(m, "Session");
session
.def(py::init())
.def("track", &ead::Session<PybindT>::track, "Track node")
.def("update",
[](py::object self, std::vector<ead::NodeptrT<PybindT>> nodes)
{
auto sess = self.cast<ead::Session<PybindT>*>();
std::unordered_set<ade::iTensor*> updates;
for (ead::NodeptrT<PybindT>& node : nodes)
{
updates.emplace(node->get_tensor().get());
}
sess->update(updates);
},
"Return calculated data",
py::arg("nodes") = std::vector<ead::NodeptrT<PybindT>>{});
// ==== constant ====
py::class_<ead::ConstantNode<PybindT>,std::shared_ptr<ead::ConstantNode<PybindT>>> constant(
m, "Constant", node);
py::implicitly_convertible<ead::iNode<PybindT>,ead::ConstantNode<PybindT>>();
// ==== variable ====
py::class_<ead::VariableNode<PybindT>,ead::VarptrT<PybindT>> variable(
m, "Variable", node);
py::implicitly_convertible<ead::iNode<PybindT>,ead::VariableNode<PybindT>>();
variable
.def("assign",
[](py::object self, py::array data)
{
auto var = self.cast<ead::VariableNode<PybindT>*>();
ade::Shape shape;
std::vector<PybindT> vec = pyead::arr2vec(shape, data);
var->assign(vec.data(), shape);
},
"Assign numpy data array to variable");
// ==== inline functions ====
m.def("scalar_constant", &ead::make_constant_scalar<PybindT>,
"Return scalar constant node");
m.def("constant",
[](py::array data)
{
ade::Shape shape;
std::vector<PybindT> vec = pyead::arr2vec(shape, data);
return ead::make_constant(vec.data(), shape);
}, "Return constant node with data");
m.def("scalar_variable", &ead::make_variable_scalar<PybindT>,
"Return labelled variable containing numpy data array",
py::arg("scalar"), py::arg("shape"), py::arg("label") = "");
m.def("variable",
[](py::array data, std::string label)
{
ade::Shape shape;
std::vector<PybindT> vec = pyead::arr2vec(shape, data);
return ead::make_variable(vec.data(), shape, label);
},
"Return labelled variable containing numpy data array",
py::arg("data"), py::arg("label") = "");
m.def("derive", &ead::derive<PybindT>,
"Return derivative of first tensor with respect to second tensor (deprecated)");
m.def("seed",
[](size_t seed)
{
ead::get_engine().seed(seed);
},
"Seed internal RNG");
}
ade::Shape shape (void)
{
return get_tensor()->shape();
}
