Skip to content
Snippets Groups Projects
Commit 8e7fa67f authored by Maxence Naud's avatar Maxence Naud
Browse files

Upd compile_gradient function

parent e40ef693
No related branches found
No related tags found
No related merge requests found
...@@ -12,11 +12,12 @@ ...@@ -12,11 +12,12 @@
#include <memory> #include <memory>
#include <set> #include <set>
#include "aidge/data/Tensor.hpp"
#include "aidge/graph/Node.hpp" #include "aidge/graph/Node.hpp"
#include "aidge/graph/GraphView.hpp" #include "aidge/graph/GraphView.hpp"
#include "aidge/operator/OperatorTensor.hpp" #include "aidge/operator/OperatorTensor.hpp"
#include "aidge/utils/ErrorHandling.hpp" #include "aidge/utils/ErrorHandling.hpp"
#include "aidge/recipies/GraphViewHelper.hpp" #include "aidge/recipes/GraphViewHelper.hpp"
std::set<std::shared_ptr<Aidge::Tensor>> Aidge::producers(std::shared_ptr<Aidge::GraphView> graphview) { std::set<std::shared_ptr<Aidge::Tensor>> Aidge::producers(std::shared_ptr<Aidge::GraphView> graphview) {
...@@ -47,12 +48,10 @@ std::set<std::shared_ptr<Aidge::Tensor>> Aidge::parameters(std::shared_ptr<Aidge ...@@ -47,12 +48,10 @@ std::set<std::shared_ptr<Aidge::Tensor>> Aidge::parameters(std::shared_ptr<Aidge
void Aidge::compile_gradient(std::shared_ptr<Aidge::GraphView> gv) { void Aidge::compile_gradient(std::shared_ptr<Aidge::GraphView> gv) {
for (const auto& node : gv->getNodes()) { for (const auto& node : gv->getNodes()) {
// TODO: check that each node is an OperatorTensor // TODO: check that each node is an OperatorTensor
AIDGE_ASSERT(node->getOperator()->operatorType() == OperatorType::Tensor, "Cannot instanciate gradient of an Operator that doesn't use Tensor."); AIDGE_ASSERT(node->getOperator()->operatorType() == OperatorType::Tensor, "Cannot instanciate gradient of an Operator ({}) that doesn't use Tensor.", node->getOperator()->type());
const std::shared_ptr<OperatorTensor> op = std::static_pointer_cast<OperatorTensor>(node -> getOperator()); const std::shared_ptr<OperatorTensor> op = std::dynamic_pointer_cast<OperatorTensor>(node -> getOperator());
for (std::size_t o = 0; o < node -> nbOutputs(); ++o) { for (std::size_t o = 0; o < node -> nbOutputs(); ++o) {
const auto& t = op->getOutput(o); op->getOutput(o)->initGradient();
t -> grad() -> setDataType(t -> dataType());
t -> grad() -> setBackend(t -> getImpl() -> backend());
} }
} }
} }
\ No newline at end of file
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment