Refactor layer shapes and remove print statement

test_conv.mojo

@@ -73,8 +73,6 @@ fn main() raises:
     let x0 = conv_layer_one.forward(input)
     let x1 = max_pool_one.forward(x0)
     var x2 = flatten.forward(x1)
-    x2.print()
-    return
     let x3 = dense1.forward(x2)
     let x4 = dense2.forward(x3)
     var loss = x4.compute_loss["mse"](true_vals)

voodoo/autograd/node.mojo

@@ -59,7 +59,7 @@ struct Node:
 
         var cap = 1
         for i in range(len(shape)):
-            cap *= shape[i]
+            cap *= shape[i] if shape[i] > 0 else 1
 
         let strides = Vector[Int](len(shape))
         strides[len(shape) - 1] = 1

voodoo/core/layers/activation.mojo

@@ -3,8 +3,6 @@ from voodoo.utils import get_activation_code
 
 
 struct Activation[
-    in_neurons: Int,
-    out_neurons: Int,
     activation: String,
 ]():
     fn __init__(

voodoo/core/layers/conv1D.mojo

@@ -46,7 +46,7 @@ struct Conv1D[
     ) raises -> Tensor[
         TensorShape(
             x.shape[0],
-            in_channels,
+            x.shape[1],
             (x.shape[2] - kernel_width + 2 * padding) // stride + 1,
         ),
         NoneInitializer,
@@ -57,7 +57,7 @@ struct Conv1D[
         let res = x.conv_1d[
             TensorShape(
                 x.shape[0],
-                in_channels,
+                x.shape[1],
                 (x.shape[2] - kernel_width + 2 * padding) // stride + 1,
             )
         ](self.W, self.padding, self.stride)

voodoo/core/layers/conv2D.mojo

@@ -47,9 +47,10 @@ struct Conv2D[
         self, x: Tensor
     ) raises -> Tensor[
         TensorShape(
-            in_channels,
-            (x.shape[1] - kernel_width + 2 * padding) // stride + 1,
-            (x.shape[2] - kernel_height + 2 * padding) // stride + 1,
+            x.shape[0],
+            x.shape[1],
+            (x.shape[2] - kernel_width + 2 * self.padding) // self.stride + 1,
+            (x.shape[3] - kernel_height + 2 * self.padding) // self.stride + 1,
         ),
         NoneInitializer,
         NoneConstraint,
@@ -58,9 +59,10 @@ struct Conv2D[
     ]:
         let res = x.conv_2d[
             TensorShape(
-                in_channels,
-                (x.shape[1] - kernel_width + 2 * padding) // stride + 1,
-                (x.shape[2] - kernel_height + 2 * padding) // stride + 1,
+                x.shape[0],
+                x.shape[1],
+                (x.shape[2] - kernel_width + 2 * self.padding) // self.stride + 1,
+                (x.shape[3] - kernel_height + 2 * self.padding) // self.stride + 1,
             )
         ](self.W, self.padding, self.stride).compute_activation[self.activation]()
 

voodoo/core/layers/maxPool2D.mojo

@@ -19,9 +19,9 @@ struct MaxPool2D[
     ) raises -> Tensor[
         TensorShape(
             x.shape[0],
-            (x.shape[1] - kernel_width + 2 * padding) // stride + 1,
-            (x.shape[2] - kernel_height + 2 * padding) // stride + 1,
-            x.shape[3],
+            x.shape[1],
+            (x.shape[2] - kernel_width + 2 * padding) // stride + 1,
+            (x.shape[3] - kernel_height + 2 * padding) // stride + 1,
         ),
         NoneInitializer,
         NoneConstraint,
@@ -31,9 +31,9 @@ struct MaxPool2D[
         let res = x.maxpool_2d[
             TensorShape(
                 x.shape[0],
-                (x.shape[1] - kernel_width + 2 * padding) // stride + 1,
-                (x.shape[2] - kernel_height + 2 * padding) // stride + 1,
-                x.shape[3],
+                x.shape[1],
+                (x.shape[2] - kernel_width + 2 * padding) // stride + 1,
+                (x.shape[3] - kernel_height + 2 * padding) // stride + 1,
             )
         ](
             StaticIntTuple[2](kernel_width, kernel_height),