Implement onnx MeanVarianceNormalization (tinygrad#943)

awstone · Jun 6, 2023 · 664d6cc · 664d6cc
1 parent 3bb38c3
commit 664d6cc
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/extra/onnx_ops.py b/extra/onnx_ops.py
@@ -222,4 +222,9 @@ def CastLike(input, target_type):
   assert isinstance(target_type, Tensor), "can only CastLike Tensor"
   return input
 
-def Binarizer(input, threshold=0.0): return input > threshold
+def Binarizer(input, threshold=0.0): return input > threshold
+
+def MeanVarianceNormalization(input, axis=(0, 2, 3)):
+  data_mean = input.mean(axis=axis, keepdim=True)
+  std = ((input**2).mean(axis=axis, keepdim=True) - data_mean**2).sqrt()
+  return (input - data_mean) / (std + 1e-9)