added conditions for merge and split

Luecx · Luecx · commit e92b7c5fbaf7 · 2023-11-09T07:29:54.000+01:00
diff --git a/src/nn/layers/activation.h b/src/nn/layers/activation.h
@@ -39,6 +39,41 @@ struct ReLU : public Layer {
     }
 };
 
+struct Linear : public Layer {
+    Layer*            prev;
+    GradientOperation grad_op;
+    int               use_id;
+    float             scalar;
+    explicit Linear(Layer* prev, float scalar = 1)
+        : Layer(prev->size)
+        , prev(prev)
+        , scalar(scalar) {
+        prev->use();
+        use_id = prev->use();
+    }
+
+    void compile(size_t batch_size) override {
+        this->compile_suboutput(batch_size, Tape {size, batch_size});
+        this->grad_op = use_id == prev->used() ? SET : INCREMENT;
+    }
+
+    void forward() override {
+        Layer::forward();
+        operations::linear<data::GPU>(prev->dense_output.values, dense_output.values, scalar);
+    }
+
+    void backward() override {
+        Layer::backward();
+        operations::linear_bp<data::GPU>(prev->dense_output.values,
+                                          prev->dense_output.gradients,
+                                          dense_output.values,
+                                          dense_output.gradients,
+                                          scalar,
+                                          grad_op);
+    }
+};
+
+
 struct Sigmoid : public Layer {
     Layer*            prev;
     GradientOperation grad_op;
diff --git a/src/nn/layers/merge.h b/src/nn/layers/merge.h
@@ -17,15 +17,20 @@ struct Merge : public Layer {
     Layer* l1;
     Layer* l2;
 
+    int use_id_1;
+    int use_id_2;
+
     Merge(Layer* l1, Layer* l2)
         : Layer(l1->size + l2->size)
         , l1(l1)
         , l2(l2) {
-        l1->use();
-        l2->use();}
+        use_id_1 = l1->use();
+        use_id_2 = l2->use();}
 
     void compile(size_t batch_size) override {
         this->compile_suboutput(batch_size, Tape(size, batch_size));
+        ERROR(use_id_1 == l1->used());
+        ERROR(use_id_2 == l2->used());
     }
 
     void compile_suboutput(size_t batch_size, const Tape& output) override {
diff --git a/src/nn/layers/split.h b/src/nn/layers/split.h
@@ -43,9 +43,12 @@ struct SplitHead : public Layer {
 
 struct Split : public Layer {
     std::vector<SplitHead> heads;
+    Layer* prev;
+    int use_id;
+
 
     explicit Split(Layer* prev, const std::vector<size_t>& head_sizes)
-        : Layer(prev->size) {
+        : Layer(prev->size), prev(prev) {
         size_t total_assigned_size = 0;
         for (size_t i = 0; i < head_sizes.size(); ++i) {
             heads.emplace_back(prev, head_sizes[i], total_assigned_size);
@@ -55,7 +58,7 @@ struct Split : public Layer {
         size_t remaining_size = prev->size - total_assigned_size;
         heads.emplace_back(prev, remaining_size, total_assigned_size);
 
-        prev->use();
+        use_id = prev->use();
     }
 
     SplitHead* operator[](size_t index) {
@@ -66,6 +69,7 @@ struct Split : public Layer {
     }
 
     void compile(size_t batch_size) override {
+        ERROR(use_id == prev->used()); // make sure the operation would be "set"
         for (auto& head : heads) {
             head.compile(batch_size);
         }
diff --git a/src/nn/layers/weighted_sum.h b/src/nn/layers/weighted_sum.h
@@ -43,7 +43,9 @@ struct WeightedSum : public Layer {
                                           prev_2->dense_output.gradients,
                                           dense_output.gradients,
                                           alpha,
-                                          beta);
+                                          beta,
+                                          grad_op_1,
+                                          grad_op_2);
     }
 };
 
diff --git a/src/operations/activation/activations.h b/src/operations/activation/activations.h
@@ -6,8 +6,8 @@
 namespace operations {
 
 DEFINE_ACTIVATION(linear
-                  , A[ida];
-                  , B_grd[idb];);
+                  , A[ida] * scalar;
+                  , B_grd[idb] * scalar;);
 
 DEFINE_ACTIVATION(sigmoid
                   , 1.0f / (1 + exp(-A[ida] * scalar));
