Implement basic scaled operations for MLP model.

TODO

jax_scaled_arithmetics/lax/scaled_ops.py

@@ -50,3 +50,23 @@ def scaled_transpose(A: ScaledArray, permutation: Sequence[int]) -> ScaledArray:
 @core.register_scaled_lax_op
 def scaled_mul(A: ScaledArray, B: ScaledArray) -> ScaledArray:
     return ScaledArray(A.data * B.data, A.scale * B.scale)
+
+
+def scaled_add(A: ScaledArray, B: ScaledArray) -> ScaledArray:
+    output_scale = lax.sqrt(A.scale**2 + B.scale**2)
+    # check correct type output if mismatch between data and scale precision
+    output_data = (A.scale / output_scale) * A.data + (B.scale / output_scale) * B.data
+    return ScaledArray(output_data, output_scale)
+
+
+def scaled_sub(A: ScaledArray, B: ScaledArray) -> ScaledArray:
+    output_scale = lax.sqrt(A.scale**2 + B.scale**2)
+    # check correct type output if mismatch between data and scale precision
+    output_data = (A.scale / output_scale) * A.data - (B.scale / output_scale) * B.data
+    return ScaledArray(output_data, output_scale)
+
+
+def scaled_dot(A: ScaledArray, B: ScaledArray) -> ScaledArray:
+    output_scale = A.scale * B.scale * lax.sqrt(A.data.shape[-1])
+    output_data = lax.dot(A.data, B.data) / lax.sqrt(A.data.shape[-1])
+    return ScaledArray(output_data, output_scale)