Adding tile function (#880)

Author: prady0t

ci/Numba-array-api-xfails.txt

@@ -70,15 +70,13 @@ array_api_tests/test_creation_functions.py::test_empty_like
 array_api_tests/test_data_type_functions.py::test_finfo[complex64]
 array_api_tests/test_manipulation_functions.py::test_squeeze
 array_api_tests/test_has_names.py::test_has_names[utility-diff]
-array_api_tests/test_has_names.py::test_has_names[manipulation-tile]
 array_api_tests/test_has_names.py::test_has_names[manipulation-unstack]
 array_api_tests/test_has_names.py::test_has_names[statistical-cumulative_sum]
 array_api_tests/test_has_names.py::test_has_names[statistical-cumulative_prod]
 array_api_tests/test_has_names.py::test_has_names[indexing-take_along_axis]
 array_api_tests/test_has_names.py::test_has_names[searching-count_nonzero]
 array_api_tests/test_has_names.py::test_has_names[searching-searchsorted]
 array_api_tests/test_signatures.py::test_func_signature[diff]
-array_api_tests/test_signatures.py::test_func_signature[tile]
 array_api_tests/test_signatures.py::test_func_signature[unstack]
 array_api_tests/test_signatures.py::test_func_signature[take_along_axis]
 array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i is +infinity and isfinite(x2_i) and x2_i > 0) -> +infinity]
@@ -107,7 +105,6 @@ array_api_tests/test_array_object.py::test_getitem_arrays_and_ints_2[1]
 array_api_tests/test_array_object.py::test_getitem_arrays_and_ints_2[None]
 array_api_tests/test_searching_functions.py::test_count_nonzero
 array_api_tests/test_searching_functions.py::test_searchsorted
-array_api_tests/test_manipulation_functions.py::test_tile
 array_api_tests/test_signatures.py::test_func_signature[cumulative_sum]
 array_api_tests/test_signatures.py::test_func_signature[cumulative_prod]
 array_api_tests/test_manipulation_functions.py::test_unstack

sparse/numba_backend/__init__.py

@@ -127,6 +127,7 @@
     std,
     sum,
     tensordot,
+    tile,
     var,
     vecdot,
     zeros,
@@ -337,6 +338,7 @@
     "zeros",
     "zeros_like",
     "repeat",
+    "tile",
 ]
 
 

sparse/numba_backend/_common.py

@@ -3147,3 +3147,42 @@ def repeat(a, repeats, axis=None):
     if not axis_is_none:
         return a.reshape(new_shape)
     return a.reshape(new_shape).flatten()
+
+
+def tile(a, reps):
+    """
+    Constructs an array by tiling an input array.
+
+    Parameters
+    ----------
+    a : SparseArray
+        Input sparse arrays.
+    reps : int or tuple[int, ...]
+        The number of repetitions for each dimension.
+        If an integer, the same number of repetitions is applied to all dimensions.
+
+    Returns
+    -------
+    out : SparseArray
+        A tiled output array.
+    """
+    if not isinstance(a, SparseArray):
+        a = as_coo(a)
+
+    if isinstance(reps, int):
+        reps = (reps,)
+    reps = tuple(reps)
+
+    if a.ndim == 0:
+        a = a.reshape((1,))
+
+    if len(reps) < a.ndim:
+        reps = (1,) * (a.ndim - len(reps)) + reps
+    elif len(reps) > a.ndim:
+        a = a.reshape((1,) * (len(reps) - a.ndim) + a.shape)
+
+    shape = a.shape
+    ndim = len(reps)
+    a = a.reshape(tuple(np.column_stack(([1] * ndim, shape)).reshape(-1)))
+    a = a.broadcast_to(tuple(np.column_stack((reps, shape)).reshape(-1)))
+    return a.reshape(tuple(np.multiply(reps, shape)))

sparse/numba_backend/tests/test_coo.py

@@ -1956,3 +1956,28 @@ def test_repeat(ndim, repeats):
     print(f"Expected: {expected}, Actual: {actual}")
     assert actual.shape == expected.shape
     np.testing.assert_array_equal(actual, expected)
+
+
+def test_tile_invalid_input():
+    a = np.eye(3)
+    assert isinstance(sparse.tile(a, 2), sparse.COO)
+
+
+@pytest.mark.parametrize(
+    "arr,reps",
+    [
+        (np.array([1, 2, 3]), (3,)),
+        (np.array([4, 5, 6, 7]), 3),
+        (np.array(1), 3),
+        (np.array([[1, 2], [3, 4]]), (2, 2)),
+        (np.array([[[1], [2]], [[3], [4]]]), (2, 1, 2)),
+        (np.random.default_rng(0).integers(0, 10, (2, 1, 3)), (2, 2, 2)),
+        (np.random.default_rng(1).integers(0, 5, (1, 3, 1, 2)), (2, 1, 3, 1)),
+    ],
+)
+def test_tile(arr, reps):
+    sparse_arr = sparse.COO.from_numpy(arr)
+    expected = np.tile(arr, reps)
+    result = sparse.tile(sparse_arr, reps).todense()
+
+    np.testing.assert_array_equal(result, expected, err_msg=f"Mismatch for shape={arr.shape}, reps={reps}")

sparse/numba_backend/tests/test_namespace.py

@@ -156,6 +156,7 @@ def test_namespace():
         "tan",
         "tanh",
         "tensordot",
+        "tile",
         "tril",
         "triu",
         "trunc",