[manipulation routines] Implement broadcast_to() for NDArray (#202)

Implements `broadcast_to()` for `NDArray`. Add tests.

It can broadcast an ndarray of any shape to any compatible shape. The
data will be copied into the new array. An example goes as follows.

```mojo
from numojo.prelude import *
from python import Python
fn main() raises:
    var np = Python.import_module("numpy")
    var a = nm.random.rand(Shape(2, 3))
    print(a)
    print(nm.routines.manipulation.broadcast_to(a, Shape(2, 2, 3)))
    print(np.broadcast_to(a.to_numpy(), (2, 2, 3)))
```

```console
[[0.8073 0.5361 0.4442]
 [0.9378 0.1910 0.2421]]
2D-array  Shape(2,3)  Strides(3,1)  DType: f64  C-cont: True  F-cont: False  own data: True

[[[0.8073 0.5361 0.4442]
  [0.9378 0.1910 0.2421]]
 [[0.8073 0.5361 0.4442]
  [0.9378 0.1910 0.2421]]]
3D-array  Shape(2,2,3)  Strides(6,3,1)  DType: f64  C-cont: True  F-cont: False  own data: True

[[[0.8074 0.5361 0.4442]
  [0.9378 0.1911 0.2421]]
 [[0.8074 0.5361 0.4442]
  [0.9378 0.1911 0.2421]]]
```

Author: forfudan

.github/workflows/run_tests.yaml

@@ -31,7 +31,7 @@ jobs:
 
       - name: Install magic
         run: |
-          curl -ssL https://magic.modular.com/ff414efd-16ac-4bf3-8efc-50b059272ab6 | bash
+          curl -ssL https://magic.modular.com/deb181c4-455c-4abe-a263-afcff49ccf67 | bash
       
       - name: Add path
         run: |

.github/workflows/test_pre_commit.yaml

@@ -26,7 +26,7 @@ jobs:
 
       - name: Install magic
         run: |
-          curl -ssL https://magic.modular.com/ff414efd-16ac-4bf3-8efc-50b059272ab6 | bash
+          curl -ssL https://magic.modular.com/deb181c4-455c-4abe-a263-afcff49ccf67 | bash
     
       - name: Add path
         run: |

numojo/__init__.mojo

@@ -194,6 +194,7 @@ from numojo.routines.manipulation import (
     reshape,
     ravel,
     transpose,
+    broadcast_to,
     flip,
 )
 

numojo/core/matrix.mojo

@@ -17,6 +17,7 @@ from python import PythonObject, Python
 from numojo.core.ndarray import NDArray
 from numojo.core.own_data import OwnData
 from numojo.core.utility import _get_offset, _update_flags
+from numojo.routines.manipulation import broadcast_to
 
 # ===----------------------------------------------------------------------===#
 # Matrix struct
@@ -1633,73 +1634,3 @@ fn _logic_func_matrix_matrix_to_matrix[
     var _B = B
 
     return C^
-
-
-fn broadcast_to[
-    dtype: DType
-](A: Matrix[dtype], shape: Tuple[Int, Int]) raises -> Matrix[dtype]:
-    """
-    Broadcase the vector to the given shape.
-
-    Example:
-
-    ```console
-    > from numojo import Matrix
-    > a = Matrix.fromstring("1 2 3", shape=(1, 3))
-    > print(mat.broadcast_to(a, (3, 3)))
-    [[1.0   2.0     3.0]
-     [1.0   2.0     3.0]
-     [1.0   2.0     3.0]]
-    > a = Matrix.fromstring("1 2 3", shape=(3, 1))
-    > print(mat.broadcast_to(a, (3, 3)))
-    [[1.0   1.0     1.0]
-     [2.0   2.0     2.0]
-     [3.0   3.0     3.0]]
-    > a = Matrix.fromstring("1", shape=(1, 1))
-    > print(mat.broadcast_to(a, (3, 3)))
-    [[1.0   1.0     1.0]
-     [1.0   1.0     1.0]
-     [1.0   1.0     1.0]]
-    > a = Matrix.fromstring("1 2", shape=(1, 2))
-    > print(mat.broadcast_to(a, (1, 2)))
-    [[1.0   2.0]]
-    > a = Matrix.fromstring("1 2 3 4", shape=(2, 2))
-    > print(mat.broadcast_to(a, (4, 2)))
-    Unhandled exception caught during execution: Cannot broadcast shape 2x2 to shape 4x2!
-    ```
-    """
-
-    var B = Matrix[dtype](shape)
-    if (A.shape[0] == shape[0]) and (A.shape[1] == shape[1]):
-        B = A
-    elif (A.shape[0] == 1) and (A.shape[1] == 1):
-        B = Matrix.full[dtype](shape, A[0, 0])
-    elif (A.shape[0] == 1) and (A.shape[1] == shape[1]):
-        for i in range(shape[0]):
-            memcpy(
-                dest=B._buf.ptr.offset(shape[1] * i),
-                src=A._buf.ptr,
-                count=shape[1],
-            )
-    elif (A.shape[1] == 1) and (A.shape[0] == shape[0]):
-        for i in range(shape[0]):
-            for j in range(shape[1]):
-                B._store(i, j, A._buf.ptr[i])
-    else:
-        var message = String(
-            "Cannot broadcast shape {}x{} to shape {}x{}!"
-        ).format(A.shape[0], A.shape[1], shape[0], shape[1])
-        raise Error(message)
-    return B^
-
-
-fn broadcast_to[
-    dtype: DType
-](A: Scalar[dtype], shape: Tuple[Int, Int]) raises -> Matrix[dtype]:
-    """
-    Broadcase the scalar to the given shape.
-    """
-
-    var B = Matrix[dtype](shape)
-    B = Matrix.full[dtype](shape, A)
-    return B^

numojo/core/ndarray.mojo

@@ -75,6 +75,7 @@ from numojo.routines.manipulation import reshape, ravel
 #       RawData type is just a wrapper of `UnsafePointer`.
 #       RefData type has an extra property `indices`: getitem(i) -> A[I[i]].
 # TODO: Rename some variables or methods that should not be exposed to users.
+# TODO: Remove 0-d array. Raise errors if operations result in 0-d array.
 # ===----------------------------------------------------------------------===#
 
 
@@ -2323,10 +2324,7 @@ struct NDArray[dtype: DType = DType.float64](
                 min_value,
                 abs(val),
             )
-        number_of_digits = max(
-            int(log10(float(max_value))) + 1,
-            abs(int(log10(float(min_value)))) + 1,
-        )
+        number_of_digits = int(log10(float(max_value))) + 1
         number_of_digits_small_values = abs(int(log10(float(min_value)))) + 1
 
         if dtype.is_floating_point():

numojo/prelude.mojo

@@ -20,9 +20,11 @@ from numojo.prelude import *
 ```
 """
 
-from numojo.core.ndarray import NDArray
-from numojo.core.matrix import Matrix
+import numojo as nm
+
 from numojo.core.item import Item, item
+from numojo.core.matrix import Matrix
+from numojo.core.ndarray import NDArray
 from numojo.core.ndshape import Shape, NDArrayShape
 
 from numojo.core.complex.complex_dtype import CDType

numojo/routines/manipulation.mojo

@@ -1,6 +1,12 @@
+# ===----------------------------------------------------------------------=== #
+# Distributed under the Apache 2.0 License with LLVM Exceptions.
+# See LICENSE and the LLVM License for more information.
+# https://github.com/Mojo-Numerics-and-Algorithms-group/NuMojo/blob/main/LICENSE
+# https://llvm.org/LICENSE.txt
+# ===----------------------------------------------------------------------=== #
+
 """
 Array manipulation routines.
-
 """
 
 from memory import UnsafePointer, memcpy
@@ -12,7 +18,12 @@ from numojo.core.ndshape import NDArrayShape, Shape
 from numojo.core.ndstrides import NDArrayStrides
 import numojo.core.matrix as matrix
 from numojo.core.matrix import Matrix
-from numojo.core.utility import _list_of_flipped_range
+from numojo.core.utility import _list_of_flipped_range, _get_offset
+
+# ===----------------------------------------------------------------------=== #
+# TODO:
+# - When `OwnData` is supported, re-write `broadcast_to()`.`
+# ===----------------------------------------------------------------------=== #
 
 # ===----------------------------------------------------------------------=== #
 # Basic operations
@@ -272,6 +283,141 @@ fn transpose[dtype: DType](A: Matrix[dtype]) -> Matrix[dtype]:
     return B^
 
 
+# ===----------------------------------------------------------------------=== #
+# Changing number of dimensions
+# ===----------------------------------------------------------------------=== #
+
+
+fn broadcast_to[
+    dtype: DType
+](a: NDArray[dtype], shape: NDArrayShape) raises -> NDArray[dtype]:
+    if a.shape.ndim > shape.ndim:
+        raise Error(
+            String("Cannot broadcast shape {} to shape {}!").format(
+                a.shape, shape
+            )
+        )
+
+    # Check whether broadcasting is possible or not.
+    # We compare the shape from the trailing dimensions.
+
+    var b_strides = NDArrayStrides(
+        shape
+    )  # Strides of b when refer to data of a
+
+    for i in range(a.shape.ndim):
+        if a.shape[a.shape.ndim - 1 - i] == shape[shape.ndim - 1 - i]:
+            b_strides[shape.ndim - 1 - i] = a.strides[a.shape.ndim - 1 - i]
+        elif a.shape[a.shape.ndim - 1 - i] == 1:
+            b_strides[shape.ndim - 1 - i] = 0
+        else:
+            raise Error(
+                String("Cannot broadcast shape {} to shape {}!").format(
+                    a.shape, shape
+                )
+            )
+    for i in range(shape.ndim - a.shape.ndim):
+        b_strides[i] = 0
+
+    # Start broadcasting.
+    # TODO: When `OwnData` is supported, re-write this part.
+    # We just need to change the shape and strides and re-use the data.
+
+    var b = NDArray[dtype](shape)  # Construct array of targeted shape.
+    # TODO: `b.strides = b_strides` when OwnData
+
+    # Iterate all items in the new array and fill in correct values.
+    for offset in range(b.size):
+        var remainder = offset
+        var indices = Item(ndim=b.ndim, initialized=False)
+
+        for i in range(b.ndim):
+            indices[i], remainder = divmod(
+                remainder,
+                b.strides[
+                    i
+                ],  # TODO: Change b.strides to NDArrayStrides(b.shape) when OwnData
+            )
+
+        (b._buf.ptr + offset).init_pointee_copy(
+            a._buf.ptr[
+                _get_offset(indices, b_strides)
+            ]  # TODO: Change b_strides to b.strides when OwnData
+        )
+
+    return b^
+
+
+fn broadcast_to[
+    dtype: DType
+](A: Matrix[dtype], shape: Tuple[Int, Int]) raises -> Matrix[dtype]:
+    """
+    Broadcasts the vector to the given shape.
+
+    Example:
+
+    ```console
+    > from numojo import Matrix
+    > a = Matrix.fromstring("1 2 3", shape=(1, 3))
+    > print(mat.broadcast_to(a, (3, 3)))
+    [[1.0   2.0     3.0]
+     [1.0   2.0     3.0]
+     [1.0   2.0     3.0]]
+    > a = Matrix.fromstring("1 2 3", shape=(3, 1))
+    > print(mat.broadcast_to(a, (3, 3)))
+    [[1.0   1.0     1.0]
+     [2.0   2.0     2.0]
+     [3.0   3.0     3.0]]
+    > a = Matrix.fromstring("1", shape=(1, 1))
+    > print(mat.broadcast_to(a, (3, 3)))
+    [[1.0   1.0     1.0]
+     [1.0   1.0     1.0]
+     [1.0   1.0     1.0]]
+    > a = Matrix.fromstring("1 2", shape=(1, 2))
+    > print(mat.broadcast_to(a, (1, 2)))
+    [[1.0   2.0]]
+    > a = Matrix.fromstring("1 2 3 4", shape=(2, 2))
+    > print(mat.broadcast_to(a, (4, 2)))
+    Unhandled exception caught during execution: Cannot broadcast shape 2x2 to shape 4x2!
+    ```
+    """
+
+    var B = Matrix[dtype](shape)
+    if (A.shape[0] == shape[0]) and (A.shape[1] == shape[1]):
+        B = A
+    elif (A.shape[0] == 1) and (A.shape[1] == 1):
+        B = Matrix.full[dtype](shape, A[0, 0])
+    elif (A.shape[0] == 1) and (A.shape[1] == shape[1]):
+        for i in range(shape[0]):
+            memcpy(
+                dest=B._buf.ptr.offset(shape[1] * i),
+                src=A._buf.ptr,
+                count=shape[1],
+            )
+    elif (A.shape[1] == 1) and (A.shape[0] == shape[0]):
+        for i in range(shape[0]):
+            for j in range(shape[1]):
+                B._store(i, j, A._buf.ptr[i])
+    else:
+        var message = String(
+            "Cannot broadcast shape {}x{} to shape {}x{}!"
+        ).format(A.shape[0], A.shape[1], shape[0], shape[1])
+        raise Error(message)
+    return B^
+
+
+fn broadcast_to[
+    dtype: DType
+](A: Scalar[dtype], shape: Tuple[Int, Int]) raises -> Matrix[dtype]:
+    """
+    Broadcasts the scalar to the given shape.
+    """
+
+    var B = Matrix[dtype](shape)
+    B = Matrix.full[dtype](shape, A)
+    return B^
+
+
 # ===----------------------------------------------------------------------=== #
 # Rearranging elements
 # ===----------------------------------------------------------------------=== #

tests/routines/test_manipulation.mojo

@@ -75,3 +75,19 @@ def test_transpose():
         np.transpose(Anp, [1, 3, 0, 2]),
         "4-d `transpose` with arbitrary `axes` is broken.",
     )
+
+
+def test_broadcast():
+    var np = Python.import_module("numpy")
+    var a = nm.random.rand(Shape(2, 1, 3))
+    var Anp = a.to_numpy()
+    check(
+        nm.broadcast_to(a, Shape(2, 2, 3)),
+        np.broadcast_to(a.to_numpy(), (2, 2, 3)),
+        "`broadcast_to` fails.",
+    )
+    check(
+        nm.broadcast_to(a, Shape(2, 2, 2, 3)),
+        np.broadcast_to(a.to_numpy(), (2, 2, 2, 3)),
+        "`broadcast_to` fails.",
+    )