Implement reduce_sum/prod/max/min scaled translation rules.

graphcore-research · Nov 21, 2023 · f214af2 · f214af2
1 parent 642a58e
commit f214af2
Showing 1 changed file with 16 additions and 0 deletions.
diff --git a/jax_scaled_arithmetics/lax/scaled_ops.py b/jax_scaled_arithmetics/lax/scaled_ops.py
@@ -107,3 +107,19 @@ def scaled_dot_general(
         / contracting_rescale
     )
     return ScaledArray(output_data, output_scale)
+
+
+def scaled_reduce_sum():
+    pass
+
+
+def scaled_reduce_prod():
+    pass
+
+
+def scaled_reduce_max():
+    pass
+
+
+def scaled_reduce_min():
+    pass