fixes and small changes to constraints / initializers

Benny-Nottonson · Mar 14, 2024 · be8129d · be8129d
1 parent 088edc0
commit be8129d
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 98 deletions.
diff --git a/voodoo/autograd/node.mojo b/voodoo/autograd/node.mojo
@@ -204,6 +204,9 @@ struct Node:
     fn get_other_params(self) -> Vector[Int]:
         return self._other_params
 
+    fn is_zero(self) -> Bool:
+        return sum(Buffer[DType.float32](self._data_ptr[0], self._cap)) == 0.0
+
     fn fill(self, val: Float32):
         for i in range(self._cap):
             self._data_ptr[0].store(i, val)

diff --git a/voodoo/core/constraints.mojo b/voodoo/core/constraints.mojo
@@ -11,10 +11,6 @@ trait Constraint(CollectionElement):
     fn constrain[shape: Vector[Int]](data: DTypePointer[DType.float32]) -> None:
         ...
 
-    @staticmethod
-    fn key() -> String:
-        ...
-
 
 @register_passable("trivial")
 struct MaxNorm[max_value: Float32](Constraint):
@@ -41,10 +37,6 @@ struct MaxNorm[max_value: Float32](Constraint):
 
         vectorize[vec, NELTS](num_elements)
 
-    @staticmethod
-    fn key() -> String:
-        return "MaxNorm"
-
 
 @register_passable("trivial")
 struct MinMaxNorm[min_value: Float32, max_value: Float32](Constraint):
@@ -76,10 +68,6 @@ struct MinMaxNorm[min_value: Float32, max_value: Float32](Constraint):
 
         vectorize[vec, NELTS](num_elements)
 
-    @staticmethod
-    fn key() -> String:
-        return "MinMaxNorm"
-
 
 @register_passable("trivial")
 struct NonNeg[](Constraint):
@@ -97,10 +85,6 @@ struct NonNeg[](Constraint):
 
         vectorize[vec, NELTS](num_elements)
 
-    @staticmethod
-    fn key() -> String:
-        return "NonNeg"
-
 
 @register_passable("trivial")
 struct RadialConstraint[](Constraint):
@@ -124,10 +108,6 @@ struct RadialConstraint[](Constraint):
 
         vectorize[vec, NELTS](num_elements)
 
-    @staticmethod
-    fn key() -> String:
-        return "RadialConstraint"
-
 
 @register_passable("trivial")
 struct UnitNorm[](Constraint):
@@ -146,10 +126,6 @@ struct UnitNorm[](Constraint):
 
         vectorize[vec, NELTS](num_elements)
 
-    @staticmethod
-    fn key() -> String:
-        return "UnitNorm"
-
 
 @register_passable("trivial")
 struct NoneConstraint[](Constraint):
@@ -160,7 +136,3 @@ struct NoneConstraint[](Constraint):
     @staticmethod
     fn constrain[shape: Vector[Int]](data: DTypePointer[DType.float32]):
         ...
-
-    @staticmethod
-    fn key() -> String:
-        return "NoneConstraint"
diff --git a/voodoo/core/initializers.mojo b/voodoo/core/initializers.mojo
@@ -15,10 +15,6 @@ trait Initializer(CollectionElement):
     fn initialize[shape: Vector[Int]](data: DTypePointer[DType.float32]):
         ...
 
-    @staticmethod
-    fn key() -> String:
-        ...
-
 
 @register_passable("trivial")
 struct Constant[value: Float64](Initializer):
@@ -36,10 +32,6 @@ struct Constant[value: Float64](Initializer):
 
         vectorize[vec, NELTS](reduce_vector_mul[shape]())
 
-    @staticmethod
-    fn key() -> String:
-        return "Constant"
-
 
 @register_passable("trivial")
 struct Zeros[](Initializer):
@@ -57,10 +49,6 @@ struct Zeros[](Initializer):
 
         vectorize[vec, NELTS](reduce_vector_mul[shape]())
 
-    @staticmethod
-    fn key() -> String:
-        return "Zeroes"
-
 
 @register_passable("trivial")
 struct Ones[](Initializer):
@@ -78,10 +66,6 @@ struct Ones[](Initializer):
 
         vectorize[vec, NELTS](reduce_vector_mul[shape]())
 
-    @staticmethod
-    fn key() -> String:
-        return "Ones"
-
 
 @register_passable("trivial")
 struct GlorotNormal[input_units: Float64, output_units: Float64](
@@ -101,10 +85,6 @@ struct GlorotNormal[input_units: Float64, output_units: Float64](
             (2.0 / (input_units + output_units)) ** 0.5,
         )
 
-    @staticmethod
-    fn key() -> String:
-        return "GlorotNormal"
-
 
 @register_passable("trivial")
 struct GlorotUniform[input_units: Float64, output_units: Float64](
@@ -127,10 +107,6 @@ struct GlorotUniform[input_units: Float64, output_units: Float64](
 
         vectorize[vec, NELTS](reduce_vector_mul[shape]())
 
-    @staticmethod
-    fn key() -> String:
-        return "GlorotUniform"
-
 
 @register_passable("trivial")
 struct HeNormal[input_units: Float64](Initializer):
@@ -143,10 +119,6 @@ struct HeNormal[input_units: Float64](Initializer):
         seed()
         randn(data, reduce_vector_mul[shape](), 0.0, (2.0 / input_units) ** 0.5)
 
-    @staticmethod
-    fn key() -> String:
-        return "HeNormal"
-
 
 @register_passable("trivial")
 struct HeUniform[input_units: Float64](Initializer):
@@ -167,10 +139,6 @@ struct HeUniform[input_units: Float64](Initializer):
 
         vectorize[vec, NELTS](reduce_vector_mul[shape]())
 
-    @staticmethod
-    fn key() -> String:
-        return "HeUniform"
-
 
 @register_passable("trivial")
 struct Identity[](Initializer):
@@ -192,10 +160,6 @@ struct Identity[](Initializer):
 
         vectorize[vec, NELTS](n * m)
 
-    @staticmethod
-    fn key() -> String:
-        return "Identity"
-
 
 @register_passable("trivial")
 struct LecunNormal[input_units: Float64](Initializer):
@@ -208,10 +172,6 @@ struct LecunNormal[input_units: Float64](Initializer):
         seed()
         randn(data, reduce_vector_mul[shape](), 0.0, (1.0 / input_units) ** 0.5)
 
-    @staticmethod
-    fn key() -> String:
-        return "LecunNormal"
-
 
 @register_passable("trivial")
 struct LecunUniform[input_units: Float64](Initializer):
@@ -232,10 +192,6 @@ struct LecunUniform[input_units: Float64](Initializer):
 
         vectorize[vec, NELTS](reduce_vector_mul[shape]())
 
-    @staticmethod
-    fn key() -> String:
-        return "LecunUniform"
-
 
 @register_passable("trivial")
 struct RandomNormal[mean: Float64, std: Float64](Initializer):
@@ -248,10 +204,6 @@ struct RandomNormal[mean: Float64, std: Float64](Initializer):
         seed()
         randn(data, reduce_vector_mul[shape](), mean, std)
 
-    @staticmethod
-    fn key() -> String:
-        return "RandomNormal"
-
 
 @register_passable("trivial")
 struct RandomUniform[low: Float64, high: Float64](Initializer):
@@ -269,10 +221,6 @@ struct RandomUniform[low: Float64, high: Float64](Initializer):
 
         vectorize[vec, NELTS](reduce_vector_mul[shape]())
 
-    @staticmethod
-    fn key() -> String:
-        return "RandomUniform"
-
 
 @register_passable("trivial")
 struct TruncatedNormal[mean: Float64, std: Float64](Initializer):
@@ -295,10 +243,6 @@ struct TruncatedNormal[mean: Float64, std: Float64](Initializer):
 
         vectorize[vec, NELTS](reduce_vector_mul[shape]())
 
-    @staticmethod
-    fn key() -> String:
-        return "TruncatedNormal"
-
 
 @register_passable("trivial")
 struct NoneInitializer[](Initializer):
@@ -309,7 +253,3 @@ struct NoneInitializer[](Initializer):
     @staticmethod
     fn initialize[shape: Vector[Int]](data: DTypePointer[DType.float32]):
         ...
-
-    @staticmethod
-    fn key() -> String:
-        return "NoneInitializer"
diff --git a/voodoo/core/tensor.mojo b/voodoo/core/tensor.mojo
@@ -38,16 +38,7 @@ struct Tensor[
         self.node = self.graph.node[False, is_static, is_single, -1](
             shape, Vector[Int]()
         )
-
-        @parameter
-        if initializer.key() != "NoneInitializer":
-            initializer.initialize[shape](self.node.get_data())
-
-            self.node.set_computed(True)
-
-            @parameter
-            if constraint.key() != "NoneConstraint":
-                constraint.constrain[shape](self.node.get_data())
+        self.refresh()
 
     fn __copyinit__(inout self, other: Self):
         self.graph = other.graph
@@ -99,6 +90,7 @@ struct Tensor[
 
     fn refresh(self) raises:
         initializer.initialize[shape](self.node.get_data())
+        constraint.constrain[shape](self.node.get_data())
 
     fn fill(owned self, val: Float32) -> Self:
         self.node.fill(val)