diff --git a/error.txt b/error.txt
new file mode 100644
index 0000000..cd529e6
--- /dev/null
+++ b/error.txt
@@ -0,0 +1,44 @@
+-- Configuring done (0.3s)
+-- Generating done (0.0s)
+-- Build files have been written to: /home/wslarch/Documents/Coding/DEV/soft/soft-cuda/build
+
+real	0m0.356s
+user	0m0.157s
+sys	0m0.181s
+[  5%] Built target gtest
+[ 10%] Building CXX object CMakeFiles/soft_lib.dir/src/backend_cpu/base/debug.cpp.o
+[ 10%] Building CXX object CMakeFiles/soft_lib.dir/src/backend_cpu/math/bias_add.cpp.o
+[ 13%] Building CXX object CMakeFiles/soft_lib.dir/src/backend_cpu/math/mean.cpp.o
+[ 16%] Building CXX object CMakeFiles/soft_lib.dir/src/backend_cpu/math/mse.cpp.o
+[ 18%] Building CXX object CMakeFiles/soft_lib.dir/src/backend_cpu/math/mul.cpp.o
+[ 21%] Building CXX object CMakeFiles/soft_lib.dir/src/backend_cpu/math/relu.cpp.o
+[ 24%] Building CXX object CMakeFiles/soft_lib.dir/src/backend_cpu/math/add.cpp.o
+[ 29%] Building CXX object CMakeFiles/soft_lib.dir/src/backend_cpu/backprop/backprop_b.cpp.o
+[ 29%] Building CXX object CMakeFiles/soft_lib.dir/src/backend_cpu/math/scalar.cpp.o
+[ 32%] Building CXX object CMakeFiles/soft_lib.dir/src/backend_cpu/math/square.cpp.o
+[ 35%] Building CXX object CMakeFiles/soft_lib.dir/src/backend_cpu/math/transpose.cpp.o
+[ 37%] Building CXX object CMakeFiles/soft_lib.dir/src/backend_cpu/math/sub.cpp.o
+[ 40%] Building CUDA object CMakeFiles/soft_lib.dir/src/backend_gpu/math/matmul.cu.o
+[ 43%] Building CUDA object CMakeFiles/soft_lib.dir/src/backend_gpu/math/add.cu.o
+[ 48%] Built target gtest_main
+[ 54%] Built target gmock
+[ 56%] Building CXX object CMakeFiles/soft_lib.dir/src/core/JSON/json_utils.cpp.o
+[ 62%] Built target gmock_main
+[ 64%] Building CXX object CMakeFiles/soft_lib.dir/src/core/graph/DAGbuild.cpp.o
+[ 67%] Building CUDA object CMakeFiles/soft_lib.dir/src/core/graph/assignBackend.cu.o
+[ 70%] Building CXX object CMakeFiles/soft_lib.dir/src/core/graph/train.cpp.o
+[ 72%] Building CUDA object CMakeFiles/soft_lib.dir/src/core/pool/pool.cu.o
+[ 75%] Building CXX object CMakeFiles/soft_lib.dir/src/core/tensor/tensor.cpp.o
+/home/wslarch/Documents/Coding/DEV/soft/soft-cuda/src/core/tensor/tensor.cpp: In function ‘bool tensor_evaluate(tensor_pool_t*, tensor_t*, float*, float*, float*)’:
+/home/wslarch/Documents/Coding/DEV/soft/soft-cuda/src/core/tensor/tensor.cpp:176:19: error: ‘tensor_op_square’ was not declared in this scope; did you mean ‘tensor_square’?
+  176 |         success = tensor_op_square(pool, t);
+      |                   ^~~~~~~~~~~~~~~~
+      |                   tensor_square
+make[2]: *** [CMakeFiles/soft_lib.dir/build.make:363: CMakeFiles/soft_lib.dir/src/core/tensor/tensor.cpp.o] Error 1
+make[2]: *** Waiting for unfinished jobs....
+make[1]: *** [CMakeFiles/Makefile2:181: CMakeFiles/soft_lib.dir/all] Error 2
+make: *** [Makefile:146: all] Error 2
+
+real	0m11.154s
+user	1m11.334s
+sys	0m16.624s
diff --git a/include/soft-cuda/tensor/api.h b/include/soft-cuda/tensor/api.h
index befb6e9..9ab3db1 100644
--- a/include/soft-cuda/tensor/api.h
+++ b/include/soft-cuda/tensor/api.h
@@ -84,7 +84,7 @@ uint32_t tensor_id(tensor_t *t);
  */
 // DONE
 tensor_t *tensor_create(tensor_pool_t *pool, tensor_dtype_t dtype, uint32_t num_dims,
-                        uint32_t *dims, void *elems);
+                        uint32_t *dims, void *elems, bool grad_status = true);
 
 /*
  * Establish a new memory arena
@@ -186,6 +186,7 @@ void tensor_print_data(tensor_t *t);
 // DONE
 tensor_t *tensor_mul(tensor_pool_t *pool, tensor_t *x, tensor_t *y);
 
+tensor_t *tensor_square(tensor_pool_t *pool, tensor_t *x);
 /* The naive version of matrix multiplication
  * @param pool       Pointer to the tensor pool
  * @param x            Pointer to the tensor which will be mulptiplied
@@ -221,6 +222,16 @@ tensor_t *tensor_add(tensor_pool_t *pool, tensor_t *x, tensor_t *y);
 // DONE
 tensor_t *tensor_add_bias(tensor_pool_t *pool, const tensor_t *xw, const tensor_t *bias);
 
+/*
+ * Do matrix subtraction
+ * @param out          Pointer to the tensor where result will be stored
+ * @param x            Pointer to the tensor which will be subtracted from 
+ * @param y            Pointer to the tensor which will be subtracted
+ *
+ * @return             Returns a tensor object with operation set.
+ * */
+// DONE
+tensor_t *tensor_sub(tensor_pool_t *pool, tensor_t *a, tensor_t *b);
 /////////////////////////////////////////////////////////////
 /// DEPRECEATED tensor_mul operation handles it automatically
 // /*
@@ -242,23 +253,25 @@ tensor_t *tensor_add_bias(tensor_pool_t *pool, const tensor_t *xw, const tensor_
 // DONE
 tensor_t *tensor_relu(tensor_pool_t *pool, tensor_t *a);
 
+// DONE
+tensor_t *tensor_mean(tensor_pool_t *pool, tensor_t *a);
 // Compares result
 // return how correct we were b/w 0-1
 tensor_t *tensor_mse_loss(tensor_pool_t *pool, tensor_t *predictions, tensor_t *target);
 
-// Fills an existing tensor with normally distributed random numbers.
+// // Fills an existing tensor with normally distributed random numbers.
 bool tensor_fill_random_normal(tensor_t *t, float mean, float std_dev);
-
-// Fused operation combining Softmax and Cross-Entropy for stability
-tensor_t *tensor_cross_entropy_loss(tensor_pool_t *pool, const tensor_t *predictions,
-                                    const tensor_t *targets);
+//
+// // Fused operation combining Softmax and Cross-Entropy for stability
+// tensor_t *tensor_cross_entropy_loss(tensor_pool_t *pool, const tensor_t *predictions,
+//                                     const tensor_t *targets);
 
 // ***********************************************************************************
 // TODO: HAVE TO UPDATED FUNC SIG SPEC 
 // Evalutes the operation(Forward) with depth=1
 // @return             boolean flag for status
 // DONE
-bool tensor_evaluate( tensor_pool_t *pool,tensor_t *t,  float *d_a, float *d_b, float *d_res);
+bool tensor_evaluate( tensor_pool_t *pool,tensor_t *t,  float *d_a = nullptr, float *d_b = nullptr, float *d_res= nullptr);
 // DONE
 bool tensor_evaluate_GPU( tensor_pool_t *pool,tensor_t *t,  float *d_a, float *d_b, float *d_res);
 
@@ -273,9 +286,7 @@ bool tensor_evaluate_GPU( tensor_pool_t *pool,tensor_t *t,  float *d_a, float *d
 bool tensor_backward(tensor_pool_t *pool, tensor_t *t);
 
 // The Optimizer
-void tensor_sgd_template(tensor_pool_t *static_weights_pool, double learning_rate);
-
-// API UNSTABLE
+void tensor_sgd(std::vector<execution_node_t *> &nodes, float learning_rate);
 /////////////////////////////////////////////////////////////
 // GRAPH OPERATIONS
 
@@ -321,6 +332,7 @@ bool verifyIfDAG(tensor_pool_t *pool, tensor_t *t, std::vector<execution_node_t
 // DONE
 void assignBackendGraph(tensor_pool_t *pool,std::vector<execution_node_t *> &nodes);
 
+void assignGradMemory(tensor_pool_t *pool_grad_cpu, tensor_pool_t *pool_grad_gpu, std::vector<execution_node_t *> &nodes);
 /* @params  Take execution_node_t which you wanna know
  * @returns the postion of the execution_node_t after verifyIfDAG
  * */
@@ -340,14 +352,15 @@ void printExecutionNode(execution_node_t *et);
  * */
 bool tensor_graph_forward_evaluate(tensor_pool_t *pool_cpu, tensor_pool_t *pool_gpu, std::vector<execution_node_t *> &nodes);
 
-/*
- * Evaluate the whole graph backward operation.
- * @params            graph struct for ops sequence
- * @return            boolean status flag
- * */
-bool tensor_graph_backward_evaluate(tensor_graph_t *g);
-
 //////////////////////////////////////////////////////////////////////////
+/////////////////////////////////////////////////////////////////////////
+// BACKWARD PASS FUNCTIONS
+
+void gradInitializer(std::vector<execution_node_t *> &nodes);
+
+bool tensor_graph_backward(std::vector<execution_node_t *> &nodes);
+
+void assignGradMemory(tensor_pool_t *pool_grad_cpu, tensor_pool_t *pool_grad_gpu, std::vector<execution_node_t *> &nodes);
 
 #ifdef __cplusplus
 }
diff --git a/main.cpp b/main.cpp
index 4ede92f..28a19cb 100644
--- a/main.cpp
+++ b/main.cpp
@@ -2,75 +2,123 @@
 #include "soft-cuda/tensor/debug_api.h"
 #include <iostream>
 #include <vector>
+#include <cassert>
 
 using namespace std;
 
 int main() {
-    // Create pools
-    tensor_pool_t *pool = tensor_pool_create(1024 * 1024);
-    tensor_pool_t *pool_2 = tensor_pool_create(1024 * 1024);
-    tensor_pool_t *pool_gpu = tensor_pool_create(1024 * 1024, true);
+    tensor_pool_t *pool          = tensor_pool_create(1024 * 1024);
+    tensor_pool_t *pool_meta     = tensor_pool_create(1024 * 1024);
+    tensor_pool_t *pool_grad_cpu = tensor_pool_create(1024 * 1024);
+
+    tensor_pool_t *pool_gpu      = tensor_pool_create(1024 * 1024, true);
+    tensor_pool_t *pool_grad_gpu = tensor_pool_create(1024 * 1024, true);
+
     assert(pool != NULL);
     assert(pool_gpu != NULL);
 
     cout << "=========================================== \n";
-    cout << "=============TESTING DAG VERIFICATION============== \n";
+    cout << "========= XOR IMPLEMENTATION (4-NEURON) ==== \n";
+
+    float val_X[8]{0,0, 0,1, 1,0, 1,1};
+    float val_Y[4]{0, 1, 1, 0};
     
-    float val_a[900]{};
-    float val_b[900]{};
-    float val_c[900]{};
-    float val_d[900]{};
+    float val_W1[8]{}; 
+    float val_W2[4]{}; 
+    float val_b1[4]{}; 
+    float val_b2[1]{}; 
+
+    uint32_t dims_X[]  = {4, 2};
+    uint32_t dims_Y[]  = {4, 1};
+    uint32_t dims_W1[] = {2, 4}; 
+    uint32_t dims_b1[] = {1, 4};
+    uint32_t dims_W2[] = {4, 1};
+    uint32_t dims_b2[] = {1, 1};
+
+    tensor_t *X  = tensor_create(pool, tensor_dtype_t::FLOAT32_T, 2, dims_X,  val_X, false);
+    tensor_t *Y  = tensor_create(pool, tensor_dtype_t::FLOAT32_T, 2, dims_Y,  val_Y, false);
     
-    // TACTICAL FIX: Make 'e' the exact same size to avoid a GPU segfault 
-    // before we implement a dedicated GPU broadcasting kernel!
-    float val_e[900]{};
-
-    uint32_t dims_a[] = {30, 30};
-    uint32_t dims_b[] = {30, 30};
-    uint32_t dims_c[] = {30, 30};
-    uint32_t dims_d[] = {30, 30};
-    uint32_t dims_e[] = {30, 30}; // Matched dimensions!
-
-    tensor_t *a = tensor_create(pool, tensor_dtype_t::FLOAT32_T, 2, dims_a, val_a);
-    tensor_t *b = tensor_create(pool, tensor_dtype_t::FLOAT32_T, 2, dims_b, val_b);
-    tensor_t *c = tensor_create(pool, tensor_dtype_t::FLOAT32_T, 2, dims_c, val_c);
-    tensor_t *d = tensor_create(pool, tensor_dtype_t::FLOAT32_T, 2, dims_d, val_d);
-    tensor_t *e = tensor_create(pool, tensor_dtype_t::FLOAT32_T, 2, dims_e, val_e);
-    tensor_fill_random_normal(a, 10.1, 5.7);
-    tensor_fill_random_normal(b, 10.1, 5.7);
-    tensor_fill_random_normal(c, 10.1, 5.7);
-    tensor_fill_random_normal(d, 10.1, 5.7);
-    tensor_fill_random_normal(e, 10.1, 5.7);
-    tensor_t *f = tensor_mul(pool, a, b);
-    tensor_t *g = tensor_mul(pool, b, c);
-    tensor_t *h = tensor_mul(pool, e, f);
-    tensor_t *i = tensor_mul(pool, h, a);
-    tensor_t *j = tensor_mul(pool, i, g);
+    tensor_t *W1 = tensor_create(pool, tensor_dtype_t::FLOAT32_T, 2, dims_W1, val_W1, true);
+    tensor_t *W2 = tensor_create(pool, tensor_dtype_t::FLOAT32_T, 2, dims_W2, val_W2, true);
+    tensor_t *b1 = tensor_create(pool, tensor_dtype_t::FLOAT32_T, 2, dims_b1, val_b1, true);
+    tensor_t *b2 = tensor_create(pool, tensor_dtype_t::FLOAT32_T, 2, dims_b2, val_b2, true);
+
+    tensor_fill_random_normal(W1, 0.5, 0.2);
+    tensor_fill_random_normal(W2, 0.5, 0.2);
+    tensor_fill_random_normal(b1, 0.0, 0.1);
+    tensor_fill_random_normal(b2, 0.0, 0.1);
 
     cout << "=========================================== \n";
-    cout << "==============LAZY EVAL DONE=============== \n";
+    cout << "============== DEFINING OPS =============== \n";
+
+    tensor_t *H_pred_bef = tensor_mul_naive(pool, X, W1);
+    tensor_t *H_pred     = tensor_add(pool, H_pred_bef, b1);
+    tensor_t *H          = tensor_relu(pool, H_pred);
+    tensor_t *Y_pred_bef = tensor_mul_naive(pool, H, W2);
+    tensor_t *Y_pred     = tensor_add(pool, Y_pred_bef, b2);
     
-    std::vector<execution_node_t*> seq;
-    bool oki = verifyIfDAG(pool_2, j, seq);
+    tensor_t *L_sub      = tensor_sub(pool, Y_pred, Y);
+    tensor_t *L_squ      = tensor_square(pool, L_sub);
+    tensor_t *mse        = tensor_mean(pool, L_squ);
     
-    // Ensure assignDevice is returning device_type::GPU under the hood!
+    cout << "=========================================== \n";
+    cout << "============== LAZY EVAL DONE ============= \n";
+
+    cout << "=========================================== \n";
+    cout << "=========== PREPARING THE KITCHEN ========= \n";
+    std::vector<execution_node_t*> seq;
+    bool oki = verifyIfDAG(pool_meta, mse, seq);
     assignBackendGraph(pool_gpu, seq);
-    tensor_graph_forward_evaluate(pool, pool_gpu, seq);
-    
+    assignGradMemory(pool_grad_cpu, pool_grad_gpu, seq);
+
     if (oki) {
         cout << "=========================================== \n";
-        cout << "========= VRAM TRACE ============= \n";
+        cout << "============= STARTING TRAINING =========== \n";
         
-            execution_node_t *node = seq.back();
+        for (int i = 0; i <= 10000; i++) {
+            tensor_graph_forward_evaluate(pool, pool_gpu, seq);
+            gradInitializer(seq);
+            tensor_graph_backward(seq);
+            
+            if (i % 1000 == 0) {
+                std::cout << "EPOCH " << i << "\n";
+                execution_node_t *mse_node = seq.back();
+                printExecutionNode(mse_node);
+            }
             
-                // Pull it back to the CPU
-                execution_node_to_host(node);
-                printExecutionNode(node);
-                cout << "\n";
+            tensor_sgd(seq, 0.05); 
+        }
+
+        tensor_graph_forward_evaluate(pool, pool_gpu, seq);
+
+        float* inputs      = (float*)tensor_get_data(X);
+        float* targets     = (float*)tensor_get_data(Y);
+        float* predictions = (float*)tensor_get_data(Y_pred);
+
+        cout << "=========================================== \n";
+        cout << "============= XOR(4 NEURONS RESULT) ============= \n";
+        cout << "X1\tX2\t|\tTarget\t|\tPredicted\n";
+        cout << "---------------------------------------------------\n";
+
+        for (int i = 0; i < 4; i++) {
+            float x1     = inputs[i * 2 + 0];
+            float x2     = inputs[i * 2 + 1];
+            float y_true = targets[i];
+            float y_pred = predictions[i];
+
+            cout << x1 << "\t" << x2 << "\t|\t" << y_true << "\t|\t" << y_pred << "\n";
+        }
+        cout << "=========================================== \n";
+        
+    } else {
+        cout << "WARNING: DAG Verification failed. Aborting expedition.\n";
     }
 
     tensor_pool_destroy(pool);
-    tensor_pool_destroy(pool_2);
+    tensor_pool_destroy(pool_meta);
+    tensor_pool_destroy(pool_grad_cpu);
+    tensor_pool_destroy(pool_grad_gpu);
     tensor_pool_destroy(pool_gpu);
+
     return 0;
 }
diff --git a/make.sh b/make.sh
index 23c57b2..27ac66a 100755
--- a/make.sh
+++ b/make.sh
@@ -1,2 +1,2 @@
-cmake -B build -DCMAKE_BUILD_TYPE=Debug
-cmake --build build
+time cmake -B build -DCMAKE_BUILD_TYPE=Debug
+time cmake --build build -j16
diff --git a/notes/Backward/MemoryModel.md b/notes/Backward/MemoryModel.md
new file mode 100644
index 0000000..86f3724
--- /dev/null
+++ b/notes/Backward/MemoryModel.md
@@ -0,0 +1,53 @@
+## Memory model for backward pass(Unstable)
+
+The problem we have is that when we do hybrid forward pass, it causes some data to be on GPU and some to be on CPU, this I thought
+would be a problem but that is not the case.
+
+### Observation
+> Since data flow is only one directional and due to our contagious logic
+
+```
+            assignPlaceOnDeviceMemory(pool, (int32_t)(node->pos), nodes);
+            if (node->t->a == NULL)
+                goto trp;
+            if (node->t->a->device == device_type::CPU) {
+                int32_t a_idx = getTheExecutionNodeIndex(node, 0);
+                assert(a_idx != -1);
+                nodes[((size_t)a_idx)]->to_device_needed = true;
+                assignPlaceOnDeviceMemory(pool, a_idx, nodes);
+            }
+        trp:
+            if (node->t->b == NULL)
+                continue;
+            if (node->t->b->device == device_type::CPU) {
+                int32_t b_idx = getTheExecutionNodeIndex(node, 1);
+                assert(b_idx != -1);
+                nodes[((size_t)b_idx)]->to_device_needed = true;
+                assignPlaceOnDeviceMemory(pool, b_idx, nodes);
+            }
+
+
+```
+We can conclude that all the data we need for grad computation is available on each device. and since we copy the parents we can 
+confidently say that we don't need to retransfer tensors inbetween eliminating the thrashing.
+
+* So here comes the part where what we will do in
+`void assignBackendGraph(tensor_pool_t *pool,std::vector<execution_node_t *> &nodes);` 
+is we will allocate the grad space for tensors on GPU as `pool_gpu_grad` and on CPU as `pool_cpu_grad` then we assign it to the home tensor(refering to the current tensor)
+then during backward pass data is assigned to each pool and is written as it is.
+since we don't want transfer overhead during computation path we will construct another data transfering fucnction which will get grad data from gpu to host.
+and then we do optimization and then free resources.
+
+Also the device_ptr_gpu is stored in execution_node_t;
+
+===================================================================================================
+
+actually i think we can't do grad placement in the AssignBackendGraph function which we are implementing the contagious logic we will need a 4th pass i think where so we can do this maneuvour
+
+P1->GPU, P2->GPU, C1->GPU
+
+but now if P1 and P2 was bought to GPU and it's copy exist on cpu we can jump to it's DRAM storage and i think since it's address we will know due to well how we can detect it you may ask it would be if it have address on execution node and also on tensor instance right then we switch to CPU allocation right.
+
+Now for this to work we will need to walk reverse in graph so one more reason for 4th seperate pass plus this also helps that we will not be touching much of things.
+
+Now if we are already doing 4th pass it would be better to instead have a seperate explict function for this so func params also don't get bloated and confusing.
diff --git a/scripts/coldstart/coldstart.prev b/scripts/coldstart/coldstart.prev
index e69de29..e48b1cd 100644
--- a/scripts/coldstart/coldstart.prev
+++ b/scripts/coldstart/coldstart.prev
@@ -0,0 +1,36 @@
+Based on the git log and the recent activity, here is the reconstruction of your mental context.
+
+### 1. What I Was Doing
+I have been building out the core components of a tensor/autograd library. I just successfully finished the **Mean Squared Error (MSE)** loss function and verified it with smoke tests. Immediately after that, I shifted focus to **gradient memory management**, specifically writing a function to handle how gradient memory is assigned/allocated to tensors.
+
+### 2. Active TODOs (Grouped)
+*   **Autograd Engine/Memory:**
+    *   Integrate `assign_grad_memory` into the tensor initialization or the backward pass flow.
+    *   Ensure gradients are properly zeroed out before the next backward pass (standard `zero_grad` functionality).
+*   **Loss Functions:**
+    *   (Optional) Implement Cross-Entropy or other standard loss functions now that the MSE template is proven.
+*   **Tooling/DevOps:**
+    *   The repetitive "Running summarizer" commits suggest I am fine-tuning an automated workflow or documentation script.
+
+### 3. Current State of the Code
+*   **Implemented:** 
+    *   MSE Loss function (logic is solid and passed basic verification).
+    *   `make.sh` utility now includes a time field for performance tracking.
+    *   A function to assign gradient memory (`assign_grad_memory`).
+*   **Partially Done:** 
+    *   The plumbing for backpropagation. While the loss exists, the full "Backward" chain for a model likely needs the newly created gradient memory function to be fully integrated.
+*   **Missing:** 
+    *   A comprehensive test for a full training step (Forward -> MSE -> Backward -> Optimizer Step).
+
+### 4. My Likely Next Steps
+1.  **Verify Memory Assignment:** Write a small test script to ensure `assign_grad_memory` actually allocates the correct shapes and doesn't leak memory.
+2.  **Connect MSE to Backprop:** Implement the `.backward()` method for the MSE node, ensuring it utilizes the new gradient memory allocation.
+3.  **End-to-End Smoke Test:** Run a simple linear regression or a single-neuron test to see if the gradient flows from the MSE loss back to the input weights correctly.
+
+### 5. Important Context I Should Not Forget
+*   **Performance Tracking:** I added a `time` field to `make.sh`, which implies I'm starting to care about how long the build or the tests take. I should keep an eye on this as memory management becomes more complex.
+*   **Manual Memory Management:** Since I'm writing `assign_grad_memory` manually, I need to be careful about whether gradients are being allocated lazily or eagerly at tensor creation.
+
+### 6. Risks / Tricky Areas
+*   **Gradient Accumulation:** If `assign_grad_memory` is called multiple times or incorrectly, I might overwrite gradients instead of accumulating them, or double-allocate memory.
+*   **Smoke Test Scope:** Passing "smoke tests" for MSE proves the math is right, but it doesn't prove the engine can handle a complex graph yet. Be wary of moving to complex architectures too fast.
diff --git a/scripts/coldstart/coldstart.sh b/scripts/coldstart/coldstart.sh
index 8b3c45f..dc394d7 100755
--- a/scripts/coldstart/coldstart.sh
+++ b/scripts/coldstart/coldstart.sh
@@ -1,6 +1,6 @@
 cp coldstart.tmp coldstart.prev 2>/dev/null
 
-git add ..
+git add ../..
 
 git commit -m "chore/ Running summarizer"
 date=$(date +%Y-%m-%d_%H-%M-%S)
@@ -11,7 +11,7 @@ filename="agent_summary_${date}.log"
   echo "=== GIT DIFF ==="
   git diff HEAD
   echo "=== TODO ==="
-  grep -r TODO ../src/ 2>/dev/null
+  grep -r TODO ../../src 2>/dev/null
   echo "=== PREVIOUS SESSION SUMMARY ==="
   [ -f coldstart.prev ] && cat coldstart.prev
 } | python coldstart_agent.py > logs/${filename}
@@ -19,6 +19,6 @@ filename="agent_summary_${date}.log"
 cat logs/${filename} > coldstart.tmp
 
 
-git add ..
+git add .
 
 git commit -m "chore/ Summarizer completion logging"
diff --git a/src/backend_cpu/backprop/backprop_b.cpp b/src/backend_cpu/backprop/backprop_b.cpp
new file mode 100644
index 0000000..c8f2712
--- /dev/null
+++ b/src/backend_cpu/backprop/backprop_b.cpp
@@ -0,0 +1,357 @@
+#include "internal_header.h"
+
+
+// We are gonna write a backprop fucntion for matmul cause that is easy thing to do
+// Since our backprop is an eagar evaluation we will make tensor objects inside the function
+// transpose it execute it as well and then work with it
+
+// Now I am slightly confused but I think it would have gone something like this
+// Since we are going to get execution node we could have a decision made here if it's data will be
+// written on device_ptr of node_.grad->data
+// We are gonna do this cause making decsion on higher level will require us to make two kernels for similar work.
+void gradInitializer(std::vector<execution_node_t *> &nodes) {
+    for (auto &node : nodes) {
+        if (!node->t->grad_compute) continue;
+        memset(node->t->grad->data, 0,node->t->grad->nvalues * sizeof(float));
+        // TODO: We are not zeroing the GPU will have to do that
+    }
+    if (!nodes.empty()) {
+        execution_node_t *root = nodes.back();
+        if (root->t != nullptr && root->t->grad != nullptr && root->t->grad->data != nullptr) {
+            ((float*)root->t->grad->data)[0] = 1.0f;
+        }
+    }
+}
+// Since this will be eagar evalution we are gonna execute it here now and we are gonna return a
+// boolean.
+// @return: boolean
+// @params: execution_node_t
+// Since all the memory is already executed we just need the execution node
+
+bool backprop__(std::vector<execution_node_t *> &nodes) {
+    for (int i = (int)nodes.size() - 1; i >= 0; i--) {
+        execution_node_t *node = nodes[(size_t)i];
+        // Cause it's not needed
+        if (!node->t->grad_compute)
+            continue;
+        // Cause it's super child
+        if (node->t->op == tensor_op_t::NONE)
+            continue;
+        bool success{false};
+        if (node->device_ptr_grad != NULL) {
+            success = backprop_gpu(node);
+        } else {
+            success = backprop_cpu(node);
+        }
+        if (!success) {
+            debug("backprop__: FUBAR at node pos=%d\n", (int)node->pos);
+            return false;
+        }
+    }
+    return true;
+}
+
+// ================================================================================ 
+// CPU BACKWARD DISPATCHER
+// ================================================================================ 
+
+bool backprop_cpu(execution_node_t *node) {
+    assert(node != NULL);
+    assert(node->t != NULL);
+    bool success = false;
+
+    switch (node->t->op) {
+    case tensor_op_t::NONE:
+        // Leaf node
+        success = true;
+        break;
+    case tensor_op_t::CAST:
+        // TODO: Implement here
+        success = false;
+        break;
+    // case tensor_op_t::MUL_MATRIX:
+    //     assert(node->t->a != NULL);
+    //     assert(node->t->b != NULL);
+    //     assert(node->t->a->dtype == node->t->b->dtype);
+    //     success = tensor_mul_grad_op_matrix(node);
+    //     break;
+    case tensor_op_t::MUL_SCALAR:
+        assert(node->t->a != NULL);
+        assert(node->t->b != NULL);
+        assert(node->t->a->dtype == node->t->b->dtype);
+        success = tensor_mul_grad_op_scalar(node);
+        break;
+    case tensor_op_t::TRANSPOSE:
+        assert(node->t->a != NULL);
+        success = tensor_tranpose_grad_op_matrix(node);
+        break;
+    case tensor_op_t::NAIVE_MATRIX_MUL:
+        assert(node->t->a != NULL);
+        assert(node->t->b != NULL);
+        assert(node->t->a->dtype == node->t->b->dtype);
+        success = tensor_mul_grad_op_matrix_naive(node);
+        break;
+    case tensor_op_t::ADD:
+        assert(node->t->a != NULL);
+        assert(node->t->b != NULL);
+        assert(node->t->a->dtype == node->t->b->dtype);
+        success = tensor_grad_op_add(node);
+        break;
+    case tensor_op_t::BROADCAST_ADD:
+        assert(node->t->a != NULL);
+        assert(node->t->b != NULL);
+        assert(node->t->a->dtype == node->t->b->dtype);
+        success = tensor_grad_op_broadcasting_add(node);
+        break;
+    case tensor_op_t::RELU:
+        assert(node->t->a != NULL);
+        success = tensor_grad_op_relu(node);
+        break;
+    case tensor_op_t::SUB:
+        assert(node->t->a != NULL);
+        assert(node->t->b != NULL);
+        assert(node->t->a->dtype == node->t->b->dtype);
+        success = tensor_grad_op_sub(node);
+        break;
+    case tensor_op_t::MEAN:
+        assert(node->t->a != NULL);
+        success = tensor_grad_op_mean(node);
+        break;
+    case tensor_op_t::SQUARE:
+        assert(node->t->a != NULL);
+        success = tensor_grad_op_square(node);
+        break;
+    default:
+        assert(false);
+    }
+    if (success) {
+        debug("backprop_cpu: success for op=%d\n", (int)node->t->op);
+    } else {
+        debug("backprop_cpu: FUBAR for op=%d\n", (int)node->t->op);
+    }
+    return success;
+}
+
+// =============================================================================
+// GPU BACKWARD DISPATCHER (STUB)
+// =============================================================================
+bool backprop_gpu([[maybe_unused]] execution_node_t *node) {
+    assert(node != NULL);
+    assert(node->t != NULL);
+    // TODO: Implement GPU backward kernels
+    debug("backprop_gpu: not yet implemented, falling back\n");
+    return false;
+}
+// TODO: ONE PROBLEM IS WE NEED MULTIPLE CASE FOR 
+bool tensor_grad_op_add(execution_node_t *node) {
+    tensor_t *out_grad = node->t->grad;
+    tensor_t *a = node->t->a;
+    tensor_t *b = node->t->b;
+
+    assert(out_grad != NULL && a != NULL && b != NULL);
+    float *g_out = (float *)out_grad->data;
+    float *g_a = a->grad != NULL ? (float *)a->grad->data : NULL;
+    float *g_b = b->grad != NULL ? (float *)b->grad->data : NULL;
+    uint32_t n = node->t->nvalues;
+
+    for (uint32_t i = 0; i < n; i++) {
+        if (g_a) g_a[i] += g_out[i];
+        if (g_b) g_b[i] += g_out[i];
+    }
+    return true;
+}
+
+bool tensor_grad_op_broadcasting_add(execution_node_t *node) {
+    tensor_t *out_grad = node->t->grad;
+    tensor_t *a = node->t->a;
+    tensor_t *b = node->t->b;
+
+    assert(out_grad != NULL && a != NULL && b != NULL);
+    float *g_out = (float *)out_grad->data;
+    float *g_a = a->grad != NULL ? (float *)a->grad->data : NULL;
+    float *g_b = b->grad != NULL ? (float *)b->grad->data : NULL;
+
+    uint32_t rows = node->t->dims[0];
+    uint32_t cols = node->t->dims[1];
+    // We are assuming that b would be the broadcasted array
+    uint32_t val = node->t->nvalues;
+    if (g_a) {
+        for (uint32_t i = 0; i < val; i++) {
+            g_a[i] += g_out[i];
+        }
+    }
+
+    if (g_b) {
+        // For broadcasting we accumulate btw
+        for (uint32_t i = 0; i < rows; i++) {
+          for (uint32_t j = 0; j < cols; j++) {
+            uint32_t b_idx = i * b->broadcast_stride[0]  + j * b->broadcast_stride[1];
+            g_b[b_idx] += g_out[i*cols + j];
+          }
+        }
+    }
+    return true;
+}
+
+bool tensor_grad_op_sub(execution_node_t *node) {
+    tensor_t *out_grad = node->t->grad;
+    tensor_t *a        = node->t->a;
+    tensor_t *b        = node->t->b;
+    assert(out_grad != NULL && a != NULL && b != NULL);
+
+    float *g_out = (float *)out_grad->data;
+    float *g_a   = a->grad != NULL ? (float *)a->grad->data : NULL;
+    float *g_b   = b->grad != NULL ? (float *)b->grad->data : NULL;
+    uint32_t n   = node->t->nvalues;
+
+    for (uint32_t i = 0; i < n; i++) {
+        if (g_a) g_a[i] += g_out[i];
+        if (g_b) g_b[i] -= g_out[i];
+    }
+    return true;
+}
+
+
+
+bool tensor_grad_op_relu(execution_node_t *node) {
+    tensor_t *out_grad = node->t->grad;
+    tensor_t *a        = node->t->a;
+    assert(out_grad != NULL && a != NULL);
+
+    if (a->grad != NULL) {
+        float *g_out  = (float *)out_grad->data;
+        float *g_a    = (float *)a->grad->data;
+        float *a_data = (float *)a->data;   
+        uint32_t n    = node->t->nvalues;
+
+        for (uint32_t i = 0; i < n; i++) {
+            if (a_data[i] > 0.0f) {
+                g_a[i] += g_out[i];
+            }
+        }
+    }
+    return true;
+}
+
+bool tensor_grad_op_mean(execution_node_t *node) {
+    tensor_t *out_grad = node->t->grad;
+    tensor_t *a        = node->t->a;
+    assert(out_grad != NULL && a != NULL);
+
+    if (a->grad != NULL) {
+        // out->grad is a scalar — one value
+        float upstream    = ((float *)out_grad->data)[0];
+        float *g_a        = (float *)a->grad->data;
+        uint32_t n        = a->nvalues;
+        float scale       = upstream / (float)n;
+
+        for (uint32_t i = 0; i < n; i++) {
+            g_a[i] += scale;
+        }
+    }
+    return true;
+}
+
+bool tensor_mul_grad_op_scalar(execution_node_t *node) {
+    tensor_t *out_grad = node->t->grad;
+    tensor_t *a        = node->t->a;
+    tensor_t *b        = node->t->b;   
+    assert(out_grad != NULL && a != NULL && b != NULL);
+
+    if (a->grad != NULL) {
+        float *g_out  = (float *)out_grad->data;
+        float *g_a    = (float *)a->grad->data;
+        float  s      = tensor_float32_value(b);
+        uint32_t n    = a->nvalues;
+
+        for (uint32_t i = 0; i < n; i++) {
+            g_a[i] += g_out[i] * s;
+        }
+    }
+    return true;
+}
+
+
+bool tensor_tranpose_grad_op_matrix(execution_node_t *node) {
+    tensor_t *out_grad = node->t->grad;
+    tensor_t *a        = node->t->a;
+    assert(out_grad != NULL && a != NULL);
+
+    if (a->grad != NULL) {
+        uint32_t rows    = a->dims[0];   
+        uint32_t cols    = a->dims[1];  
+        float *g_out     = (float *)out_grad->data;
+        float *g_a       = (float *)a->grad->data;
+
+        for (uint32_t i = 0; i < rows; i++) {
+            for (uint32_t j = 0; j < cols; j++) {
+                g_a[i * cols + j] += g_out[j * rows + i];
+            }
+        }
+    }
+    return true;
+}
+
+bool tensor_mul_grad_op_matrix_naive(execution_node_t *node) {
+    tensor_t *out_grad = node->t->grad;
+    tensor_t *a        = node->t->a;
+    tensor_t *b        = node->t->b;
+    assert(out_grad != NULL && a != NULL && b != NULL);
+
+    uint32_t M = a->dims[0];
+    uint32_t K = a->dims[1];
+    uint32_t N = b->dims[1];
+
+    float *g_out = (float *)out_grad->data;    // (M, N)
+    float *a_data = (float *)a->data;          // (M, K)
+    float *b_data = (float *)b->data;          // (K, N)
+    float *g_a    = a->grad != NULL ? (float *)a->grad->data : NULL;    // (M, K)
+    float *g_b    = b->grad != NULL ? (float *)b->grad->data : NULL;    // (K, N)
+
+    if (g_a) {
+        // dL/dA[i,k] += sum_j  dL/dOut[i,j] * B[k,j]
+        for (uint32_t i = 0; i < M; i++) {
+            for (uint32_t k = 0; k < K; k++) {
+                float sum = 0.0f;
+                for (uint32_t j = 0; j < N; j++) {
+                    sum += g_out[i * N + j] * b_data[k * N + j];
+                }
+                g_a[i * K + k] += sum;
+            }
+        }
+    }
+
+    if (g_b) {
+        // dL/dB[k,j] += sum_i  A[i,k] * dL/dOut[i,j]
+        for (uint32_t k = 0; k < K; k++) {
+            for (uint32_t j = 0; j < N; j++) {
+                float sum = 0.0f;
+                for (uint32_t i = 0; i < M; i++) {
+                    sum += a_data[i * K + k] * g_out[i * N + j];
+                }
+                g_b[k * N + j] += sum;
+            }
+        }
+    }
+    return true;
+}
+
+bool tensor_grad_op_square(execution_node_t *node) {
+    tensor_t *out_grad = node->t->grad;
+    tensor_t *a        = node->t->a;
+    assert(out_grad != NULL && a != NULL);
+
+    if (a->grad != NULL) {
+        float *g_out  = (float *)out_grad->data;
+        float *g_a    = (float *)a->grad->data;
+        float *a_data = (float *)a->data;   
+        uint32_t n    = node->t->nvalues;
+
+        for (uint32_t i = 0; i < n; i++) {
+            g_a[i] += g_out[i] * 2.0f * a_data[i];
+        }
+    }
+    return true;
+}
+
+
diff --git a/src/backend_cpu/math/add_bias.cpp b/src/backend_cpu/backprop/matmul_b.cpp
similarity index 100%
rename from src/backend_cpu/math/add_bias.cpp
rename to src/backend_cpu/backprop/matmul_b.cpp
diff --git a/src/backend_cpu/include/backprop/backprop_b.h b/src/backend_cpu/include/backprop/backprop_b.h
new file mode 100644
index 0000000..6dd53fb
--- /dev/null
+++ b/src/backend_cpu/include/backprop/backprop_b.h
@@ -0,0 +1,27 @@
+#pragma once
+
+void gradInitializer(std::vector<execution_node_t *> &nodes);
+
+bool backprop__(std::vector<execution_node_t *> &nodes);
+
+bool backprop_cpu(execution_node_t *node);
+
+bool backprop_gpu([[maybe_unused]] execution_node_t *node);
+
+bool tensor_grad_op_add(execution_node_t *node);
+
+bool tensor_grad_op_broadcasting_add(execution_node_t *node);
+
+bool tensor_grad_op_sub(execution_node_t *node);
+
+bool tensor_grad_op_relu(execution_node_t *node);
+
+bool tensor_grad_op_mean(execution_node_t *node);
+
+bool tensor_mul_grad_op_scalar(execution_node_t *node);
+
+bool tensor_tranpose_grad_op_matrix(execution_node_t *node);
+
+bool tensor_mul_grad_op_matrix_naive(execution_node_t *node);
+
+bool tensor_grad_op_square(execution_node_t *node);
diff --git a/src/backend_cpu/include/mean.h b/src/backend_cpu/include/mean.h
new file mode 100644
index 0000000..e82f37f
--- /dev/null
+++ b/src/backend_cpu/include/mean.h
@@ -0,0 +1,2 @@
+#pragma once
+bool tensor_op_mean(tensor_pool_t *pool, tensor_t *t);
diff --git a/src/backend_cpu/include/square.h b/src/backend_cpu/include/square.h
new file mode 100644
index 0000000..ceb9c33
--- /dev/null
+++ b/src/backend_cpu/include/square.h
@@ -0,0 +1,7 @@
+#pragma once
+
+#include "internal_header.h"
+
+tensor_t *tensor_square(tensor_pool_t *pool, tensor_t *x);
+
+bool tensor_op_square(tensor_pool_t *pool, tensor_t *t);
diff --git a/src/backend_cpu/include/sub.h b/src/backend_cpu/include/sub.h
new file mode 100644
index 0000000..a82defb
--- /dev/null
+++ b/src/backend_cpu/include/sub.h
@@ -0,0 +1,6 @@
+#pragma once
+
+tensor_t *tensor_sub(tensor_pool_t *pool, tensor_t *a, tensor_t *b);
+
+
+bool tensor_op_sub(tensor_pool_t *pool, tensor_t *t);
diff --git a/src/backend_cpu/math/add.cpp b/src/backend_cpu/math/add.cpp
index 54b5d77..a110607 100644
--- a/src/backend_cpu/math/add.cpp
+++ b/src/backend_cpu/math/add.cpp
@@ -1,78 +1,82 @@
-#include "internal_header.h"
-
-static tensor_t *tensor_broadcast_add(tensor_pool_t *pool, tensor_t *a, tensor_t *b) {
-    assert(pool != NULL);
-    assert(a != NULL);
-    assert(b != NULL);
-
-    if (b->dims[0] == 1) {
-        b->broadcast_stride[0] = 0;
-        b->broadcast_stride[1] = b->stride[1];
-    } else {
-        b->broadcast_stride[1] = 0;
-        b->broadcast_stride[0] = b->stride[0];
-    }
-
-    tensor_t *t = tensor_dtype_create(pool, a->dtype, a->ndims, a->dims, NULL);
-    if (t == NULL) {
-        return NULL;
-    }
-
-    t->op = tensor_op_t::BROADCAST_ADD;
-    t->a = a;
-    t->b = b;
-    return t;
-}
-
-// We are baking lots of assumptions in this one haha
-tensor_t *tensor_add(tensor_pool_t *pool, tensor_t *x, tensor_t *y) {
-    assert(pool != NULL);
-    assert(x != NULL);
-    assert(y != NULL);
-
-    if (x->dims[0] == 1 || x->dims[1] == 1) {
-        return tensor_broadcast_add(pool, y, x);
-    }
-    if (y->dims[0] == 1 || y->dims[1] == 1) {
-        return tensor_broadcast_add(pool, x, y);
-    }
-
-    tensor_t *t = tensor_dtype_create(pool, x->dtype, x->ndims, x->dims, NULL);
-    if (t == NULL) {
-        return NULL;
-    }
-    t->op = tensor_op_t::ADD;
-    t->a = x;
-    t->b = y;
-    return t;
-}
-
-bool tensor_op_add(tensor_pool_t *pool, tensor_t *t) {
-    assert(pool != NULL);
-    assert(t != NULL);
-
-    for (uint32_t i = 0; i < t->dims[0] * t->dims[1]; i++) {
-        ((float *)t->data)[i] = ((float *)t->a->data)[i] + ((float *)t->b->data)[i];
-    }
-    return true;
-}
-
-bool tensor_op_broadcasting_add(tensor_pool_t *pool, tensor_t *t) {
-    assert(pool != NULL);
-    assert(t != NULL);
-    assert(t->a != NULL);
-    assert(t->b != NULL);
-
-    for (uint32_t i = 0; i < t->dims[0] * t->dims[1]; i++) {
-        ((float *)t->data)[i] = ((float *)t->a->data)[i] + get_sum_for_col_op(t->b, i, t);
-    }
-    return true;
-}
-
-float get_sum_for_col_op(tensor_t *b, uint32_t i, tensor_t *t) {
-    uint32_t row = i / (t->dims[1]);
-    uint32_t col = i % (t->dims[1]);
-    float val =
-        ((float *)b->data)[row * ((b->broadcast_stride)[0]) + col * ((b->broadcast_stride)[1])];
-    return val;
-}
+#include "internal_header.h"
+
+static tensor_t *tensor_broadcast_add(tensor_pool_t *pool, tensor_t *a, tensor_t *b) {
+    assert(pool != NULL);
+    assert(a != NULL);
+    assert(b != NULL);
+
+    if (b->dims[0] == 1) {
+        b->broadcast_stride[0] = 0;
+        b->broadcast_stride[1] = b->stride[1];
+    } else {
+        b->broadcast_stride[1] = 0;
+        b->broadcast_stride[0] = b->stride[0];
+    }
+
+    tensor_t *t = tensor_dtype_create(pool, a->dtype, a->ndims, a->dims, NULL);
+    if (t == NULL) {
+        return NULL;
+    }
+
+    t->op = tensor_op_t::BROADCAST_ADD;
+    t->a = a;
+    t->b = b;
+    return t;
+}
+
+// We are baking lots of assumptions in this one haha
+tensor_t *tensor_add(tensor_pool_t *pool, tensor_t *x, tensor_t *y) {
+    assert(pool != NULL);
+    assert(x != NULL);
+    assert(y != NULL);
+
+    if (x->nvalues == y->nvalues) {
+        tensor_t *t = tensor_dtype_create(pool, x->dtype, x->ndims, x->dims, NULL);
+        if (t == NULL) return NULL;
+        t->op = tensor_op_t::ADD;
+        t->a = x;
+        t->b = y;
+        return t;
+    }
+
+    if (y->nvalues < x->nvalues && (y->dims[0] == 1 || y->dims[1] == 1)) {
+        return tensor_broadcast_add(pool, x, y);
+    }
+    
+    if (x->nvalues < y->nvalues && (x->dims[0] == 1 || x->dims[1] == 1)) {
+        return tensor_broadcast_add(pool, y, x);
+    }
+
+    return NULL; // Unsupported broadcast shape
+}
+
+bool tensor_op_add(tensor_pool_t *pool, tensor_t *t) {
+    assert(pool != NULL);
+    assert(t != NULL);
+
+    for (uint32_t i = 0; i < t->dims[0] * t->dims[1]; i++) {
+        ((float *)t->data)[i] = ((float *)t->a->data)[i] + ((float *)t->b->data)[i];
+    }
+    return true;
+}
+
+bool tensor_op_broadcasting_add(tensor_pool_t *pool, tensor_t *t) {
+    assert(pool != NULL);
+    assert(t != NULL);
+    assert(t->a != NULL);
+    assert(t->b != NULL);
+
+    for (uint32_t i = 0; i < t->dims[0] * t->dims[1]; i++) {
+        ((float *)t->data)[i] = ((float *)t->a->data)[i] + get_sum_for_col_op(t->b, i, t);
+    }
+    return true;
+}
+
+float get_sum_for_col_op(tensor_t *b, uint32_t i, tensor_t *t) {
+    uint32_t row = i / (t->dims[1]);
+    uint32_t col = i % (t->dims[1]);
+    float val =
+        ((float *)b->data)[row * ((b->broadcast_stride)[0]) + col * ((b->broadcast_stride)[1])];
+    return val;
+}
+
diff --git a/src/backend_cpu/math/mean.cpp b/src/backend_cpu/math/mean.cpp
new file mode 100644
index 0000000..7fddeb6
--- /dev/null
+++ b/src/backend_cpu/math/mean.cpp
@@ -0,0 +1,27 @@
+#include "internal_header.h"
+
+tensor_t *tensor_mean(tensor_pool_t *pool, tensor_t *a) {
+    assert(pool != NULL);
+    assert(a != NULL);
+
+    tensor_t *t = tensor_dtype_create(pool, a->dtype, 0, 0, NULL);
+    if (t == NULL) {
+        return NULL;
+    }
+
+    t->op = tensor_op_t::MEAN;
+    t->a = a;
+    return t;
+}
+
+bool tensor_op_mean(tensor_pool_t *pool, tensor_t *t) {
+    assert(pool != NULL);
+    assert(t != NULL);
+    float sum{};
+    for (uint32_t i = 0; i < t->a->nvalues; i++) {
+        sum += ((float *)t->a->data)[i];
+    }
+    ((float *)t->data)[0] = sum/(float)(t->a->nvalues);
+    ((float *)t->grad->data)[0] = 1.0f;
+    return true;
+}
diff --git a/src/backend_cpu/math/mse.cpp b/src/backend_cpu/math/mse.cpp
new file mode 100644
index 0000000..82bba6a
--- /dev/null
+++ b/src/backend_cpu/math/mse.cpp
@@ -0,0 +1,2 @@
+#include "internal_header.h"
+
diff --git a/src/backend_cpu/math/square.cpp b/src/backend_cpu/math/square.cpp
new file mode 100644
index 0000000..a828114
--- /dev/null
+++ b/src/backend_cpu/math/square.cpp
@@ -0,0 +1,32 @@
+#include "internal_header.h"
+
+tensor_t *tensor_square(tensor_pool_t *pool, tensor_t *x) {
+    assert(pool != NULL);
+    assert(x != NULL);
+
+    tensor_t *t = tensor_dtype_create(pool, x->dtype, x->ndims, x->dims, NULL);
+    if (t == NULL) {
+        return NULL;
+    }
+
+    t->op = tensor_op_t::SQUARE;
+    t->a = x;
+    
+    return t;
+}
+
+bool tensor_op_square(tensor_pool_t *pool, tensor_t *t) {
+    assert(pool != NULL);
+    assert(t != NULL);
+    assert(t->a != NULL);
+
+    float *in = (float *)t->a->data;
+    float *out = (float *)t->data;
+    uint32_t n = t->nvalues;
+
+    for (uint32_t i = 0; i < n; i++) {
+        out[i] = in[i] * in[i];
+    }
+    return true;
+}
+
diff --git a/src/backend_cpu/math/sub.cpp b/src/backend_cpu/math/sub.cpp
new file mode 100644
index 0000000..3bd133c
--- /dev/null
+++ b/src/backend_cpu/math/sub.cpp
@@ -0,0 +1,27 @@
+#include "internal_header.h"
+
+tensor_t *tensor_sub(tensor_pool_t *pool, tensor_t *a, tensor_t *b) {
+    assert(pool != NULL);
+    assert(a != NULL);
+    assert(b != NULL);
+
+    tensor_t *t = tensor_dtype_create(pool, a->dtype, a->ndims, a->dims, NULL);
+    if (t == NULL) {
+        return NULL;
+    }
+
+    t->op = tensor_op_t::SUB;
+    t->a = a;
+    t->b = b;
+    return t;
+}
+
+bool tensor_op_sub(tensor_pool_t *pool, tensor_t *t) {
+    assert(pool != NULL);
+    assert(t != NULL);
+
+    for (uint32_t i = 0; i < t->dims[0] * t->dims[1]; i++) {
+        ((float *)t->data)[i] = ((float *)t->a->data)[i] - ((float *)t->b->data)[i];
+    }
+    return true;
+}
diff --git a/src/core/graph/assignBackend.cu b/src/core/graph/assignBackend.cu
index 007c09c..4d285d1 100644
--- a/src/core/graph/assignBackend.cu
+++ b/src/core/graph/assignBackend.cu
@@ -1,5 +1,5 @@
 #include "internal_header.h"
-#include "vector"
+// #include "vector"
 #include <iostream>
 #include <cuda_runtime.h>
 
@@ -23,6 +23,7 @@ void assignBackend(execution_node_t *e, json &data) {
  */
 device_type assignDevice([[maybe_unused]] uint8_t ndims, [[maybe_unused]] uint32_t *dims,
                          [[maybe_unused]] tensor_op_t op, uint32_t nvalues, json &data) {
+    return device_type::CPU;
     if(op == tensor_op_t::NONE) {
         return device_type::CPU;
     }
@@ -185,6 +186,7 @@ void printExecutionNode(execution_node_t *et) {
     tensor_print_data(et->t);
     std::cout << et->to_device_needed << "\n";
     std::cout << et->device_ptr << "\n";
+    std::cout << et->device_ptr_grad << "\n";
     std::cout << (et->backend_fn != NULL) << "\n";
 }
 
@@ -208,3 +210,29 @@ bool execution_node_to_host(execution_node_t *node) {
     }
     return true;
 }
+
+// What we can do is keep it default and then have user declare leaf to NULL. As otherwise it would
+// be difficult to implement the propogation logic
+// NOTE: Since we are making the memory here and device_ptr_grad is here to what we will do is take
+// the whole execution node and then pass it to the eagar OPS and then we will see if it does have
+// `device_ptr_grad` if so then it's fine right other wise we will directly write to
+// node->t->grad->data right. Then after all the ops we will make a new function which will scan the
+// whole graph and then if it sees the device_ptr_grad is not null it will initiate data transfer.
+void assignGradMemory(tensor_pool_t *pool_grad_cpu, tensor_pool_t *pool_grad_gpu, std::vector<execution_node_t *> &nodes) {
+    for(auto &node : nodes) {
+        // We are not assigning grad instance during tensor create.
+        // Maybe we can do that but I feel it might cause recursion problem even if with assigning base case of leaf nodes
+        // there is the problem that during creation time it could cause function call overhead and whatnot.
+        // So we are going to make a tensor instance here and assign it here i guess. 
+        
+        if(node->t->grad_compute == false) continue;
+
+        node->t->grad =  tensor_dtype_create(pool_grad_cpu, node->t->dtype, node->t->ndims, node->t->dims, NULL);
+
+        if(node->backend_fn == tensor_evaluate_GPU) {
+            size_t size = node->t->nvalues * sizeof(float) ;
+            uint32_t id;
+            node->device_ptr_grad = tensor_pool_alloc(pool_grad_gpu, size, &id);
+        }
+    }
+}
diff --git a/src/core/graph/train.cpp b/src/core/graph/train.cpp
new file mode 100644
index 0000000..f0c01a9
--- /dev/null
+++ b/src/core/graph/train.cpp
@@ -0,0 +1,26 @@
+#include "internal_header.h"
+#include <vector>
+
+bool tensor_graph_backward(std::vector<execution_node_t *> &nodes) {
+    return backprop__(nodes);
+}
+
+
+void tensor_sgd(std::vector<execution_node_t *> &nodes, float learning_rate) {
+    for (auto &node : nodes) {
+        tensor_t *t = node->t;
+
+        if (t->op != tensor_op_t::NONE) continue;
+        if (!t->grad_compute)           continue;
+        if (t->grad == NULL)            continue;
+
+        float *w    = (float *)t->data;
+        float *grad = (float *)t->grad->data;
+        uint32_t n  = t->nvalues;
+
+        for (uint32_t i = 0; i < n; i++) {
+            w[i] -= learning_rate * grad[i];
+            // grad[i] = 0.0f;   
+        }
+    }
+}
diff --git a/src/core/include/graph/DAGbuild.h b/src/core/include/graph/DAGbuild.h
index 707d324..180a1c6 100644
--- a/src/core/include/graph/DAGbuild.h
+++ b/src/core/include/graph/DAGbuild.h
@@ -14,6 +14,9 @@ struct execution_node {
     // Pointer to device VRAM alloc, NULL if not needed
     void *device_ptr;
 
+    // Pointer to device VRAM alloc, will be the gradient pointer storing the result of backward pass for data i.e on GPU
+    void *device_ptr_grad;
+    
     // Boolean flag storing weather it will need to be transfered based upon reading the childs OPS
     bool to_device_needed;
 
diff --git a/src/core/include/graph/assignBackend.h b/src/core/include/graph/assignBackend.h
index ee34501..91f8d3f 100644
--- a/src/core/include/graph/assignBackend.h
+++ b/src/core/include/graph/assignBackend.h
@@ -1,7 +1,9 @@
 #pragma once
 #include <vector>
 #include "../third_party/json.hpp"
+
 using json = nlohmann::json;
+
 void assignBackend(execution_node_t *e);
 
 device_type assignDevice(uint8_t ndims, uint32_t *dims, tensor_op_t op, uint32_t nvalues, json &data);
@@ -9,3 +11,6 @@ device_type assignDevice(uint8_t ndims, uint32_t *dims, tensor_op_t op, uint32_t
 int32_t getTheExecutionNodeIndex(uint32_t id, std::vector<execution_node_t *> &nodes);
 
 void assignBackendGraph(tensor_pool_t *pool, std::vector<execution_node_t *> &nodes);
+
+
+void assignGradMemory(tensor_pool_t *pool_grad_cpu, tensor_pool_t *pool_grad_gpu, std::vector<execution_node_t *> &nodes);
diff --git a/src/core/include/tensor/tensor.h b/src/core/include/tensor/tensor.h
index c6f05f9..05ee054 100644
--- a/src/core/include/tensor/tensor.h
+++ b/src/core/include/tensor/tensor.h
@@ -12,6 +12,9 @@ enum class tensor_op_t {
     ADD,              // Add a and b
     BROADCAST_ADD,    // Performs broadcasting during addition
     RELU,             // Activation function
+    SUB,              // Subtract two tensor of same shape
+    MEAN,             // Returns a scalar value mean of the tensor
+    SQUARE,           // Square the tensor element-wise
 };
 
 struct tensor_instance {
@@ -47,7 +50,11 @@ struct tensor_instance {
 
     // If tensor is transposed
     bool is_transposed;
-
+    
+    // Writing it for common interface and ease of access
+    // grad is a data-only tensor. Fields op, a, b,
+    // stateTracker, grad_compute, grad are always NULL/zero.
+    // Only data, nvalues, ndims, dims, dtype are valid.
     // For storing autograd result
     tensor_t *grad;
 
@@ -59,7 +66,7 @@ struct tensor_instance {
 // a scalar is created. If elems is NULL then the tensor is created with zeros.
 // If elems is not NULL then it is assigned as initial values.
 tensor_t *tensor_dtype_create(tensor_pool_t *pool, tensor_dtype_t dtype, uint32_t num_dims,
-                              uint32_t *dims, void *elems);
+                              uint32_t *dims, void *elems, bool grad_status = true);
 
 // Evaluate the tensor, return true on success
 bool tensor_evaluate( tensor_pool_t *pool,tensor_t *t,  float *d_a, float *d_b, float *d_res);
@@ -68,5 +75,5 @@ bool tensor_evaluate( tensor_pool_t *pool,tensor_t *t,  float *d_a, float *d_b,
 size_t tensor_dtype_sizeof(tensor_dtype_t dtype);
 
 tensor_t *tensor_create(tensor_pool_t *pool, tensor_dtype_t dtype, uint32_t num_dims,
-                        uint32_t *dims, void *elems);
+                        uint32_t *dims, void *elems, bool grad_status);
 #endif // !PRIVATE_TENSOR_H
diff --git a/src/core/tensor/tensor.cpp b/src/core/tensor/tensor.cpp
index b8761de..f90cfe9 100644
--- a/src/core/tensor/tensor.cpp
+++ b/src/core/tensor/tensor.cpp
@@ -33,7 +33,7 @@ size_t tensor_dtype_sizeof(tensor_dtype_t dtype) {
 }
 
 tensor_t *tensor_dtype_create(tensor_pool_t *pool, tensor_dtype_t dtype, uint32_t ndims,
-                              uint32_t *dims, void *elems) {
+                              uint32_t *dims, void *elems, bool grad_status) {
     assert(pool != NULL);
     assert(tensor_dtype_sizeof(dtype) > 0);
 
@@ -99,7 +99,7 @@ tensor_t *tensor_dtype_create(tensor_pool_t *pool, tensor_dtype_t dtype, uint32_
     t->b = NULL;
     t->stateTracker = 0;
     t->device = device_type::CPU;
-    t->grad_compute = false;
+    t->grad_compute = grad_status;
     t->is_transposed = false;
     // Return success
     return t;
@@ -107,8 +107,8 @@ tensor_t *tensor_dtype_create(tensor_pool_t *pool, tensor_dtype_t dtype, uint32_
 
 // Create float32 tensor
 tensor_t *tensor_create(tensor_pool_t *pool, tensor_dtype_t dtype, uint32_t num_dims,
-                        uint32_t *dims, void *elems) {
-    return tensor_dtype_create(pool, dtype, num_dims, dims, elems);
+                        uint32_t *dims, void *elems, bool grad_status) {
+    return tensor_dtype_create(pool, dtype, num_dims, dims, elems, grad_status);
 }
 
 bool tensor_evaluate(tensor_pool_t *pool, tensor_t *t,  [[maybe_unused]]float *d_a, [[maybe_unused]]float *d_b, [[maybe_unused]]float *d_res) {
@@ -161,6 +161,20 @@ bool tensor_evaluate(tensor_pool_t *pool, tensor_t *t,  [[maybe_unused]]float *d
         assert(t->a != NULL);
         success = tensor_op_relu(pool, t);
         break;
+    case tensor_op_t::SUB:
+        assert(t->a != NULL);
+        assert(t->b != NULL);
+        assert(t->a->dtype == t->b->dtype);
+        success = tensor_op_sub(pool, t);
+        break;
+    case tensor_op_t::MEAN:
+        assert(t->a != NULL);
+        success = tensor_op_mean(pool, t);
+        break;
+    case tensor_op_t::SQUARE:
+        assert(t->a != NULL);
+        success = tensor_op_square(pool, t);
+        break;
     default:
         assert(false);
     }
@@ -221,6 +235,10 @@ uint8_t tensor_get_ndims(tensor_t *t) { return t->ndims; }
 uint32_t *tensor_get_dims(tensor_t *t) { return t->dims; }
 
 void tensor_print_data(tensor_t *t) {
+    if (t->ndims == 0) {
+        std::cout << ((float *)t->data)[0] << "\n";
+        return;
+    }
     for (uint32_t i = 0; i < t->dims[0]; i++) {
         for (uint32_t j = 0; j < t->dims[1]; j++) {
             uint32_t index = i * t->dims[1] + j;
diff --git a/src/internal_header.h b/src/internal_header.h
index eb2f722..625597a 100644
--- a/src/internal_header.h
+++ b/src/internal_header.h
@@ -7,11 +7,15 @@
 #include "./core/include/tensor/tensor.h"
 #include "./core/include/pool/pool.h"
 #include "backend_cpu/include/add.h"
+#include "backend_cpu/include/sub.h"
+#include "backend_cpu/include/mean.h"
 #include "backend_cpu/include/debug.h"
+#include "backend_cpu/include/square.h"
 #include "backend_cpu/include/mul.h"
 #include "backend_cpu/include/scalar.h"
 #include "backend_cpu/include/transpose.h"
 #include "backend_cpu/include/relu.h"
+#include "backend_cpu/include/backprop/backprop_b.h"
 #include "./core/include/graph/DAGbuild.h"
 #include "./core/include/graph/assignBackend.h"
 #include "./core/include/JSON/json_utils.h"
diff --git a/summary.txt b/summary.txt
new file mode 100644
index 0000000..d009af4
--- /dev/null
+++ b/summary.txt
@@ -0,0 +1,10 @@
+0EPOCH13
+1000EPOCH13
+2000EPOCH13
+3000EPOCH13
+4000EPOCH13
+5000EPOCH13
+6000EPOCH13
+7000EPOCH13
+8000EPOCH13
+9000EPOCH13