Fixed some bugs and failed assertions that only appeared in debug mode

Author: RobBa

src/backend/computational_graph/elementwise_mul_node.cpp

@@ -15,6 +15,7 @@ using namespace std;
 using namespace graph;
 
 vector<shared_ptr<Tensor>> ElementwiseMulNode::backward(const Tensor& upstreamGrad) {
+    assert(!upstreamGrad.getRequiresGrad());
     return {
         make_shared<Tensor>(upstreamGrad * (*parents[1])), 
         make_shared<Tensor>(upstreamGrad * (*parents[0]))

src/backend/computational_graph/graph_creation.cpp

@@ -21,35 +21,35 @@ using namespace std;
 shared_ptr<Tensor> graph::mul(const shared_ptr<Tensor> left, const shared_ptr<Tensor> right) {
   auto res = make_shared<Tensor>((*left) * (*right));
   if(left->getRequiresGrad() || right->getRequiresGrad()){
-    assert(res->getRequiresGrad());
     res->setCgNode(make_shared<graph::ElementwiseMulNode>(left, right));
+    assert(res->getRequiresGrad());
   }
   return res;
 }
 
 shared_ptr<Tensor> graph::add(const shared_ptr<Tensor> left, const shared_ptr<Tensor> right) {
   auto res = make_shared<Tensor>(*left + *right);
   if(left->getRequiresGrad() || right->getRequiresGrad()){
-    assert(res->getRequiresGrad());
     res->setCgNode(make_shared<graph::AddNode>(left, right));
+    assert(res->getRequiresGrad());
   }
   return res;
 }
 
 shared_ptr<Tensor> graph::matmul(const shared_ptr<Tensor> left, const shared_ptr<Tensor> right) {
   auto res = make_shared<Tensor>(left->matmul(*right));
   if(left->getRequiresGrad() || right->getRequiresGrad()){
-    assert(res->getRequiresGrad());
     res->setCgNode(make_shared<graph::MatMulNode>(left, right));
+    assert(res->getRequiresGrad());
   }
   return res;
 }
 
 shared_ptr<Tensor> graph::mul(const shared_ptr<Tensor> t, ftype scalar) {
   auto res = make_shared<Tensor>((*t) * scalar);
   if(t->getRequiresGrad()){
-    assert(res->getRequiresGrad());
     res->setCgNode(std::make_shared<graph::ScalarMulNode>(t, scalar));
+    assert(res->getRequiresGrad());
   }
   return res;
 }
@@ -61,8 +61,8 @@ shared_ptr<Tensor> graph::mul(ftype scalar, const shared_ptr<Tensor> t) {
 shared_ptr<Tensor> graph::add(const shared_ptr<Tensor> t, ftype scalar) {
   auto res = make_shared<Tensor>((*t) + scalar);
   if(t->getRequiresGrad()){
-    assert(res->getRequiresGrad());
     res->setCgNode(std::make_shared<graph::ScalarAddNode>(t));
+    assert(res->getRequiresGrad());
   }
   return res;
 }
@@ -74,17 +74,67 @@ shared_ptr<Tensor> graph::add(ftype scalar, const shared_ptr<Tensor> t) {
 shared_ptr<Tensor> graph::sub(const shared_ptr<Tensor> t, ftype scalar) {
   auto res = make_shared<Tensor>((*t) - scalar);
   if(t->getRequiresGrad()){
-    assert(res->getRequiresGrad());
     res->setCgNode(std::make_shared<graph::ScalarAddNode>(t));
+    assert(res->getRequiresGrad());
   }
   return res;
 }
 
 shared_ptr<Tensor> graph::div(const shared_ptr<Tensor> t, ftype scalar) {
   auto res = make_shared<Tensor>((*t) / scalar);
   if(t->getRequiresGrad()){
-    assert(res->getRequiresGrad());
     res->setCgNode(std::make_shared<graph::ScalarMulNode>(t, 1 / scalar));
+    assert(res->getRequiresGrad());
   }
   return res;
+}
+
+/**
+ * @brief Special linear indexing, see getItem() overloads in tensor. 
+ * Used to keep the computational graph intact.
+ * E.g. if we have something like 
+ * 
+ * loss = loss + other.get(i), we need to make sure get(i) can map to computational graph.
+ */
+shared_ptr<Tensor> graph::getAsShared(const shared_ptr<Tensor>& t, tensorSize_t idx) {
+  ftype val = t->getItem(idx);
+  return make_shared<Tensor>(std::vector<tensorDim_t>{1}, std::vector<ftype>{val}, 
+                             t->getDevice(), t->getRequiresGrad()); 
+}
+
+/**
+ * @brief Special linear indexing, see getItem() overloads in tensor. 
+ * Used to keep the computational graph intact.
+ * E.g. if we have something like 
+ * 
+ * loss = loss + other.get(i), we need to make sure get(i) can map to computational graph.
+ */
+std::shared_ptr<Tensor> graph::getAsShared(const Tensor& t, tensorSize_t idx) {
+  ftype val = t.getItem(idx);
+  return make_shared<Tensor>(std::vector<tensorDim_t>{1}, std::vector<ftype>{val}, 
+                             t.getDevice(), t.getRequiresGrad()); 
+}
+
+/**
+ * @brief Used to keep the computational graph intact.
+ * E.g. if we have something like 
+ * 
+ * loss = loss + other.get(i), we need to make sure get(i) can map to computational graph.
+ */
+shared_ptr<Tensor> graph::getAsShared(const shared_ptr<Tensor>& t, vector<tensorDim_t>&& idx) {
+  ftype val = t->getItem(std::move(idx));
+  return make_shared<Tensor>(std::vector<tensorDim_t>{1}, std::vector<ftype>{val}, 
+                             t->getDevice(), t->getRequiresGrad()); 
+}
+
+/**
+ * @brief Used to keep the computational graph intact.
+ * E.g. if we have something like 
+ * 
+ * loss = loss + other.get(i), we need to make sure get(i) can map to computational graph.
+ */
+std::shared_ptr<Tensor> graph::getAsShared(const Tensor& t, std::vector<tensorDim_t>&& idx) {
+  ftype val = t.getItem(std::move(idx));
+  return make_shared<Tensor>(std::vector<tensorDim_t>{1}, std::vector<ftype>{val}, 
+                             t.getDevice(), t.getRequiresGrad()); 
 }

src/backend/computational_graph/graph_creation.h

@@ -29,6 +29,12 @@ namespace graph {
   std::shared_ptr<Tensor> add(ftype scalar, const std::shared_ptr<Tensor> left);
 
   std::shared_ptr<Tensor> sub(const std::shared_ptr<Tensor> left, ftype scalar);
-  std::shared_ptr<Tensor> div(const std::shared_ptr<Tensor> left, ftype scalar);  
+  std::shared_ptr<Tensor> div(const std::shared_ptr<Tensor> left, ftype scalar);
+
+  std::shared_ptr<Tensor> getAsShared(const std::shared_ptr<Tensor>& t, tensorSize_t idx);
+  std::shared_ptr<Tensor> getAsShared(const Tensor& t, tensorSize_t idx);
+
+  std::shared_ptr<Tensor> getAsShared(const std::shared_ptr<Tensor>& t, std::vector<tensorDim_t>&& idx);
+  std::shared_ptr<Tensor> getAsShared(const Tensor& t, std::vector<tensorDim_t>&& idx);
 }
 

src/backend/computational_graph/matmul_node.cpp

@@ -15,6 +15,7 @@ using namespace std;
 using namespace graph;
 
 vector<shared_ptr<Tensor>> MatMulNode::backward(const Tensor& upstreamGrad) {
+    assert(!upstreamGrad.getRequiresGrad());
     return {
         make_shared<Tensor>(upstreamGrad.matmul(parents[1]->transpose(-2, -1))), 
         make_shared<Tensor>(parents[0]->transpose(-2, -1).matmul(upstreamGrad))

src/backend/computational_graph/relu_node.cpp

@@ -17,12 +17,14 @@ using namespace std;
 using namespace graph;
 
 vector<shared_ptr<Tensor>> ReLuNode::backward(const Tensor& upstreamGrad) {
+  assert(!upstreamGrad.getRequiresGrad());
+
   constexpr ftype zero = 0.0;
 
   auto res = make_shared<Tensor>(upstreamGrad.getDims().toVector(), upstreamGrad.getDevice(), false);
   for(tensorSize_t i=0; i<upstreamGrad.getSize(); i++){
-    auto v = upstreamGrad.get(i);
-    res->set(v > zero ? v : zero, i);
+    auto v = upstreamGrad.getItem(i);
+    res->setItem(v > zero ? v : zero, i);
   }
   return {std::move(res)};
 }

src/backend/computational_graph/scalar_op_nodes.cpp

@@ -17,13 +17,16 @@ using namespace std;
 using namespace graph;
 
 vector<shared_ptr<Tensor>> graph::ScalarAddNode::backward(const Tensor& upstreamGrad) {
+  assert(!upstreamGrad.getRequiresGrad());
   return {make_shared<Tensor>(upstreamGrad.createDeepCopy())};
 }
 
 vector<shared_ptr<Tensor>> graph::ScalarMulNode::backward(const Tensor& upstreamGrad) {
+  assert(!upstreamGrad.getRequiresGrad());
+
   auto res = make_shared<Tensor>(upstreamGrad.createDeepCopy());
   for(tensorSize_t i=0; i<res->getSize(); i++){
-    res->set(res->get(i) * factor, i);
+    res->setItem(res->getItem(i) * factor, i);
   }
   return {std::move(res)};
 }

src/backend/data_modeling/dim_type.cpp

@@ -22,7 +22,7 @@ tensorDim_t Dimension::multVector(const std::vector<tensorDim_t>& dims) const no
 #ifndef NDEBUG
   utility::SafeArithmetics_t<tensorSize_t> mult(1);
   for(auto dim: dims){
-    mult * dim;
+    mult = mult * dim;
   }
 
   res = mult.value;
@@ -85,7 +85,7 @@ Dimension& Dimension::operator=(Dimension&& other) noexcept {
 ostream& operator<<(ostream& os, const Dimension& d) noexcept {
   os << "(";
   for(int i=0; i<d.nDims(); i++){
-    os << d.get(i);
+    os << d.getItem(i);
 
     if(i+1<d.nDims()){
       os << ",";

src/backend/data_modeling/dim_type.h

@@ -54,7 +54,7 @@ class Dimension final {
       return size;
     }
 
-    tensorDim_t get(int idx) const {
+    tensorDim_t getItem(int idx) const {
       assert(size!=0);
       if(idx<0){
         idx = dims.size() + idx; // -1 is last idx, -2 second last and so forth