Implement scaled_concatenate LAX operation. (#15)

`scaled_concatenate` is using a max rescaling.

jax_scaled_arithmetics/lax/scaled_ops.py

@@ -1,6 +1,7 @@
 # Copyright (c) 2023 Graphcore Ltd. All rights reserved.
 from typing import Optional, Sequence
 
+import jax.numpy as jnp
 from jax import lax
 
 from jax_scaled_arithmetics import core
@@ -20,6 +21,18 @@ def scaled_convert_element_type(A: ScaledArray, new_dtype: DTypeLike, weak_type:
     return ScaledArray(lax.convert_element_type(A.data, new_dtype=new_dtype), A.scale)
 
 
+@core.register_scaled_lax_op
+def scaled_concatenate(operands: Sequence[ScaledArray], dimension: int) -> ScaledArray:
+    # TODO: inputs checking (dtype and cie).
+    scales = jnp.array([v.scale for v in operands])
+    # Max rescaling of the collection of operands.
+    # TODO: explore alternative strategies?
+    scale_max = jnp.max(scales)
+    datas = [v.data * (v.scale / scale_max) for v in operands]
+    data_concat = lax.concatenate(datas, dimension=dimension)
+    return ScaledArray(data_concat, scale_max)
+
+
 @core.register_scaled_lax_op
 def scaled_slice(
     A: ScaledArray, start_indices: Sequence[int], limit_indices: Sequence[int], strides: Optional[Sequence[int]] = None

tests/lax/test_scaled_ops.py

@@ -6,6 +6,7 @@
 from jax_scaled_arithmetics.core import ScaledArray, scaled_array
 from jax_scaled_arithmetics.lax import (
     scaled_broadcast_in_dim,
+    scaled_concatenate,
     scaled_convert_element_type,
     scaled_mul,
     scaled_slice,
@@ -21,6 +22,14 @@ def test__scaled_broadcast_in_dim__proper_scaling(self):
         npt.assert_array_equal(z.scale, x.scale)
         npt.assert_array_almost_equal(z.data, x.data.reshape((5, 1)))
 
+    def test__scaled_concatenate__proper_scaling(self):
+        x = scaled_array(np.random.rand(2, 3), 0.5, dtype=np.float32)
+        y = scaled_array(np.random.rand(5, 3), 2, dtype=np.float32)
+        z = scaled_concatenate([x, y], dimension=0)
+        assert isinstance(z, ScaledArray)
+        npt.assert_array_equal(z.scale, y.scale)
+        npt.assert_array_almost_equal(z, np.concatenate([x, y], axis=0))
+
     def test__scaled_convert_element_type__proper_scaling(self):
         x = scaled_array(np.random.rand(5), 2, dtype=np.float32)
         z = scaled_convert_element_type(x, new_dtype=np.float16)