refactor

test_functional.mojo

@@ -1,12 +1,7 @@
 from time.time import now
 from tensor import TensorShape
 
-from voodoo.initializers import HeNormal, RandomUniform
-from voodoo import (
-    Tensor,
-    get_activation_code,
-    get_loss_code,
-)
+from voodoo.core import Tensor, HeNormal, RandomUniform, SGD
 from voodoo.utils import (
     info,
     clear,
@@ -53,7 +48,7 @@ fn main() raises:
         var computed_loss = loss.forward_static()
         avg_loss += computed_loss[0]
         loss.backward()
-        loss.optimize["sgd", 0.01]()
+        loss.optimize[SGD[0.01]]()
 
         if epoch % every == 0:
             var bar = String("")

test_layer.mojo

@@ -1,14 +1,12 @@
-from voodoo import Tensor, get_loss_code, Graph
-from voodoo.layers.Dense import Dense
-from voodoo.layers.Dropout import Dropout
-from voodoo.layers.LeakyReLu import LeakyReLu
+from time.time import now
+from tensor import TensorShape
+
+from voodoo.core import Tensor, HeNormal, RandomUniform, SGD
+from voodoo.core.layers import Dense, LeakyReLu
 from voodoo.utils import (
     info,
     clear,
 )
-from time.time import now
-from tensor import TensorShape
-from voodoo.initializers import HeNormal, RandomUniform
 
 
 fn nanoseconds_to_seconds(t: Int) -> Float64:
@@ -26,7 +24,6 @@ fn main() raises:
         weight_initializer = HeNormal[1],
         bias_initializer = HeNormal[32],
     ]()
-    let dropout_layer = Dropout[dropout_rate=0.05]()
     let leaky_relu = LeakyReLu[
         in_neurons=32,
         out_neurons=32,
@@ -54,7 +51,6 @@ fn main() raises:
     let true_vals = Tensor[data_shape, RandomUniform[0, 1]]()
 
     var x = input_layer.forward(input)
-    x = dropout_layer.forward(x)
     x = leaky_relu.forward(x)
     x = dense_layer.forward(x)
     x = output_layer.forward(x)
@@ -71,7 +67,7 @@ fn main() raises:
         var computed_loss = loss.forward_static()
         avg_loss += computed_loss[0]
         loss.backward()
-        loss.optimize["sgd", 0.01]()
+        loss.optimize[SGD[0.01]]()
 
         if epoch % every == 0:
             var bar = String("")

voodoo/__init__.mojo

@@ -1,5 +0,0 @@
-from voodoo.node import Node
-from voodoo.graph import Graph
-from voodoo.tensor import Tensor
-from voodoo.code_lookup import get_activation_code, get_loss_code
-from voodoo.utils import Vector

voodoo/autograd/__init__.mojo

@@ -0,0 +1,2 @@
+from .graph import Graph
+from .node import Node

voodoo/graph.mojo → voodoo/autograd/graph.mojo

@@ -2,12 +2,13 @@ from memory import memset_zero, memset
 from math import log2, exp2, ceil, round
 from tensor import TensorShape
 
-from voodoo.node import Node
-from voodoo.kernels import KERNELS
-from voodoo.utils import Vector, get_broadcasted_shape_for_ew_op, warn
-from voodoo.kernels.optimizers import SGD
 from voodoo.constants import MEMORY_POOL_SIZE, OP_TUPLE, BINARY_OP, UNARY_OP
-from voodoo.operator_codes import (
+from voodoo.autograd.kernels import KERNELS
+from voodoo.core import Optimizer
+from voodoo.utils import (
+    Vector,
+    get_broadcasted_shape_for_ew_op,
+    warn,
     copy_code,
     mmul_code,
     conv1d_code,
@@ -513,12 +514,8 @@ struct Graph:
         for i in range(len(self._grad_nodes_order)):
             _ = self.backward_recursive(self._nodes[self._grad_nodes_order[i]])
 
-    fn optimizer_step[type: String, learning_rate: Float32](self) raises:
-        if type == "sgd":
-            SGD[learning_rate].step(self._nodes)
-        else:
-            warn("Invalid optimizer: " + type + " using sgd\n")
-            SGD[learning_rate].step(self._nodes)
+    fn optimizer_step[optimizer: Optimizer](self):
+        optimizer.step(self._nodes)
 
     fn copy(inout self, parent1: Node) raises -> Node:
         return self.node[False, True, False, copy_code](

voodoo/autograd/kernels/__init__.mojo

@@ -0,0 +1 @@
+from .kernels import KERNELS

voodoo/kernels/activations.mojo → voodoo/autograd/kernels/activations.mojo

@@ -1,7 +1,7 @@
 from math import exp, log, abs, tanh, cosh, erf
 
 from voodoo.constants import F32_MAX, NELTS
-from voodoo.kernels.generics import GenericActivation
+from voodoo.autograd.kernels.generics import GenericActivation
 
 
 trait Activation:

voodoo/kernels/arithmetic.mojo → voodoo/autograd/kernels/arithmetic.mojo

@@ -1,6 +1,6 @@
 from math import sqrt, exp2, exp, log2, sin, cos, tan, log, asin, acos, atan, sinh, cosh
 
-from voodoo.kernels.generics import GenericArithmetic, GenericBinaryArithmetic
+from voodoo.autograd.kernels.generics import GenericArithmetic, GenericBinaryArithmetic
 
 
 trait Aritmetic:

voodoo/kernels/conv.mojo → voodoo/autograd/kernels/conv.mojo

@@ -1,7 +1,8 @@
 from algorithm import vectorize, tile
 from math import max
 
-from voodoo import Node
+from voodoo.autograd import Node
+from voodoo.utils import Vector
 from voodoo.constants import NELTS, PREFETCH_READ, PREFETCH_WRITE
 
 alias tile_sizes = VariadicList[Int](32, 16, 8, 4, 2, 1)

voodoo/kernels/generics.mojo → voodoo/autograd/kernels/generics.mojo

@@ -6,6 +6,7 @@ from voodoo.utils import (
     strides_a,
     strides_b,
     recursive_broadcast,
+    Vector,
 )
 from voodoo.constants import NELTS, PREFETCH_READ, PREFETCH_WRITE
 

voodoo/kernels/loader.mojo → voodoo/autograd/kernels/kernels.mojo

@@ -1,14 +1,14 @@
-from voodoo.kernels.operations import (
+from .operations import (
     Copy,
     Reshape,
     Transpose,
     Sum,
     Dropout,
 )
-from voodoo.kernels.matmul import MMul
-from voodoo.kernels.maxpool import MaxPool1D, MaxPool2D
-from voodoo.kernels.conv import Conv1D, Conv2D
-from voodoo.kernels.activations import (
+from .matmul import MMul
+from .maxpool import MaxPool1D, MaxPool2D
+from .conv import Conv1D, Conv2D
+from .activations import (
     Relu,
     Sigmoid,
     Softplus,
@@ -25,7 +25,7 @@ from voodoo.kernels.activations import (
     Linear,
     Mish,
 )
-from voodoo.kernels.arithmetic import (
+from .arithmetic import (
     Sqrt,
     Abs,
     Exp2,
@@ -45,9 +45,10 @@ from voodoo.kernels.arithmetic import (
     Div,
     Pow,
 )
-from voodoo.operator_codes import *
-from voodoo.kernels.losses import MSE, MAE, MAPE, MSLE
+from .losses import MSE, MAE, MAPE, MSLE
+
 from voodoo.constants import F32_MAX, UNARY_OP, BINARY_OP, OP_TUPLE, NU, NB
+from voodoo.utils.operator_codes import *
 
 
 @register_passable("trivial")

voodoo/kernels/losses.mojo → voodoo/autograd/kernels/losses.mojo

@@ -1,7 +1,7 @@
 from math import log, abs
 
 from voodoo.constants import EPSILON
-from voodoo.kernels.generics import GenericLoss
+from voodoo.autograd.kernels.generics import GenericLoss
 
 
 trait Loss:

voodoo/kernels/matmul.mojo → voodoo/autograd/kernels/matmul.mojo

@@ -1,7 +1,7 @@
 from algorithm import vectorize
 from math import max, min
 
-from voodoo import Node
+from voodoo.autograd import Node
 from voodoo.utils import recursive_broadcast
 from voodoo.constants import PREFETCH_READ, PREFETCH_WRITE, F32_MAX, NELTS
 

voodoo/kernels/maxpool.mojo → voodoo/autograd/kernels/maxpool.mojo

@@ -1,7 +1,7 @@
 from algorithm import vectorize
 from math import max
 
-from voodoo import Node
+from voodoo.autograd import Node
 from voodoo.constants import PREFETCH_READ, PREFETCH_WRITE, F32_MAX, NELTS
 
 

voodoo/kernels/operations.mojo → voodoo/autograd/kernels/operations.mojo

@@ -1,7 +1,7 @@
 from random import random_float64
 from algorithm import vectorize
 
-from voodoo import Node
+from voodoo.autograd import Node
 from voodoo.constants import NELTS
 
 

voodoo/constants.mojo

@@ -1,6 +1,8 @@
 from math.limit import inf
 from sys.intrinsics import PrefetchOptions
 
+from voodoo.autograd import Node
+
 alias PI = 3.141592653589793238462643383279502884197169399375105820974944592307816406286208998628034825342
 alias MEMORY_POOL_SIZE = 20
 alias EPSILON = 1e-8

voodoo/core/__init__.mojo

@@ -0,0 +1,27 @@
+from .tensor import Tensor
+from .initializers import (
+    Initializer,
+    Constant,
+    Zeros,
+    Ones,
+    GlorotNormal,
+    GlorotUniform,
+    HeNormal,
+    HeUniform,
+    LecunNormal,
+    LecunUniform,
+    RandomNormal,
+    RandomUniform,
+    TruncatedNormal,
+    NoneInitializer,
+)
+from .constraints import (
+    Constraint,
+    MaxNorm,
+    MinMaxNorm,
+    NonNeg,
+    RadialConstraint,
+    UnitNorm,
+    NoneConstraint,
+)
+from .optimizers import Optimizer, SGD

voodoo/constraints.mojo → voodoo/core/constraints.mojo

@@ -2,7 +2,7 @@ from algorithm import vectorize
 from math import sqrt, max, abs
 from tensor import TensorShape
 
-from voodoo.utils import reduce_vector_mul
+from voodoo.utils import Vector, reduce_vector_mul
 from voodoo.constants import NELTS, EPSILON