warp-lang 1.0.1__py3-none-manylinux2014_aarch64.whl → 1.1.0__py3-none-manylinux2014_aarch64.whl

warp/tests/test_indexedarray.py

@@ -1,1134 +1,1134 @@
-# Copyright (c) 2023 NVIDIA CORPORATION.  All rights reserved.
-# NVIDIA CORPORATION and its licensors retain all intellectual property
-# and proprietary rights in and to this software, related documentation
-# and any modifications thereto.  Any use, reproduction, disclosure or
-# distribution of this software and related documentation without an express
-# license agreement from NVIDIA CORPORATION is strictly prohibited.
-
-import unittest
-from typing import Any
-
-import numpy as np
-
-import warp as wp
-from warp.tests.test_array import FillStruct
-from warp.tests.unittest_utils import *
-
-wp.init()
-
-
-@wp.kernel
-def kernel_1d(a: wp.indexedarray(dtype=float), expected: wp.array(dtype=float)):
-    i = wp.tid()
-
-    wp.expect_eq(a[i], expected[i])
-
-    a[i] = 2.0 * a[i]
-
-    wp.atomic_add(a, i, 1.0)
-
-    wp.expect_eq(a[i], 2.0 * expected[i] + 1.0)
-
-
-def test_indexedarray_1d(test, device):
-    values = np.arange(10, dtype=np.float32)
-    arr = wp.array(data=values, device=device)
-
-    indices = wp.array([1, 3, 5, 7, 9], dtype=int, device=device)
-
-    iarr = wp.indexedarray1d(arr, [indices])
-
-    test.assertEqual(iarr.dtype, arr.dtype)
-    test.assertEqual(iarr.ndim, 1)
-    test.assertEqual(iarr.shape, (5,))
-    test.assertEqual(iarr.size, 5)
-
-    expected_arr = wp.array(data=[1, 3, 5, 7, 9], dtype=float, device=device)
-
-    wp.launch(kernel_1d, dim=iarr.size, inputs=[iarr, expected_arr], device=device)
-
-
-@wp.kernel
-def kernel_2d(a: wp.indexedarray2d(dtype=float), expected: wp.array2d(dtype=float)):
-    i, j = wp.tid()
-
-    # check expected values
-    wp.expect_eq(a[i, j], expected[i, j])
-
-    # test wp.view()
-    wp.expect_eq(a[i][j], a[i, j])
-
-    a[i, j] = 2.0 * a[i, j]
-
-    wp.atomic_add(a, i, j, 1.0)
-
-    wp.expect_eq(a[i, j], 2.0 * expected[i, j] + 1.0)
-
-
-def test_indexedarray_2d(test, device):
-    values = np.arange(100, dtype=np.float32).reshape((10, 10))
-    arr = wp.array(data=values, device=device)
-
-    indices0 = wp.array([1, 3], dtype=int, device=device)
-    indices1 = wp.array([2, 4, 8], dtype=int, device=device)
-
-    iarr = wp.indexedarray2d(arr, [indices0, indices1])
-
-    test.assertEqual(iarr.dtype, arr.dtype)
-    test.assertEqual(iarr.ndim, 2)
-    test.assertEqual(iarr.shape, (2, 3))
-    test.assertEqual(iarr.size, 6)
-
-    expected_values = [[12, 14, 18], [32, 34, 38]]
-    expected_arr = wp.array(data=expected_values, dtype=float, device=device)
-
-    wp.launch(kernel_2d, dim=iarr.shape, inputs=[iarr, expected_arr], device=device)
-
-
-@wp.kernel
-def kernel_3d(a: wp.indexedarray3d(dtype=float), expected: wp.array3d(dtype=float)):
-    i, j, k = wp.tid()
-
-    # check expected values
-    wp.expect_eq(a[i, j, k], expected[i, j, k])
-
-    # test wp.view()
-    wp.expect_eq(a[i][j][k], a[i, j, k])
-    wp.expect_eq(a[i, j][k], a[i, j, k])
-    wp.expect_eq(a[i][j, k], a[i, j, k])
-
-    a[i, j, k] = 2.0 * a[i, j, k]
-
-    wp.atomic_add(a, i, j, k, 1.0)
-
-    wp.expect_eq(a[i, j, k], 2.0 * expected[i, j, k] + 1.0)
-
-
-def test_indexedarray_3d(test, device):
-    values = np.arange(1000, dtype=np.float32).reshape((10, 10, 10))
-    arr = wp.array(data=values, device=device)
-
-    indices0 = wp.array([1, 3], dtype=int, device=device)
-    indices1 = wp.array([2, 4, 8], dtype=int, device=device)
-    indices2 = wp.array([0, 5], dtype=int, device=device)
-
-    iarr = wp.indexedarray3d(arr, [indices0, indices1, indices2])
-
-    test.assertEqual(iarr.dtype, arr.dtype)
-    test.assertEqual(iarr.ndim, 3)
-    test.assertEqual(iarr.shape, (2, 3, 2))
-    test.assertEqual(iarr.size, 12)
-
-    expected_values = [
-        [[120, 125], [140, 145], [180, 185]],
-        [[320, 325], [340, 345], [380, 385]],
-    ]
-    expected_arr = wp.array(data=expected_values, dtype=float, device=device)
-
-    wp.launch(kernel_3d, dim=iarr.shape, inputs=[iarr, expected_arr], device=device)
-
-
-@wp.kernel
-def kernel_4d(a: wp.indexedarray4d(dtype=float), expected: wp.array4d(dtype=float)):
-    i, j, k, l = wp.tid()
-
-    # check expected values
-    wp.expect_eq(a[i, j, k, l], expected[i, j, k, l])
-
-    # test wp.view()
-    wp.expect_eq(a[i][j][k][l], a[i, j, k, l])
-    wp.expect_eq(a[i][j, k, l], a[i, j, k, l])
-    wp.expect_eq(a[i, j][k, l], a[i, j, k, l])
-    wp.expect_eq(a[i, j, k][l], a[i, j, k, l])
-
-    a[i, j, k, l] = 2.0 * a[i, j, k, l]
-
-    wp.atomic_add(a, i, j, k, l, 1.0)
-
-    wp.expect_eq(a[i, j, k, l], 2.0 * expected[i, j, k, l] + 1.0)
-
-
-def test_indexedarray_4d(test, device):
-    values = np.arange(10000, dtype=np.float32).reshape((10, 10, 10, 10))
-    arr = wp.array(data=values, device=device)
-
-    indices0 = wp.array([1, 3], dtype=int, device=device)
-    indices1 = wp.array([2, 4, 8], dtype=int, device=device)
-    indices2 = wp.array([0, 5], dtype=int, device=device)
-    indices3 = wp.array([6, 7, 9], dtype=int, device=device)
-
-    iarr = wp.indexedarray4d(arr, [indices0, indices1, indices2, indices3])
-
-    test.assertEqual(iarr.dtype, arr.dtype)
-    test.assertEqual(iarr.ndim, 4)
-    test.assertEqual(iarr.shape, (2, 3, 2, 3))
-    test.assertEqual(iarr.size, 36)
-
-    expected_values = [
-        [
-            [[1206, 1207, 1209], [1256, 1257, 1259]],
-            [[1406, 1407, 1409], [1456, 1457, 1459]],
-            [[1806, 1807, 1809], [1856, 1857, 1859]],
-        ],
-        [
-            [[3206, 3207, 3209], [3256, 3257, 3259]],
-            [[3406, 3407, 3409], [3456, 3457, 3459]],
-            [[3806, 3807, 3809], [3856, 3857, 3859]],
-        ],
-    ]
-    expected_arr = wp.array(data=expected_values, dtype=float, device=device)
-
-    wp.launch(kernel_4d, dim=iarr.shape, inputs=[iarr, expected_arr], device=device)
-
-
-def test_indexedarray_mixed(test, device):
-    # [[[ 0,  1,  2,  3],
-    #   [ 4,  5,  6,  7],
-    #   [ 8,  9, 10, 11],
-    #   [12, 13, 14, 15]],
-    #  [[16, 17, 18, 19],
-    #   [20, 21, 22, 23],
-    #   [24, 25, 26, 27],
-    #   [28, 29, 30, 31]],
-    #  [[32, 33, 34, 35],
-    #   [36, 37, 38, 39],
-    #   [40, 41, 42, 43],
-    #   [44, 45, 46, 47],
-    #  [[48, 49, 50, 51],
-    #   [52, 53, 54, 55],
-    #   [56, 57, 58, 59],
-    #   [60, 61, 62, 63]]]]
-    values = np.arange(64, dtype=np.float32).reshape((4, 4, 4))
-
-    indices = wp.array([0, 3], dtype=int, device=device)
-
-    # -----
-
-    arr = wp.array(data=values, device=device)
-    iarr = wp.indexedarray(arr, [indices, None, None])
-    test.assertEqual(iarr.dtype, arr.dtype)
-    test.assertEqual(iarr.ndim, 3)
-    test.assertEqual(iarr.shape, (2, 4, 4))
-    test.assertEqual(iarr.size, 32)
-
-    expected_values = [
-        [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11], [12, 13, 14, 15]],
-        [[48, 49, 50, 51], [52, 53, 54, 55], [56, 57, 58, 59], [60, 61, 62, 63]],
-    ]
-    expected_arr = wp.array(data=expected_values, dtype=float, device=device)
-
-    wp.launch(kernel_3d, dim=iarr.shape, inputs=[iarr, expected_arr], device=device)
-
-    # -----
-
-    arr = wp.array(data=values, device=device)
-    iarr = wp.indexedarray(arr, [indices, indices, None])
-    test.assertEqual(iarr.dtype, arr.dtype)
-    test.assertEqual(iarr.ndim, 3)
-    test.assertEqual(iarr.shape, (2, 2, 4))
-    test.assertEqual(iarr.size, 16)
-
-    expected_values = [[[0, 1, 2, 3], [12, 13, 14, 15]], [[48, 49, 50, 51], [60, 61, 62, 63]]]
-    expected_arr = wp.array(data=expected_values, dtype=float, device=device)
-
-    wp.launch(kernel_3d, dim=iarr.shape, inputs=[iarr, expected_arr], device=device)
-
-    # -----
-
-    arr = wp.array(data=values, device=device)
-    iarr = wp.indexedarray(arr, [indices, None, indices])
-    test.assertEqual(iarr.dtype, arr.dtype)
-    test.assertEqual(iarr.ndim, 3)
-    test.assertEqual(iarr.shape, (2, 4, 2))
-    test.assertEqual(iarr.size, 16)
-
-    expected_values = [[[0, 3], [4, 7], [8, 11], [12, 15]], [[48, 51], [52, 55], [56, 59], [60, 63]]]
-    expected_arr = wp.array(data=expected_values, dtype=float, device=device)
-
-    wp.launch(kernel_3d, dim=iarr.shape, inputs=[iarr, expected_arr], device=device)
-
-    # -----
-
-    arr = wp.array(data=values, device=device)
-    iarr = wp.indexedarray(arr, [None, indices, indices])
-    test.assertEqual(iarr.dtype, arr.dtype)
-    test.assertEqual(iarr.ndim, 3)
-    test.assertEqual(iarr.shape, (4, 2, 2))
-    test.assertEqual(iarr.size, 16)
-
-    expected_values = [[[0, 3], [12, 15]], [[16, 19], [28, 31]], [[32, 35], [44, 47]], [[48, 51], [60, 63]]]
-    expected_arr = wp.array(data=expected_values, dtype=float, device=device)
-
-    wp.launch(kernel_3d, dim=iarr.shape, inputs=[iarr, expected_arr], device=device)
-
-
-vec2i = wp.types.vector(length=2, dtype=wp.int32)
-vec3i = wp.types.vector(length=3, dtype=wp.int32)
-vec4i = wp.types.vector(length=4, dtype=wp.int32)
-
-
-@wp.kernel
-def shape_kernel_1d(arr: wp.indexedarray1d(dtype=float), expected: int):
-    wp.expect_eq(arr.shape[0], expected)
-
-
-@wp.kernel
-def shape_kernel_2d(arr: wp.indexedarray2d(dtype=float), expected: vec2i):
-    wp.expect_eq(arr.shape[0], expected[0])
-    wp.expect_eq(arr.shape[1], expected[1])
-
-    # 1d slice
-    view = arr[0]
-    wp.expect_eq(view.shape[0], expected[1])
-
-
-@wp.kernel
-def shape_kernel_3d(arr: wp.indexedarray3d(dtype=float), expected: vec3i):
-    wp.expect_eq(arr.shape[0], expected[0])
-    wp.expect_eq(arr.shape[1], expected[1])
-    wp.expect_eq(arr.shape[2], expected[2])
-
-    # 2d slice
-    view2 = arr[0]
-    wp.expect_eq(view2.shape[0], expected[1])
-    wp.expect_eq(view2.shape[1], expected[2])
-
-    # 1d slice
-    view1 = arr[0, 0]
-    wp.expect_eq(view1.shape[0], expected[2])
-
-
-@wp.kernel
-def shape_kernel_4d(arr: wp.indexedarray4d(dtype=float), expected: vec4i):
-    wp.expect_eq(arr.shape[0], expected[0])
-    wp.expect_eq(arr.shape[1], expected[1])
-    wp.expect_eq(arr.shape[2], expected[2])
-    wp.expect_eq(arr.shape[3], expected[3])
-
-    # 3d slice
-    view3 = arr[0]
-    wp.expect_eq(view3.shape[0], expected[1])
-    wp.expect_eq(view3.shape[1], expected[2])
-    wp.expect_eq(view3.shape[2], expected[3])
-
-    # 2d slice
-    view2 = arr[0, 0]
-    wp.expect_eq(view2.shape[0], expected[2])
-    wp.expect_eq(view2.shape[1], expected[3])
-
-    # 1d slice
-    view1 = arr[0, 0, 0]
-    wp.expect_eq(view1.shape[0], expected[3])
-
-
-def test_indexedarray_shape(test, device):
-    with wp.ScopedDevice(device):
-        data1 = wp.zeros(10, dtype=float)
-        data2 = wp.zeros((10, 20), dtype=float)
-        data3 = wp.zeros((10, 20, 30), dtype=float)
-        data4 = wp.zeros((10, 20, 30, 40), dtype=float)
-
-        indices1 = wp.array(data=[2, 7], dtype=int)
-        indices2 = wp.array(data=[2, 7, 12, 17], dtype=int)
-        indices3 = wp.array(data=[2, 7, 12, 17, 22, 27], dtype=int)
-        indices4 = wp.array(data=[2, 7, 12, 17, 22, 27, 32, 37], dtype=int)
-
-        ia1 = wp.indexedarray(data1, [indices1])
-        wp.launch(shape_kernel_1d, dim=1, inputs=[ia1, 2])
-
-        ia2_1 = wp.indexedarray(data2, [indices1, None])
-        ia2_2 = wp.indexedarray(data2, [None, indices2])
-        ia2_3 = wp.indexedarray(data2, [indices1, indices2])
-        wp.launch(shape_kernel_2d, dim=1, inputs=[ia2_1, vec2i(2, 20)])
-        wp.launch(shape_kernel_2d, dim=1, inputs=[ia2_2, vec2i(10, 4)])
-        wp.launch(shape_kernel_2d, dim=1, inputs=[ia2_3, vec2i(2, 4)])
-
-        ia3_1 = wp.indexedarray(data3, [indices1, None, None])
-        ia3_2 = wp.indexedarray(data3, [None, indices2, None])
-        ia3_3 = wp.indexedarray(data3, [None, None, indices3])
-        ia3_4 = wp.indexedarray(data3, [indices1, indices2, None])
-        ia3_5 = wp.indexedarray(data3, [indices1, None, indices3])
-        ia3_6 = wp.indexedarray(data3, [None, indices2, indices3])
-        ia3_7 = wp.indexedarray(data3, [indices1, indices2, indices3])
-        wp.launch(shape_kernel_3d, dim=1, inputs=[ia3_1, vec3i(2, 20, 30)])
-        wp.launch(shape_kernel_3d, dim=1, inputs=[ia3_2, vec3i(10, 4, 30)])
-        wp.launch(shape_kernel_3d, dim=1, inputs=[ia3_3, vec3i(10, 20, 6)])
-        wp.launch(shape_kernel_3d, dim=1, inputs=[ia3_4, vec3i(2, 4, 30)])
-        wp.launch(shape_kernel_3d, dim=1, inputs=[ia3_5, vec3i(2, 20, 6)])
-        wp.launch(shape_kernel_3d, dim=1, inputs=[ia3_6, vec3i(10, 4, 6)])
-        wp.launch(shape_kernel_3d, dim=1, inputs=[ia3_7, vec3i(2, 4, 6)])
-
-        ia4_1 = wp.indexedarray(data4, [indices1, None, None, None])
-        ia4_2 = wp.indexedarray(data4, [indices1, None, None, indices4])
-        ia4_3 = wp.indexedarray(data4, [None, indices2, indices3, None])
-        ia4_4 = wp.indexedarray(data4, [indices1, indices2, indices3, indices4])
-        wp.launch(shape_kernel_4d, dim=1, inputs=[ia4_1, vec4i(2, 20, 30, 40)])
-        wp.launch(shape_kernel_4d, dim=1, inputs=[ia4_2, vec4i(2, 20, 30, 8)])
-        wp.launch(shape_kernel_4d, dim=1, inputs=[ia4_3, vec4i(10, 4, 6, 40)])
-        wp.launch(shape_kernel_4d, dim=1, inputs=[ia4_4, vec4i(2, 4, 6, 8)])
-
-        wp.synchronize_device(device)
-
-
-def test_indexedarray_getitem(test, device):
-    with wp.ScopedDevice(device):
-        data = wp.array(data=np.arange(1000, dtype=np.int32).reshape((10, 10, 10)))
-
-        I = wp.array(data=[0, 1, 2], dtype=int)
-
-        # use constructor
-        a1 = wp.indexedarray(data, [None, None, I])
-        a2 = wp.indexedarray(data, [None, I])
-        a3 = wp.indexedarray(data, [None, I, I])
-        a4 = wp.indexedarray(data, [I])
-        a5 = wp.indexedarray(data, [I, None, I])
-        a6 = wp.indexedarray(data, [I, I])
-        a7 = wp.indexedarray(data, [I, I, I])
-
-        # use array.__getitem__()
-        b1 = data[:, :, I]
-        b2 = data[:, I]
-        b3 = data[:, I, I]
-        b4 = data[I]
-        b5 = data[I, :, I]
-        b6 = data[I, I]
-        b7 = data[I, I, I]
-
-        test.assertEqual(type(a1), type(b1))
-        test.assertEqual(type(a2), type(b2))
-        test.assertEqual(type(a3), type(b3))
-        test.assertEqual(type(a4), type(b4))
-        test.assertEqual(type(a5), type(b5))
-        test.assertEqual(type(a6), type(b6))
-        test.assertEqual(type(a7), type(b7))
-
-        assert_np_equal(a1.numpy(), b1.numpy())
-        assert_np_equal(a2.numpy(), b2.numpy())
-        assert_np_equal(a3.numpy(), b3.numpy())
-        assert_np_equal(a4.numpy(), b4.numpy())
-        assert_np_equal(a5.numpy(), b5.numpy())
-        assert_np_equal(a6.numpy(), b6.numpy())
-        assert_np_equal(a7.numpy(), b7.numpy())
-
-
-def test_indexedarray_slicing(test, device):
-    with wp.ScopedDevice(device):
-        data = wp.array(data=np.arange(1000, dtype=np.int32).reshape((10, 10, 10)))
-
-        # test equivalence of slicing and indexing the same range
-        s = slice(0, 3)
-        I = wp.array(data=[0, 1, 2], dtype=int)
-
-        a0 = data[s, s, s]
-        test.assertEqual(type(a0), wp.array)
-        a1 = data[s, s, I]
-        test.assertEqual(type(a1), wp.indexedarray)
-        a2 = data[s, I, s]
-        test.assertEqual(type(a2), wp.indexedarray)
-        a3 = data[s, I, I]
-        test.assertEqual(type(a3), wp.indexedarray)
-        a4 = data[I, s, s]
-        test.assertEqual(type(a4), wp.indexedarray)
-        a5 = data[I, s, I]
-        test.assertEqual(type(a5), wp.indexedarray)
-        a6 = data[I, I, s]
-        test.assertEqual(type(a6), wp.indexedarray)
-        a7 = data[I, I, I]
-        test.assertEqual(type(a7), wp.indexedarray)
-
-        expected = a0.numpy()
-
-        assert_np_equal(a1.numpy(), expected)
-        assert_np_equal(a2.numpy(), expected)
-        assert_np_equal(a3.numpy(), expected)
-        assert_np_equal(a4.numpy(), expected)
-        assert_np_equal(a5.numpy(), expected)
-        assert_np_equal(a6.numpy(), expected)
-        assert_np_equal(a7.numpy(), expected)
-
-
-# generic increment kernels that work with any array (regular or indexed)
-@wp.kernel
-def inc_1d(a: Any):
-    i = wp.tid()
-    a[i] = a[i] + 1
-
-
-@wp.kernel
-def inc_2d(a: Any):
-    i, j = wp.tid()
-    a[i, j] = a[i, j] + 1
-
-
-@wp.kernel
-def inc_3d(a: Any):
-    i, j, k = wp.tid()
-    a[i, j, k] = a[i, j, k] + 1
-
-
-@wp.kernel
-def inc_4d(a: Any):
-    i, j, k, l = wp.tid()
-    a[i, j, k, l] = a[i, j, k, l] + 1
-
-
-# optional overloads to avoid module reloading
-wp.overload(inc_1d, [wp.array1d(dtype=int)])
-wp.overload(inc_2d, [wp.array2d(dtype=int)])
-wp.overload(inc_3d, [wp.array3d(dtype=int)])
-wp.overload(inc_4d, [wp.array4d(dtype=int)])
-
-wp.overload(inc_1d, [wp.indexedarray1d(dtype=int)])
-wp.overload(inc_2d, [wp.indexedarray2d(dtype=int)])
-wp.overload(inc_3d, [wp.indexedarray3d(dtype=int)])
-wp.overload(inc_4d, [wp.indexedarray4d(dtype=int)])
-
-
-def test_indexedarray_generics(test, device):
-    with wp.ScopedDevice(device):
-        data1 = wp.zeros((5,), dtype=int)
-        data2 = wp.zeros((5, 5), dtype=int)
-        data3 = wp.zeros((5, 5, 5), dtype=int)
-        data4 = wp.zeros((5, 5, 5, 5), dtype=int)
-
-        indices = wp.array(data=[0, 4], dtype=int)
-
-        ia1 = wp.indexedarray(data1, [indices])
-        ia2 = wp.indexedarray(data2, [indices, indices])
-        ia3 = wp.indexedarray(data3, [indices, indices, indices])
-        ia4 = wp.indexedarray(data4, [indices, indices, indices, indices])
-
-        wp.launch(inc_1d, dim=data1.shape, inputs=[data1])
-        wp.launch(inc_2d, dim=data2.shape, inputs=[data2])
-        wp.launch(inc_3d, dim=data3.shape, inputs=[data3])
-        wp.launch(inc_4d, dim=data4.shape, inputs=[data4])
-
-        wp.launch(inc_1d, dim=ia1.shape, inputs=[ia1])
-        wp.launch(inc_2d, dim=ia2.shape, inputs=[ia2])
-        wp.launch(inc_3d, dim=ia3.shape, inputs=[ia3])
-        wp.launch(inc_4d, dim=ia4.shape, inputs=[ia4])
-
-        expected1 = np.ones(5, dtype=np.int32)
-        expected1[0] = 2
-        expected1[4] = 2
-
-        expected2 = np.ones((5, 5), dtype=np.int32)
-        expected2[0, 0] = 2
-        expected2[0, 4] = 2
-        expected2[4, 0] = 2
-        expected2[4, 4] = 2
-
-        expected3 = np.ones((5, 5, 5), dtype=np.int32)
-        expected3[0, 0, 0] = 2
-        expected3[0, 0, 4] = 2
-        expected3[0, 4, 0] = 2
-        expected3[0, 4, 4] = 2
-        expected3[4, 0, 0] = 2
-        expected3[4, 0, 4] = 2
-        expected3[4, 4, 0] = 2
-        expected3[4, 4, 4] = 2
-
-        expected4 = np.ones((5, 5, 5, 5), dtype=np.int32)
-        expected4[0, 0, 0, 0] = 2
-        expected4[0, 0, 0, 4] = 2
-        expected4[0, 0, 4, 0] = 2
-        expected4[0, 0, 4, 4] = 2
-        expected4[0, 4, 0, 0] = 2
-        expected4[0, 4, 0, 4] = 2
-        expected4[0, 4, 4, 0] = 2
-        expected4[0, 4, 4, 4] = 2
-        expected4[4, 0, 0, 0] = 2
-        expected4[4, 0, 0, 4] = 2
-        expected4[4, 0, 4, 0] = 2
-        expected4[4, 0, 4, 4] = 2
-        expected4[4, 4, 0, 0] = 2
-        expected4[4, 4, 0, 4] = 2
-        expected4[4, 4, 4, 0] = 2
-        expected4[4, 4, 4, 4] = 2
-
-        assert_np_equal(data1.numpy(), expected1)
-        assert_np_equal(data2.numpy(), expected2)
-        assert_np_equal(data3.numpy(), expected3)
-        assert_np_equal(data4.numpy(), expected4)
-
-        assert_np_equal(ia1.numpy(), np.full((2,), 2, dtype=np.int32))
-        assert_np_equal(ia2.numpy(), np.full((2, 2), 2, dtype=np.int32))
-        assert_np_equal(ia3.numpy(), np.full((2, 2, 2), 2, dtype=np.int32))
-        assert_np_equal(ia4.numpy(), np.full((2, 2, 2, 2), 2, dtype=np.int32))
-
-
-def test_indexedarray_empty(test, device):
-    # Test whether common operations work with empty (zero-sized) indexed arrays
-    # without throwing exceptions.
-
-    def test_empty_ops(ndim, nrows, ncols, wptype, nptype):
-        data_shape = (1,) * ndim
-        dtype_shape = ()
-
-        if wptype in wp.types.scalar_types:
-            # scalar, vector, or matrix
-            if ncols > 0:
-                if nrows > 0:
-                    wptype = wp.types.matrix((nrows, ncols), wptype)
-                else:
-                    wptype = wp.types.vector(ncols, wptype)
-                dtype_shape = wptype._shape_
-            fill_value = wptype(42)
-        else:
-            # struct
-            fill_value = wptype()
-
-        # create a data array
-        data = wp.empty(data_shape, dtype=wptype, device=device, requires_grad=True)
-
-        # create a zero-sized array of indices
-        indices = wp.empty(0, dtype=int, device=device)
-
-        a = data[indices]
-
-        # we expect dim to be zero for the empty indexed array, unchanged otherwise
-        expected_shape = (0, *data_shape[1:])
-
-        test.assertEqual(a.size, 0)
-        test.assertEqual(a.shape, expected_shape)
-
-        # all of these methods should succeed with zero-sized arrays
-        a.zero_()
-        a.fill_(fill_value)
-        b = a.contiguous()
-
-        b = wp.empty_like(a)
-        b = wp.zeros_like(a)
-        b = wp.full_like(a, fill_value)
-        b = wp.clone(a)
-
-        wp.copy(a, b)
-        a.assign(b)
-
-        na = a.numpy()
-        test.assertEqual(na.size, 0)
-        test.assertEqual(na.shape, (*expected_shape, *dtype_shape))
-        test.assertEqual(na.dtype, nptype)
-
-        test.assertEqual(a.list(), [])
-
-    for ndim in range(1, 5):
-        # test with scalars, vectors, and matrices
-        for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
-            # scalars
-            test_empty_ops(ndim, 0, 0, wptype, nptype)
-
-            for ncols in [2, 3, 4, 5]:
-                # vectors
-                test_empty_ops(ndim, 0, ncols, wptype, nptype)
-                # square matrices
-                test_empty_ops(ndim, ncols, ncols, wptype, nptype)
-
-            # non-square matrices
-            test_empty_ops(ndim, 2, 3, wptype, nptype)
-            test_empty_ops(ndim, 3, 2, wptype, nptype)
-            test_empty_ops(ndim, 3, 4, wptype, nptype)
-            test_empty_ops(ndim, 4, 3, wptype, nptype)
-
-        # test with structs
-        test_empty_ops(ndim, 0, 0, FillStruct, FillStruct.numpy_dtype())
-
-
-def test_indexedarray_fill_scalar(test, device):
-    dim_x = 4
-
-    for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
-        data1 = wp.zeros(dim_x, dtype=wptype, device=device)
-        data2 = wp.zeros((dim_x, dim_x), dtype=wptype, device=device)
-        data3 = wp.zeros((dim_x, dim_x, dim_x), dtype=wptype, device=device)
-        data4 = wp.zeros((dim_x, dim_x, dim_x, dim_x), dtype=wptype, device=device)
-
-        indices = wp.array(np.arange(0, dim_x, 2, dtype=np.int32), device=device)
-
-        a1 = data1[indices]
-        a2 = data2[indices]
-        a3 = data3[indices]
-        a4 = data4[indices]
-
-        assert_np_equal(a1.numpy(), np.zeros(a1.shape, dtype=nptype))
-        assert_np_equal(a2.numpy(), np.zeros(a2.shape, dtype=nptype))
-        assert_np_equal(a3.numpy(), np.zeros(a3.shape, dtype=nptype))
-        assert_np_equal(a4.numpy(), np.zeros(a4.shape, dtype=nptype))
-
-        # fill with int value
-        fill_value = 42
-
-        a1.fill_(fill_value)
-        a2.fill_(fill_value)
-        a3.fill_(fill_value)
-        a4.fill_(fill_value)
-
-        assert_np_equal(a1.numpy(), np.full(a1.shape, fill_value, dtype=nptype))
-        assert_np_equal(a2.numpy(), np.full(a2.shape, fill_value, dtype=nptype))
-        assert_np_equal(a3.numpy(), np.full(a3.shape, fill_value, dtype=nptype))
-        assert_np_equal(a4.numpy(), np.full(a4.shape, fill_value, dtype=nptype))
-
-        a1.zero_()
-        a2.zero_()
-        a3.zero_()
-        a4.zero_()
-
-        assert_np_equal(a1.numpy(), np.zeros(a1.shape, dtype=nptype))
-        assert_np_equal(a2.numpy(), np.zeros(a2.shape, dtype=nptype))
-        assert_np_equal(a3.numpy(), np.zeros(a3.shape, dtype=nptype))
-        assert_np_equal(a4.numpy(), np.zeros(a4.shape, dtype=nptype))
-
-        if wptype in wp.types.float_types:
-            # fill with float value
-            fill_value = 13.37
-
-            a1.fill_(fill_value)
-            a2.fill_(fill_value)
-            a3.fill_(fill_value)
-            a4.fill_(fill_value)
-
-            assert_np_equal(a1.numpy(), np.full(a1.shape, fill_value, dtype=nptype))
-            assert_np_equal(a2.numpy(), np.full(a2.shape, fill_value, dtype=nptype))
-            assert_np_equal(a3.numpy(), np.full(a3.shape, fill_value, dtype=nptype))
-            assert_np_equal(a4.numpy(), np.full(a4.shape, fill_value, dtype=nptype))
-
-        # fill with Warp scalar value
-        fill_value = wptype(17)
-
-        a1.fill_(fill_value)
-        a2.fill_(fill_value)
-        a3.fill_(fill_value)
-        a4.fill_(fill_value)
-
-        assert_np_equal(a1.numpy(), np.full(a1.shape, fill_value.value, dtype=nptype))
-        assert_np_equal(a2.numpy(), np.full(a2.shape, fill_value.value, dtype=nptype))
-        assert_np_equal(a3.numpy(), np.full(a3.shape, fill_value.value, dtype=nptype))
-        assert_np_equal(a4.numpy(), np.full(a4.shape, fill_value.value, dtype=nptype))
-
-
-def test_indexedarray_fill_vector(test, device):
-    # test filling a vector array with scalar or vector values (vec_type, list, or numpy array)
-
-    dim_x = 4
-
-    for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
-        # vector types
-        vector_types = [
-            wp.types.vector(2, wptype),
-            wp.types.vector(3, wptype),
-            wp.types.vector(4, wptype),
-            wp.types.vector(5, wptype),
-        ]
-
-        for vec_type in vector_types:
-            vec_len = vec_type._length_
-
-            data1 = wp.zeros(dim_x, dtype=vec_type, device=device)
-            data2 = wp.zeros((dim_x, dim_x), dtype=vec_type, device=device)
-            data3 = wp.zeros((dim_x, dim_x, dim_x), dtype=vec_type, device=device)
-            data4 = wp.zeros((dim_x, dim_x, dim_x, dim_x), dtype=vec_type, device=device)
-
-            indices = wp.array(np.arange(0, dim_x, 2, dtype=np.int32), device=device)
-
-            a1 = data1[indices]
-            a2 = data2[indices]
-            a3 = data3[indices]
-            a4 = data4[indices]
-
-            assert_np_equal(a1.numpy(), np.zeros((*a1.shape, vec_len), dtype=nptype))
-            assert_np_equal(a2.numpy(), np.zeros((*a2.shape, vec_len), dtype=nptype))
-            assert_np_equal(a3.numpy(), np.zeros((*a3.shape, vec_len), dtype=nptype))
-            assert_np_equal(a4.numpy(), np.zeros((*a4.shape, vec_len), dtype=nptype))
-
-            # fill with int scalar
-            fill_value = 42
-
-            a1.fill_(fill_value)
-            a2.fill_(fill_value)
-            a3.fill_(fill_value)
-            a4.fill_(fill_value)
-
-            assert_np_equal(a1.numpy(), np.full((*a1.shape, vec_len), fill_value, dtype=nptype))
-            assert_np_equal(a2.numpy(), np.full((*a2.shape, vec_len), fill_value, dtype=nptype))
-            assert_np_equal(a3.numpy(), np.full((*a3.shape, vec_len), fill_value, dtype=nptype))
-            assert_np_equal(a4.numpy(), np.full((*a4.shape, vec_len), fill_value, dtype=nptype))
-
-            # test zeroing
-            a1.zero_()
-            a2.zero_()
-            a3.zero_()
-            a4.zero_()
-
-            assert_np_equal(a1.numpy(), np.zeros((*a1.shape, vec_len), dtype=nptype))
-            assert_np_equal(a2.numpy(), np.zeros((*a2.shape, vec_len), dtype=nptype))
-            assert_np_equal(a3.numpy(), np.zeros((*a3.shape, vec_len), dtype=nptype))
-            assert_np_equal(a4.numpy(), np.zeros((*a4.shape, vec_len), dtype=nptype))
-
-            # vector values can be passed as a list, numpy array, or Warp vector instance
-            fill_list = [17, 42, 99, 101, 127][:vec_len]
-            fill_arr = np.array(fill_list, dtype=nptype)
-            fill_vec = vec_type(fill_list)
-
-            expected1 = np.tile(fill_arr, a1.size).reshape((*a1.shape, vec_len))
-            expected2 = np.tile(fill_arr, a2.size).reshape((*a2.shape, vec_len))
-            expected3 = np.tile(fill_arr, a3.size).reshape((*a3.shape, vec_len))
-            expected4 = np.tile(fill_arr, a4.size).reshape((*a4.shape, vec_len))
-
-            # fill with list of vector length
-            a1.fill_(fill_list)
-            a2.fill_(fill_list)
-            a3.fill_(fill_list)
-            a4.fill_(fill_list)
-
-            assert_np_equal(a1.numpy(), expected1)
-            assert_np_equal(a2.numpy(), expected2)
-            assert_np_equal(a3.numpy(), expected3)
-            assert_np_equal(a4.numpy(), expected4)
-
-            # clear
-            a1.zero_()
-            a2.zero_()
-            a3.zero_()
-            a4.zero_()
-
-            # fill with numpy array of vector length
-            a1.fill_(fill_arr)
-            a2.fill_(fill_arr)
-            a3.fill_(fill_arr)
-            a4.fill_(fill_arr)
-
-            assert_np_equal(a1.numpy(), expected1)
-            assert_np_equal(a2.numpy(), expected2)
-            assert_np_equal(a3.numpy(), expected3)
-            assert_np_equal(a4.numpy(), expected4)
-
-            # clear
-            a1.zero_()
-            a2.zero_()
-            a3.zero_()
-            a4.zero_()
-
-            # fill with vec instance
-            a1.fill_(fill_vec)
-            a2.fill_(fill_vec)
-            a3.fill_(fill_vec)
-            a4.fill_(fill_vec)
-
-            assert_np_equal(a1.numpy(), expected1)
-            assert_np_equal(a2.numpy(), expected2)
-            assert_np_equal(a3.numpy(), expected3)
-            assert_np_equal(a4.numpy(), expected4)
-
-            if wptype in wp.types.float_types:
-                # fill with float scalar
-                fill_value = 13.37
-
-                a1.fill_(fill_value)
-                a2.fill_(fill_value)
-                a3.fill_(fill_value)
-                a4.fill_(fill_value)
-
-                assert_np_equal(a1.numpy(), np.full((*a1.shape, vec_len), fill_value, dtype=nptype))
-                assert_np_equal(a2.numpy(), np.full((*a2.shape, vec_len), fill_value, dtype=nptype))
-                assert_np_equal(a3.numpy(), np.full((*a3.shape, vec_len), fill_value, dtype=nptype))
-                assert_np_equal(a4.numpy(), np.full((*a4.shape, vec_len), fill_value, dtype=nptype))
-
-                # fill with float list of vector length
-                fill_list = [-2.5, -1.25, 1.25, 2.5, 5.0][:vec_len]
-
-                a1.fill_(fill_list)
-                a2.fill_(fill_list)
-                a3.fill_(fill_list)
-                a4.fill_(fill_list)
-
-                expected1 = np.tile(np.array(fill_list, dtype=nptype), a1.size).reshape((*a1.shape, vec_len))
-                expected2 = np.tile(np.array(fill_list, dtype=nptype), a2.size).reshape((*a2.shape, vec_len))
-                expected3 = np.tile(np.array(fill_list, dtype=nptype), a3.size).reshape((*a3.shape, vec_len))
-                expected4 = np.tile(np.array(fill_list, dtype=nptype), a4.size).reshape((*a4.shape, vec_len))
-
-                assert_np_equal(a1.numpy(), expected1)
-                assert_np_equal(a2.numpy(), expected2)
-                assert_np_equal(a3.numpy(), expected3)
-                assert_np_equal(a4.numpy(), expected4)
-
-
-def test_indexedarray_fill_matrix(test, device):
-    # test filling a matrix array with scalar or matrix values (mat_type, nested list, or 2d numpy array)
-
-    dim_x = 4
-
-    for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
-        # matrix types
-        matrix_types = [
-            # square matrices
-            wp.types.matrix((2, 2), wptype),
-            wp.types.matrix((3, 3), wptype),
-            wp.types.matrix((4, 4), wptype),
-            wp.types.matrix((5, 5), wptype),
-            # non-square matrices
-            wp.types.matrix((2, 3), wptype),
-            wp.types.matrix((3, 2), wptype),
-            wp.types.matrix((3, 4), wptype),
-            wp.types.matrix((4, 3), wptype),
-        ]
-
-        for mat_type in matrix_types:
-            mat_len = mat_type._length_
-            mat_shape = mat_type._shape_
-
-            data1 = wp.zeros(dim_x, dtype=mat_type, device=device)
-            data2 = wp.zeros((dim_x, dim_x), dtype=mat_type, device=device)
-            data3 = wp.zeros((dim_x, dim_x, dim_x), dtype=mat_type, device=device)
-            data4 = wp.zeros((dim_x, dim_x, dim_x, dim_x), dtype=mat_type, device=device)
-
-            indices = wp.array(np.arange(0, dim_x, 2, dtype=np.int32), device=device)
-
-            a1 = data1[indices]
-            a2 = data2[indices]
-            a3 = data3[indices]
-            a4 = data4[indices]
-
-            assert_np_equal(a1.numpy(), np.zeros((*a1.shape, *mat_shape), dtype=nptype))
-            assert_np_equal(a2.numpy(), np.zeros((*a2.shape, *mat_shape), dtype=nptype))
-            assert_np_equal(a3.numpy(), np.zeros((*a3.shape, *mat_shape), dtype=nptype))
-            assert_np_equal(a4.numpy(), np.zeros((*a4.shape, *mat_shape), dtype=nptype))
-
-            # fill with scalar
-            fill_value = 42
-
-            a1.fill_(fill_value)
-            a2.fill_(fill_value)
-            a3.fill_(fill_value)
-            a4.fill_(fill_value)
-
-            assert_np_equal(a1.numpy(), np.full((*a1.shape, *mat_shape), fill_value, dtype=nptype))
-            assert_np_equal(a2.numpy(), np.full((*a2.shape, *mat_shape), fill_value, dtype=nptype))
-            assert_np_equal(a3.numpy(), np.full((*a3.shape, *mat_shape), fill_value, dtype=nptype))
-            assert_np_equal(a4.numpy(), np.full((*a4.shape, *mat_shape), fill_value, dtype=nptype))
-
-            # test zeroing
-            a1.zero_()
-            a2.zero_()
-            a3.zero_()
-            a4.zero_()
-
-            assert_np_equal(a1.numpy(), np.zeros((*a1.shape, *mat_shape), dtype=nptype))
-            assert_np_equal(a2.numpy(), np.zeros((*a2.shape, *mat_shape), dtype=nptype))
-            assert_np_equal(a3.numpy(), np.zeros((*a3.shape, *mat_shape), dtype=nptype))
-            assert_np_equal(a4.numpy(), np.zeros((*a4.shape, *mat_shape), dtype=nptype))
-
-            # matrix values can be passed as a 1d numpy array, 2d numpy array, flat list, nested list, or Warp matrix instance
-            if wptype != wp.bool:
-                fill_arr1 = np.arange(mat_len, dtype=nptype)
-            else:
-                fill_arr1 = np.ones(mat_len, dtype=nptype)
-            fill_arr2 = fill_arr1.reshape(mat_shape)
-            fill_list1 = list(fill_arr1)
-            fill_list2 = [list(row) for row in fill_arr2]
-            fill_mat = mat_type(fill_arr1)
-
-            expected1 = np.tile(fill_arr1, a1.size).reshape((*a1.shape, *mat_shape))
-            expected2 = np.tile(fill_arr1, a2.size).reshape((*a2.shape, *mat_shape))
-            expected3 = np.tile(fill_arr1, a3.size).reshape((*a3.shape, *mat_shape))
-            expected4 = np.tile(fill_arr1, a4.size).reshape((*a4.shape, *mat_shape))
-
-            # fill with 1d numpy array
-            a1.fill_(fill_arr1)
-            a2.fill_(fill_arr1)
-            a3.fill_(fill_arr1)
-            a4.fill_(fill_arr1)
-
-            assert_np_equal(a1.numpy(), expected1)
-            assert_np_equal(a2.numpy(), expected2)
-            assert_np_equal(a3.numpy(), expected3)
-            assert_np_equal(a4.numpy(), expected4)
-
-            # clear
-            a1.zero_()
-            a2.zero_()
-            a3.zero_()
-            a4.zero_()
-
-            # fill with 2d numpy array
-            a1.fill_(fill_arr2)
-            a2.fill_(fill_arr2)
-            a3.fill_(fill_arr2)
-            a4.fill_(fill_arr2)
-
-            assert_np_equal(a1.numpy(), expected1)
-            assert_np_equal(a2.numpy(), expected2)
-            assert_np_equal(a3.numpy(), expected3)
-            assert_np_equal(a4.numpy(), expected4)
-
-            # clear
-            a1.zero_()
-            a2.zero_()
-            a3.zero_()
-            a4.zero_()
-
-            # fill with flat list
-            a1.fill_(fill_list1)
-            a2.fill_(fill_list1)
-            a3.fill_(fill_list1)
-            a4.fill_(fill_list1)
-
-            assert_np_equal(a1.numpy(), expected1)
-            assert_np_equal(a2.numpy(), expected2)
-            assert_np_equal(a3.numpy(), expected3)
-            assert_np_equal(a4.numpy(), expected4)
-
-            # clear
-            a1.zero_()
-            a2.zero_()
-            a3.zero_()
-            a4.zero_()
-
-            # fill with nested list
-            a1.fill_(fill_list2)
-            a2.fill_(fill_list2)
-            a3.fill_(fill_list2)
-            a4.fill_(fill_list2)
-
-            assert_np_equal(a1.numpy(), expected1)
-            assert_np_equal(a2.numpy(), expected2)
-            assert_np_equal(a3.numpy(), expected3)
-            assert_np_equal(a4.numpy(), expected4)
-
-            # clear
-            a1.zero_()
-            a2.zero_()
-            a3.zero_()
-            a4.zero_()
-
-            # fill with mat instance
-            a1.fill_(fill_mat)
-            a2.fill_(fill_mat)
-            a3.fill_(fill_mat)
-            a4.fill_(fill_mat)
-
-            assert_np_equal(a1.numpy(), expected1)
-            assert_np_equal(a2.numpy(), expected2)
-            assert_np_equal(a3.numpy(), expected3)
-            assert_np_equal(a4.numpy(), expected4)
-
-
-def test_indexedarray_fill_struct(test, device):
-    dim_x = 8
-
-    nptype = FillStruct.numpy_dtype()
-
-    data1 = wp.zeros(dim_x, dtype=FillStruct, device=device)
-    data2 = wp.zeros((dim_x, dim_x), dtype=FillStruct, device=device)
-    data3 = wp.zeros((dim_x, dim_x, dim_x), dtype=FillStruct, device=device)
-    data4 = wp.zeros((dim_x, dim_x, dim_x, dim_x), dtype=FillStruct, device=device)
-
-    indices = wp.array(np.arange(0, dim_x, 2, dtype=np.int32), device=device)
-
-    a1 = data1[indices]
-    a2 = data2[indices]
-    a3 = data3[indices]
-    a4 = data4[indices]
-
-    assert_np_equal(a1.numpy(), np.zeros(a1.shape, dtype=nptype))
-    assert_np_equal(a2.numpy(), np.zeros(a2.shape, dtype=nptype))
-    assert_np_equal(a3.numpy(), np.zeros(a3.shape, dtype=nptype))
-    assert_np_equal(a4.numpy(), np.zeros(a4.shape, dtype=nptype))
-
-    s = FillStruct()
-
-    # fill with default struct value (should be all zeros)
-    a1.fill_(s)
-    a2.fill_(s)
-    a3.fill_(s)
-    a4.fill_(s)
-
-    assert_np_equal(a1.numpy(), np.zeros(a1.shape, dtype=nptype))
-    assert_np_equal(a2.numpy(), np.zeros(a2.shape, dtype=nptype))
-    assert_np_equal(a3.numpy(), np.zeros(a3.shape, dtype=nptype))
-    assert_np_equal(a4.numpy(), np.zeros(a4.shape, dtype=nptype))
-
-    # scalars
-    s.i1 = -17
-    s.i2 = 42
-    s.i4 = 99
-    s.i8 = 101
-    s.f2 = -1.25
-    s.f4 = 13.37
-    s.f8 = 0.125
-    # vectors
-    s.v2 = [21, 22]
-    s.v3 = [31, 32, 33]
-    s.v4 = [41, 42, 43, 44]
-    s.v5 = [51, 52, 53, 54, 55]
-    # matrices
-    s.m2 = [[61, 62]] * 2
-    s.m3 = [[71, 72, 73]] * 3
-    s.m4 = [[81, 82, 83, 84]] * 4
-    s.m5 = [[91, 92, 93, 94, 95]] * 5
-    # arrays
-    s.a1 = wp.zeros((2,) * 1, dtype=float, device=device)
-    s.a2 = wp.zeros((2,) * 2, dtype=float, device=device)
-    s.a3 = wp.zeros((2,) * 3, dtype=float, device=device)
-    s.a4 = wp.zeros((2,) * 4, dtype=float, device=device)
-
-    # fill with custom struct value
-    a1.fill_(s)
-    a2.fill_(s)
-    a3.fill_(s)
-    a4.fill_(s)
-
-    ns = s.numpy_value()
-
-    expected1 = np.empty(a1.shape, dtype=nptype)
-    expected2 = np.empty(a2.shape, dtype=nptype)
-    expected3 = np.empty(a3.shape, dtype=nptype)
-    expected4 = np.empty(a4.shape, dtype=nptype)
-
-    expected1.fill(ns)
-    expected2.fill(ns)
-    expected3.fill(ns)
-    expected4.fill(ns)
-
-    assert_np_equal(a1.numpy(), expected1)
-    assert_np_equal(a2.numpy(), expected2)
-    assert_np_equal(a3.numpy(), expected3)
-    assert_np_equal(a4.numpy(), expected4)
-
-    # test clearing
-    a1.zero_()
-    a2.zero_()
-    a3.zero_()
-    a4.zero_()
-
-    assert_np_equal(a1.numpy(), np.zeros(a1.shape, dtype=nptype))
-    assert_np_equal(a2.numpy(), np.zeros(a2.shape, dtype=nptype))
-    assert_np_equal(a3.numpy(), np.zeros(a3.shape, dtype=nptype))
-    assert_np_equal(a4.numpy(), np.zeros(a4.shape, dtype=nptype))
-
-
-devices = get_test_devices()
-
-
-class TestIndexedArray(unittest.TestCase):
-    pass
-
-
-add_function_test(TestIndexedArray, "test_indexedarray_1d", test_indexedarray_1d, devices=devices)
-add_function_test(TestIndexedArray, "test_indexedarray_2d", test_indexedarray_2d, devices=devices)
-add_function_test(TestIndexedArray, "test_indexedarray_3d", test_indexedarray_3d, devices=devices)
-add_function_test(TestIndexedArray, "test_indexedarray_4d", test_indexedarray_4d, devices=devices)
-add_function_test(TestIndexedArray, "test_indexedarray_mixed", test_indexedarray_mixed, devices=devices)
-add_function_test(TestIndexedArray, "test_indexedarray_shape", test_indexedarray_shape, devices=devices)
-add_function_test(TestIndexedArray, "test_indexedarray_getitem", test_indexedarray_getitem, devices=devices)
-add_function_test(TestIndexedArray, "test_indexedarray_slicing", test_indexedarray_slicing, devices=devices)
-add_function_test(TestIndexedArray, "test_indexedarray_generics", test_indexedarray_generics, devices=devices)
-add_function_test(TestIndexedArray, "test_indexedarray_empty", test_indexedarray_empty, devices=devices)
-add_function_test(TestIndexedArray, "test_indexedarray_fill_scalar", test_indexedarray_fill_scalar, devices=devices)
-add_function_test(TestIndexedArray, "test_indexedarray_fill_vector", test_indexedarray_fill_vector, devices=devices)
-add_function_test(TestIndexedArray, "test_indexedarray_fill_matrix", test_indexedarray_fill_matrix, devices=devices)
-add_function_test(TestIndexedArray, "test_indexedarray_fill_struct", test_indexedarray_fill_struct, devices=devices)
-
-
-if __name__ == "__main__":
-    wp.build.clear_kernel_cache()
-    unittest.main(verbosity=2)
+# Copyright (c) 2023 NVIDIA CORPORATION.  All rights reserved.
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import unittest
+from typing import Any
+
+import numpy as np
+
+import warp as wp
+from warp.tests.test_array import FillStruct
+from warp.tests.unittest_utils import *
+
+wp.init()
+
+
+@wp.kernel
+def kernel_1d(a: wp.indexedarray(dtype=float), expected: wp.array(dtype=float)):
+    i = wp.tid()
+
+    wp.expect_eq(a[i], expected[i])
+
+    a[i] = 2.0 * a[i]
+
+    wp.atomic_add(a, i, 1.0)
+
+    wp.expect_eq(a[i], 2.0 * expected[i] + 1.0)
+
+
+def test_indexedarray_1d(test, device):
+    values = np.arange(10, dtype=np.float32)
+    arr = wp.array(data=values, device=device)
+
+    indices = wp.array([1, 3, 5, 7, 9], dtype=int, device=device)
+
+    iarr = wp.indexedarray1d(arr, [indices])
+
+    test.assertEqual(iarr.dtype, arr.dtype)
+    test.assertEqual(iarr.ndim, 1)
+    test.assertEqual(iarr.shape, (5,))
+    test.assertEqual(iarr.size, 5)
+
+    expected_arr = wp.array(data=[1, 3, 5, 7, 9], dtype=float, device=device)
+
+    wp.launch(kernel_1d, dim=iarr.size, inputs=[iarr, expected_arr], device=device)
+
+
+@wp.kernel
+def kernel_2d(a: wp.indexedarray2d(dtype=float), expected: wp.array2d(dtype=float)):
+    i, j = wp.tid()
+
+    # check expected values
+    wp.expect_eq(a[i, j], expected[i, j])
+
+    # test wp.view()
+    wp.expect_eq(a[i][j], a[i, j])
+
+    a[i, j] = 2.0 * a[i, j]
+
+    wp.atomic_add(a, i, j, 1.0)
+
+    wp.expect_eq(a[i, j], 2.0 * expected[i, j] + 1.0)
+
+
+def test_indexedarray_2d(test, device):
+    values = np.arange(100, dtype=np.float32).reshape((10, 10))
+    arr = wp.array(data=values, device=device)
+
+    indices0 = wp.array([1, 3], dtype=int, device=device)
+    indices1 = wp.array([2, 4, 8], dtype=int, device=device)
+
+    iarr = wp.indexedarray2d(arr, [indices0, indices1])
+
+    test.assertEqual(iarr.dtype, arr.dtype)
+    test.assertEqual(iarr.ndim, 2)
+    test.assertEqual(iarr.shape, (2, 3))
+    test.assertEqual(iarr.size, 6)
+
+    expected_values = [[12, 14, 18], [32, 34, 38]]
+    expected_arr = wp.array(data=expected_values, dtype=float, device=device)
+
+    wp.launch(kernel_2d, dim=iarr.shape, inputs=[iarr, expected_arr], device=device)
+
+
+@wp.kernel
+def kernel_3d(a: wp.indexedarray3d(dtype=float), expected: wp.array3d(dtype=float)):
+    i, j, k = wp.tid()
+
+    # check expected values
+    wp.expect_eq(a[i, j, k], expected[i, j, k])
+
+    # test wp.view()
+    wp.expect_eq(a[i][j][k], a[i, j, k])
+    wp.expect_eq(a[i, j][k], a[i, j, k])
+    wp.expect_eq(a[i][j, k], a[i, j, k])
+
+    a[i, j, k] = 2.0 * a[i, j, k]
+
+    wp.atomic_add(a, i, j, k, 1.0)
+
+    wp.expect_eq(a[i, j, k], 2.0 * expected[i, j, k] + 1.0)
+
+
+def test_indexedarray_3d(test, device):
+    values = np.arange(1000, dtype=np.float32).reshape((10, 10, 10))
+    arr = wp.array(data=values, device=device)
+
+    indices0 = wp.array([1, 3], dtype=int, device=device)
+    indices1 = wp.array([2, 4, 8], dtype=int, device=device)
+    indices2 = wp.array([0, 5], dtype=int, device=device)
+
+    iarr = wp.indexedarray3d(arr, [indices0, indices1, indices2])
+
+    test.assertEqual(iarr.dtype, arr.dtype)
+    test.assertEqual(iarr.ndim, 3)
+    test.assertEqual(iarr.shape, (2, 3, 2))
+    test.assertEqual(iarr.size, 12)
+
+    expected_values = [
+        [[120, 125], [140, 145], [180, 185]],
+        [[320, 325], [340, 345], [380, 385]],
+    ]
+    expected_arr = wp.array(data=expected_values, dtype=float, device=device)
+
+    wp.launch(kernel_3d, dim=iarr.shape, inputs=[iarr, expected_arr], device=device)
+
+
+@wp.kernel
+def kernel_4d(a: wp.indexedarray4d(dtype=float), expected: wp.array4d(dtype=float)):
+    i, j, k, l = wp.tid()
+
+    # check expected values
+    wp.expect_eq(a[i, j, k, l], expected[i, j, k, l])
+
+    # test wp.view()
+    wp.expect_eq(a[i][j][k][l], a[i, j, k, l])
+    wp.expect_eq(a[i][j, k, l], a[i, j, k, l])
+    wp.expect_eq(a[i, j][k, l], a[i, j, k, l])
+    wp.expect_eq(a[i, j, k][l], a[i, j, k, l])
+
+    a[i, j, k, l] = 2.0 * a[i, j, k, l]
+
+    wp.atomic_add(a, i, j, k, l, 1.0)
+
+    wp.expect_eq(a[i, j, k, l], 2.0 * expected[i, j, k, l] + 1.0)
+
+
+def test_indexedarray_4d(test, device):
+    values = np.arange(10000, dtype=np.float32).reshape((10, 10, 10, 10))
+    arr = wp.array(data=values, device=device)
+
+    indices0 = wp.array([1, 3], dtype=int, device=device)
+    indices1 = wp.array([2, 4, 8], dtype=int, device=device)
+    indices2 = wp.array([0, 5], dtype=int, device=device)
+    indices3 = wp.array([6, 7, 9], dtype=int, device=device)
+
+    iarr = wp.indexedarray4d(arr, [indices0, indices1, indices2, indices3])
+
+    test.assertEqual(iarr.dtype, arr.dtype)
+    test.assertEqual(iarr.ndim, 4)
+    test.assertEqual(iarr.shape, (2, 3, 2, 3))
+    test.assertEqual(iarr.size, 36)
+
+    expected_values = [
+        [
+            [[1206, 1207, 1209], [1256, 1257, 1259]],
+            [[1406, 1407, 1409], [1456, 1457, 1459]],
+            [[1806, 1807, 1809], [1856, 1857, 1859]],
+        ],
+        [
+            [[3206, 3207, 3209], [3256, 3257, 3259]],
+            [[3406, 3407, 3409], [3456, 3457, 3459]],
+            [[3806, 3807, 3809], [3856, 3857, 3859]],
+        ],
+    ]
+    expected_arr = wp.array(data=expected_values, dtype=float, device=device)
+
+    wp.launch(kernel_4d, dim=iarr.shape, inputs=[iarr, expected_arr], device=device)
+
+
+def test_indexedarray_mixed(test, device):
+    # [[[ 0,  1,  2,  3],
+    #   [ 4,  5,  6,  7],
+    #   [ 8,  9, 10, 11],
+    #   [12, 13, 14, 15]],
+    #  [[16, 17, 18, 19],
+    #   [20, 21, 22, 23],
+    #   [24, 25, 26, 27],
+    #   [28, 29, 30, 31]],
+    #  [[32, 33, 34, 35],
+    #   [36, 37, 38, 39],
+    #   [40, 41, 42, 43],
+    #   [44, 45, 46, 47],
+    #  [[48, 49, 50, 51],
+    #   [52, 53, 54, 55],
+    #   [56, 57, 58, 59],
+    #   [60, 61, 62, 63]]]]
+    values = np.arange(64, dtype=np.float32).reshape((4, 4, 4))
+
+    indices = wp.array([0, 3], dtype=int, device=device)
+
+    # -----
+
+    arr = wp.array(data=values, device=device)
+    iarr = wp.indexedarray(arr, [indices, None, None])
+    test.assertEqual(iarr.dtype, arr.dtype)
+    test.assertEqual(iarr.ndim, 3)
+    test.assertEqual(iarr.shape, (2, 4, 4))
+    test.assertEqual(iarr.size, 32)
+
+    expected_values = [
+        [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11], [12, 13, 14, 15]],
+        [[48, 49, 50, 51], [52, 53, 54, 55], [56, 57, 58, 59], [60, 61, 62, 63]],
+    ]
+    expected_arr = wp.array(data=expected_values, dtype=float, device=device)
+
+    wp.launch(kernel_3d, dim=iarr.shape, inputs=[iarr, expected_arr], device=device)
+
+    # -----
+
+    arr = wp.array(data=values, device=device)
+    iarr = wp.indexedarray(arr, [indices, indices, None])
+    test.assertEqual(iarr.dtype, arr.dtype)
+    test.assertEqual(iarr.ndim, 3)
+    test.assertEqual(iarr.shape, (2, 2, 4))
+    test.assertEqual(iarr.size, 16)
+
+    expected_values = [[[0, 1, 2, 3], [12, 13, 14, 15]], [[48, 49, 50, 51], [60, 61, 62, 63]]]
+    expected_arr = wp.array(data=expected_values, dtype=float, device=device)
+
+    wp.launch(kernel_3d, dim=iarr.shape, inputs=[iarr, expected_arr], device=device)
+
+    # -----
+
+    arr = wp.array(data=values, device=device)
+    iarr = wp.indexedarray(arr, [indices, None, indices])
+    test.assertEqual(iarr.dtype, arr.dtype)
+    test.assertEqual(iarr.ndim, 3)
+    test.assertEqual(iarr.shape, (2, 4, 2))
+    test.assertEqual(iarr.size, 16)
+
+    expected_values = [[[0, 3], [4, 7], [8, 11], [12, 15]], [[48, 51], [52, 55], [56, 59], [60, 63]]]
+    expected_arr = wp.array(data=expected_values, dtype=float, device=device)
+
+    wp.launch(kernel_3d, dim=iarr.shape, inputs=[iarr, expected_arr], device=device)
+
+    # -----
+
+    arr = wp.array(data=values, device=device)
+    iarr = wp.indexedarray(arr, [None, indices, indices])
+    test.assertEqual(iarr.dtype, arr.dtype)
+    test.assertEqual(iarr.ndim, 3)
+    test.assertEqual(iarr.shape, (4, 2, 2))
+    test.assertEqual(iarr.size, 16)
+
+    expected_values = [[[0, 3], [12, 15]], [[16, 19], [28, 31]], [[32, 35], [44, 47]], [[48, 51], [60, 63]]]
+    expected_arr = wp.array(data=expected_values, dtype=float, device=device)
+
+    wp.launch(kernel_3d, dim=iarr.shape, inputs=[iarr, expected_arr], device=device)
+
+
+vec2i = wp.types.vector(length=2, dtype=wp.int32)
+vec3i = wp.types.vector(length=3, dtype=wp.int32)
+vec4i = wp.types.vector(length=4, dtype=wp.int32)
+
+
+@wp.kernel
+def shape_kernel_1d(arr: wp.indexedarray1d(dtype=float), expected: int):
+    wp.expect_eq(arr.shape[0], expected)
+
+
+@wp.kernel
+def shape_kernel_2d(arr: wp.indexedarray2d(dtype=float), expected: vec2i):
+    wp.expect_eq(arr.shape[0], expected[0])
+    wp.expect_eq(arr.shape[1], expected[1])
+
+    # 1d slice
+    view = arr[0]
+    wp.expect_eq(view.shape[0], expected[1])
+
+
+@wp.kernel
+def shape_kernel_3d(arr: wp.indexedarray3d(dtype=float), expected: vec3i):
+    wp.expect_eq(arr.shape[0], expected[0])
+    wp.expect_eq(arr.shape[1], expected[1])
+    wp.expect_eq(arr.shape[2], expected[2])
+
+    # 2d slice
+    view2 = arr[0]
+    wp.expect_eq(view2.shape[0], expected[1])
+    wp.expect_eq(view2.shape[1], expected[2])
+
+    # 1d slice
+    view1 = arr[0, 0]
+    wp.expect_eq(view1.shape[0], expected[2])
+
+
+@wp.kernel
+def shape_kernel_4d(arr: wp.indexedarray4d(dtype=float), expected: vec4i):
+    wp.expect_eq(arr.shape[0], expected[0])
+    wp.expect_eq(arr.shape[1], expected[1])
+    wp.expect_eq(arr.shape[2], expected[2])
+    wp.expect_eq(arr.shape[3], expected[3])
+
+    # 3d slice
+    view3 = arr[0]
+    wp.expect_eq(view3.shape[0], expected[1])
+    wp.expect_eq(view3.shape[1], expected[2])
+    wp.expect_eq(view3.shape[2], expected[3])
+
+    # 2d slice
+    view2 = arr[0, 0]
+    wp.expect_eq(view2.shape[0], expected[2])
+    wp.expect_eq(view2.shape[1], expected[3])
+
+    # 1d slice
+    view1 = arr[0, 0, 0]
+    wp.expect_eq(view1.shape[0], expected[3])
+
+
+def test_indexedarray_shape(test, device):
+    with wp.ScopedDevice(device):
+        data1 = wp.zeros(10, dtype=float)
+        data2 = wp.zeros((10, 20), dtype=float)
+        data3 = wp.zeros((10, 20, 30), dtype=float)
+        data4 = wp.zeros((10, 20, 30, 40), dtype=float)
+
+        indices1 = wp.array(data=[2, 7], dtype=int)
+        indices2 = wp.array(data=[2, 7, 12, 17], dtype=int)
+        indices3 = wp.array(data=[2, 7, 12, 17, 22, 27], dtype=int)
+        indices4 = wp.array(data=[2, 7, 12, 17, 22, 27, 32, 37], dtype=int)
+
+        ia1 = wp.indexedarray(data1, [indices1])
+        wp.launch(shape_kernel_1d, dim=1, inputs=[ia1, 2])
+
+        ia2_1 = wp.indexedarray(data2, [indices1, None])
+        ia2_2 = wp.indexedarray(data2, [None, indices2])
+        ia2_3 = wp.indexedarray(data2, [indices1, indices2])
+        wp.launch(shape_kernel_2d, dim=1, inputs=[ia2_1, vec2i(2, 20)])
+        wp.launch(shape_kernel_2d, dim=1, inputs=[ia2_2, vec2i(10, 4)])
+        wp.launch(shape_kernel_2d, dim=1, inputs=[ia2_3, vec2i(2, 4)])
+
+        ia3_1 = wp.indexedarray(data3, [indices1, None, None])
+        ia3_2 = wp.indexedarray(data3, [None, indices2, None])
+        ia3_3 = wp.indexedarray(data3, [None, None, indices3])
+        ia3_4 = wp.indexedarray(data3, [indices1, indices2, None])
+        ia3_5 = wp.indexedarray(data3, [indices1, None, indices3])
+        ia3_6 = wp.indexedarray(data3, [None, indices2, indices3])
+        ia3_7 = wp.indexedarray(data3, [indices1, indices2, indices3])
+        wp.launch(shape_kernel_3d, dim=1, inputs=[ia3_1, vec3i(2, 20, 30)])
+        wp.launch(shape_kernel_3d, dim=1, inputs=[ia3_2, vec3i(10, 4, 30)])
+        wp.launch(shape_kernel_3d, dim=1, inputs=[ia3_3, vec3i(10, 20, 6)])
+        wp.launch(shape_kernel_3d, dim=1, inputs=[ia3_4, vec3i(2, 4, 30)])
+        wp.launch(shape_kernel_3d, dim=1, inputs=[ia3_5, vec3i(2, 20, 6)])
+        wp.launch(shape_kernel_3d, dim=1, inputs=[ia3_6, vec3i(10, 4, 6)])
+        wp.launch(shape_kernel_3d, dim=1, inputs=[ia3_7, vec3i(2, 4, 6)])
+
+        ia4_1 = wp.indexedarray(data4, [indices1, None, None, None])
+        ia4_2 = wp.indexedarray(data4, [indices1, None, None, indices4])
+        ia4_3 = wp.indexedarray(data4, [None, indices2, indices3, None])
+        ia4_4 = wp.indexedarray(data4, [indices1, indices2, indices3, indices4])
+        wp.launch(shape_kernel_4d, dim=1, inputs=[ia4_1, vec4i(2, 20, 30, 40)])
+        wp.launch(shape_kernel_4d, dim=1, inputs=[ia4_2, vec4i(2, 20, 30, 8)])
+        wp.launch(shape_kernel_4d, dim=1, inputs=[ia4_3, vec4i(10, 4, 6, 40)])
+        wp.launch(shape_kernel_4d, dim=1, inputs=[ia4_4, vec4i(2, 4, 6, 8)])
+
+        wp.synchronize_device(device)
+
+
+def test_indexedarray_getitem(test, device):
+    with wp.ScopedDevice(device):
+        data = wp.array(data=np.arange(1000, dtype=np.int32).reshape((10, 10, 10)))
+
+        I = wp.array(data=[0, 1, 2], dtype=int)
+
+        # use constructor
+        a1 = wp.indexedarray(data, [None, None, I])
+        a2 = wp.indexedarray(data, [None, I])
+        a3 = wp.indexedarray(data, [None, I, I])
+        a4 = wp.indexedarray(data, [I])
+        a5 = wp.indexedarray(data, [I, None, I])
+        a6 = wp.indexedarray(data, [I, I])
+        a7 = wp.indexedarray(data, [I, I, I])
+
+        # use array.__getitem__()
+        b1 = data[:, :, I]
+        b2 = data[:, I]
+        b3 = data[:, I, I]
+        b4 = data[I]
+        b5 = data[I, :, I]
+        b6 = data[I, I]
+        b7 = data[I, I, I]
+
+        test.assertEqual(type(a1), type(b1))
+        test.assertEqual(type(a2), type(b2))
+        test.assertEqual(type(a3), type(b3))
+        test.assertEqual(type(a4), type(b4))
+        test.assertEqual(type(a5), type(b5))
+        test.assertEqual(type(a6), type(b6))
+        test.assertEqual(type(a7), type(b7))
+
+        assert_np_equal(a1.numpy(), b1.numpy())
+        assert_np_equal(a2.numpy(), b2.numpy())
+        assert_np_equal(a3.numpy(), b3.numpy())
+        assert_np_equal(a4.numpy(), b4.numpy())
+        assert_np_equal(a5.numpy(), b5.numpy())
+        assert_np_equal(a6.numpy(), b6.numpy())
+        assert_np_equal(a7.numpy(), b7.numpy())
+
+
+def test_indexedarray_slicing(test, device):
+    with wp.ScopedDevice(device):
+        data = wp.array(data=np.arange(1000, dtype=np.int32).reshape((10, 10, 10)))
+
+        # test equivalence of slicing and indexing the same range
+        s = slice(0, 3)
+        I = wp.array(data=[0, 1, 2], dtype=int)
+
+        a0 = data[s, s, s]
+        test.assertEqual(type(a0), wp.array)
+        a1 = data[s, s, I]
+        test.assertEqual(type(a1), wp.indexedarray)
+        a2 = data[s, I, s]
+        test.assertEqual(type(a2), wp.indexedarray)
+        a3 = data[s, I, I]
+        test.assertEqual(type(a3), wp.indexedarray)
+        a4 = data[I, s, s]
+        test.assertEqual(type(a4), wp.indexedarray)
+        a5 = data[I, s, I]
+        test.assertEqual(type(a5), wp.indexedarray)
+        a6 = data[I, I, s]
+        test.assertEqual(type(a6), wp.indexedarray)
+        a7 = data[I, I, I]
+        test.assertEqual(type(a7), wp.indexedarray)
+
+        expected = a0.numpy()
+
+        assert_np_equal(a1.numpy(), expected)
+        assert_np_equal(a2.numpy(), expected)
+        assert_np_equal(a3.numpy(), expected)
+        assert_np_equal(a4.numpy(), expected)
+        assert_np_equal(a5.numpy(), expected)
+        assert_np_equal(a6.numpy(), expected)
+        assert_np_equal(a7.numpy(), expected)
+
+
+# generic increment kernels that work with any array (regular or indexed)
+@wp.kernel
+def inc_1d(a: Any):
+    i = wp.tid()
+    a[i] = a[i] + 1
+
+
+@wp.kernel
+def inc_2d(a: Any):
+    i, j = wp.tid()
+    a[i, j] = a[i, j] + 1
+
+
+@wp.kernel
+def inc_3d(a: Any):
+    i, j, k = wp.tid()
+    a[i, j, k] = a[i, j, k] + 1
+
+
+@wp.kernel
+def inc_4d(a: Any):
+    i, j, k, l = wp.tid()
+    a[i, j, k, l] = a[i, j, k, l] + 1
+
+
+# optional overloads to avoid module reloading
+wp.overload(inc_1d, [wp.array1d(dtype=int)])
+wp.overload(inc_2d, [wp.array2d(dtype=int)])
+wp.overload(inc_3d, [wp.array3d(dtype=int)])
+wp.overload(inc_4d, [wp.array4d(dtype=int)])
+
+wp.overload(inc_1d, [wp.indexedarray1d(dtype=int)])
+wp.overload(inc_2d, [wp.indexedarray2d(dtype=int)])
+wp.overload(inc_3d, [wp.indexedarray3d(dtype=int)])
+wp.overload(inc_4d, [wp.indexedarray4d(dtype=int)])
+
+
+def test_indexedarray_generics(test, device):
+    with wp.ScopedDevice(device):
+        data1 = wp.zeros((5,), dtype=int)
+        data2 = wp.zeros((5, 5), dtype=int)
+        data3 = wp.zeros((5, 5, 5), dtype=int)
+        data4 = wp.zeros((5, 5, 5, 5), dtype=int)
+
+        indices = wp.array(data=[0, 4], dtype=int)
+
+        ia1 = wp.indexedarray(data1, [indices])
+        ia2 = wp.indexedarray(data2, [indices, indices])
+        ia3 = wp.indexedarray(data3, [indices, indices, indices])
+        ia4 = wp.indexedarray(data4, [indices, indices, indices, indices])
+
+        wp.launch(inc_1d, dim=data1.shape, inputs=[data1])
+        wp.launch(inc_2d, dim=data2.shape, inputs=[data2])
+        wp.launch(inc_3d, dim=data3.shape, inputs=[data3])
+        wp.launch(inc_4d, dim=data4.shape, inputs=[data4])
+
+        wp.launch(inc_1d, dim=ia1.shape, inputs=[ia1])
+        wp.launch(inc_2d, dim=ia2.shape, inputs=[ia2])
+        wp.launch(inc_3d, dim=ia3.shape, inputs=[ia3])
+        wp.launch(inc_4d, dim=ia4.shape, inputs=[ia4])
+
+        expected1 = np.ones(5, dtype=np.int32)
+        expected1[0] = 2
+        expected1[4] = 2
+
+        expected2 = np.ones((5, 5), dtype=np.int32)
+        expected2[0, 0] = 2
+        expected2[0, 4] = 2
+        expected2[4, 0] = 2
+        expected2[4, 4] = 2
+
+        expected3 = np.ones((5, 5, 5), dtype=np.int32)
+        expected3[0, 0, 0] = 2
+        expected3[0, 0, 4] = 2
+        expected3[0, 4, 0] = 2
+        expected3[0, 4, 4] = 2
+        expected3[4, 0, 0] = 2
+        expected3[4, 0, 4] = 2
+        expected3[4, 4, 0] = 2
+        expected3[4, 4, 4] = 2
+
+        expected4 = np.ones((5, 5, 5, 5), dtype=np.int32)
+        expected4[0, 0, 0, 0] = 2
+        expected4[0, 0, 0, 4] = 2
+        expected4[0, 0, 4, 0] = 2
+        expected4[0, 0, 4, 4] = 2
+        expected4[0, 4, 0, 0] = 2
+        expected4[0, 4, 0, 4] = 2
+        expected4[0, 4, 4, 0] = 2
+        expected4[0, 4, 4, 4] = 2
+        expected4[4, 0, 0, 0] = 2
+        expected4[4, 0, 0, 4] = 2
+        expected4[4, 0, 4, 0] = 2
+        expected4[4, 0, 4, 4] = 2
+        expected4[4, 4, 0, 0] = 2
+        expected4[4, 4, 0, 4] = 2
+        expected4[4, 4, 4, 0] = 2
+        expected4[4, 4, 4, 4] = 2
+
+        assert_np_equal(data1.numpy(), expected1)
+        assert_np_equal(data2.numpy(), expected2)
+        assert_np_equal(data3.numpy(), expected3)
+        assert_np_equal(data4.numpy(), expected4)
+
+        assert_np_equal(ia1.numpy(), np.full((2,), 2, dtype=np.int32))
+        assert_np_equal(ia2.numpy(), np.full((2, 2), 2, dtype=np.int32))
+        assert_np_equal(ia3.numpy(), np.full((2, 2, 2), 2, dtype=np.int32))
+        assert_np_equal(ia4.numpy(), np.full((2, 2, 2, 2), 2, dtype=np.int32))
+
+
+def test_indexedarray_empty(test, device):
+    # Test whether common operations work with empty (zero-sized) indexed arrays
+    # without throwing exceptions.
+
+    def test_empty_ops(ndim, nrows, ncols, wptype, nptype):
+        data_shape = (1,) * ndim
+        dtype_shape = ()
+
+        if wptype in wp.types.scalar_types:
+            # scalar, vector, or matrix
+            if ncols > 0:
+                if nrows > 0:
+                    wptype = wp.types.matrix((nrows, ncols), wptype)
+                else:
+                    wptype = wp.types.vector(ncols, wptype)
+                dtype_shape = wptype._shape_
+            fill_value = wptype(42)
+        else:
+            # struct
+            fill_value = wptype()
+
+        # create a data array
+        data = wp.empty(data_shape, dtype=wptype, device=device, requires_grad=True)
+
+        # create a zero-sized array of indices
+        indices = wp.empty(0, dtype=int, device=device)
+
+        a = data[indices]
+
+        # we expect dim to be zero for the empty indexed array, unchanged otherwise
+        expected_shape = (0, *data_shape[1:])
+
+        test.assertEqual(a.size, 0)
+        test.assertEqual(a.shape, expected_shape)
+
+        # all of these methods should succeed with zero-sized arrays
+        a.zero_()
+        a.fill_(fill_value)
+        b = a.contiguous()
+
+        b = wp.empty_like(a)
+        b = wp.zeros_like(a)
+        b = wp.full_like(a, fill_value)
+        b = wp.clone(a)
+
+        wp.copy(a, b)
+        a.assign(b)
+
+        na = a.numpy()
+        test.assertEqual(na.size, 0)
+        test.assertEqual(na.shape, (*expected_shape, *dtype_shape))
+        test.assertEqual(na.dtype, nptype)
+
+        test.assertEqual(a.list(), [])
+
+    for ndim in range(1, 5):
+        # test with scalars, vectors, and matrices
+        for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
+            # scalars
+            test_empty_ops(ndim, 0, 0, wptype, nptype)
+
+            for ncols in [2, 3, 4, 5]:
+                # vectors
+                test_empty_ops(ndim, 0, ncols, wptype, nptype)
+                # square matrices
+                test_empty_ops(ndim, ncols, ncols, wptype, nptype)
+
+            # non-square matrices
+            test_empty_ops(ndim, 2, 3, wptype, nptype)
+            test_empty_ops(ndim, 3, 2, wptype, nptype)
+            test_empty_ops(ndim, 3, 4, wptype, nptype)
+            test_empty_ops(ndim, 4, 3, wptype, nptype)
+
+        # test with structs
+        test_empty_ops(ndim, 0, 0, FillStruct, FillStruct.numpy_dtype())
+
+
+def test_indexedarray_fill_scalar(test, device):
+    dim_x = 4
+
+    for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
+        data1 = wp.zeros(dim_x, dtype=wptype, device=device)
+        data2 = wp.zeros((dim_x, dim_x), dtype=wptype, device=device)
+        data3 = wp.zeros((dim_x, dim_x, dim_x), dtype=wptype, device=device)
+        data4 = wp.zeros((dim_x, dim_x, dim_x, dim_x), dtype=wptype, device=device)
+
+        indices = wp.array(np.arange(0, dim_x, 2, dtype=np.int32), device=device)
+
+        a1 = data1[indices]
+        a2 = data2[indices]
+        a3 = data3[indices]
+        a4 = data4[indices]
+
+        assert_np_equal(a1.numpy(), np.zeros(a1.shape, dtype=nptype))
+        assert_np_equal(a2.numpy(), np.zeros(a2.shape, dtype=nptype))
+        assert_np_equal(a3.numpy(), np.zeros(a3.shape, dtype=nptype))
+        assert_np_equal(a4.numpy(), np.zeros(a4.shape, dtype=nptype))
+
+        # fill with int value
+        fill_value = 42
+
+        a1.fill_(fill_value)
+        a2.fill_(fill_value)
+        a3.fill_(fill_value)
+        a4.fill_(fill_value)
+
+        assert_np_equal(a1.numpy(), np.full(a1.shape, fill_value, dtype=nptype))
+        assert_np_equal(a2.numpy(), np.full(a2.shape, fill_value, dtype=nptype))
+        assert_np_equal(a3.numpy(), np.full(a3.shape, fill_value, dtype=nptype))
+        assert_np_equal(a4.numpy(), np.full(a4.shape, fill_value, dtype=nptype))
+
+        a1.zero_()
+        a2.zero_()
+        a3.zero_()
+        a4.zero_()
+
+        assert_np_equal(a1.numpy(), np.zeros(a1.shape, dtype=nptype))
+        assert_np_equal(a2.numpy(), np.zeros(a2.shape, dtype=nptype))
+        assert_np_equal(a3.numpy(), np.zeros(a3.shape, dtype=nptype))
+        assert_np_equal(a4.numpy(), np.zeros(a4.shape, dtype=nptype))
+
+        if wptype in wp.types.float_types:
+            # fill with float value
+            fill_value = 13.37
+
+            a1.fill_(fill_value)
+            a2.fill_(fill_value)
+            a3.fill_(fill_value)
+            a4.fill_(fill_value)
+
+            assert_np_equal(a1.numpy(), np.full(a1.shape, fill_value, dtype=nptype))
+            assert_np_equal(a2.numpy(), np.full(a2.shape, fill_value, dtype=nptype))
+            assert_np_equal(a3.numpy(), np.full(a3.shape, fill_value, dtype=nptype))
+            assert_np_equal(a4.numpy(), np.full(a4.shape, fill_value, dtype=nptype))
+
+        # fill with Warp scalar value
+        fill_value = wptype(17)
+
+        a1.fill_(fill_value)
+        a2.fill_(fill_value)
+        a3.fill_(fill_value)
+        a4.fill_(fill_value)
+
+        assert_np_equal(a1.numpy(), np.full(a1.shape, fill_value.value, dtype=nptype))
+        assert_np_equal(a2.numpy(), np.full(a2.shape, fill_value.value, dtype=nptype))
+        assert_np_equal(a3.numpy(), np.full(a3.shape, fill_value.value, dtype=nptype))
+        assert_np_equal(a4.numpy(), np.full(a4.shape, fill_value.value, dtype=nptype))
+
+
+def test_indexedarray_fill_vector(test, device):
+    # test filling a vector array with scalar or vector values (vec_type, list, or numpy array)
+
+    dim_x = 4
+
+    for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
+        # vector types
+        vector_types = [
+            wp.types.vector(2, wptype),
+            wp.types.vector(3, wptype),
+            wp.types.vector(4, wptype),
+            wp.types.vector(5, wptype),
+        ]
+
+        for vec_type in vector_types:
+            vec_len = vec_type._length_
+
+            data1 = wp.zeros(dim_x, dtype=vec_type, device=device)
+            data2 = wp.zeros((dim_x, dim_x), dtype=vec_type, device=device)
+            data3 = wp.zeros((dim_x, dim_x, dim_x), dtype=vec_type, device=device)
+            data4 = wp.zeros((dim_x, dim_x, dim_x, dim_x), dtype=vec_type, device=device)
+
+            indices = wp.array(np.arange(0, dim_x, 2, dtype=np.int32), device=device)
+
+            a1 = data1[indices]
+            a2 = data2[indices]
+            a3 = data3[indices]
+            a4 = data4[indices]
+
+            assert_np_equal(a1.numpy(), np.zeros((*a1.shape, vec_len), dtype=nptype))
+            assert_np_equal(a2.numpy(), np.zeros((*a2.shape, vec_len), dtype=nptype))
+            assert_np_equal(a3.numpy(), np.zeros((*a3.shape, vec_len), dtype=nptype))
+            assert_np_equal(a4.numpy(), np.zeros((*a4.shape, vec_len), dtype=nptype))
+
+            # fill with int scalar
+            fill_value = 42
+
+            a1.fill_(fill_value)
+            a2.fill_(fill_value)
+            a3.fill_(fill_value)
+            a4.fill_(fill_value)
+
+            assert_np_equal(a1.numpy(), np.full((*a1.shape, vec_len), fill_value, dtype=nptype))
+            assert_np_equal(a2.numpy(), np.full((*a2.shape, vec_len), fill_value, dtype=nptype))
+            assert_np_equal(a3.numpy(), np.full((*a3.shape, vec_len), fill_value, dtype=nptype))
+            assert_np_equal(a4.numpy(), np.full((*a4.shape, vec_len), fill_value, dtype=nptype))
+
+            # test zeroing
+            a1.zero_()
+            a2.zero_()
+            a3.zero_()
+            a4.zero_()
+
+            assert_np_equal(a1.numpy(), np.zeros((*a1.shape, vec_len), dtype=nptype))
+            assert_np_equal(a2.numpy(), np.zeros((*a2.shape, vec_len), dtype=nptype))
+            assert_np_equal(a3.numpy(), np.zeros((*a3.shape, vec_len), dtype=nptype))
+            assert_np_equal(a4.numpy(), np.zeros((*a4.shape, vec_len), dtype=nptype))
+
+            # vector values can be passed as a list, numpy array, or Warp vector instance
+            fill_list = [17, 42, 99, 101, 127][:vec_len]
+            fill_arr = np.array(fill_list, dtype=nptype)
+            fill_vec = vec_type(fill_list)
+
+            expected1 = np.tile(fill_arr, a1.size).reshape((*a1.shape, vec_len))
+            expected2 = np.tile(fill_arr, a2.size).reshape((*a2.shape, vec_len))
+            expected3 = np.tile(fill_arr, a3.size).reshape((*a3.shape, vec_len))
+            expected4 = np.tile(fill_arr, a4.size).reshape((*a4.shape, vec_len))
+
+            # fill with list of vector length
+            a1.fill_(fill_list)
+            a2.fill_(fill_list)
+            a3.fill_(fill_list)
+            a4.fill_(fill_list)
+
+            assert_np_equal(a1.numpy(), expected1)
+            assert_np_equal(a2.numpy(), expected2)
+            assert_np_equal(a3.numpy(), expected3)
+            assert_np_equal(a4.numpy(), expected4)
+
+            # clear
+            a1.zero_()
+            a2.zero_()
+            a3.zero_()
+            a4.zero_()
+
+            # fill with numpy array of vector length
+            a1.fill_(fill_arr)
+            a2.fill_(fill_arr)
+            a3.fill_(fill_arr)
+            a4.fill_(fill_arr)
+
+            assert_np_equal(a1.numpy(), expected1)
+            assert_np_equal(a2.numpy(), expected2)
+            assert_np_equal(a3.numpy(), expected3)
+            assert_np_equal(a4.numpy(), expected4)
+
+            # clear
+            a1.zero_()
+            a2.zero_()
+            a3.zero_()
+            a4.zero_()
+
+            # fill with vec instance
+            a1.fill_(fill_vec)
+            a2.fill_(fill_vec)
+            a3.fill_(fill_vec)
+            a4.fill_(fill_vec)
+
+            assert_np_equal(a1.numpy(), expected1)
+            assert_np_equal(a2.numpy(), expected2)
+            assert_np_equal(a3.numpy(), expected3)
+            assert_np_equal(a4.numpy(), expected4)
+
+            if wptype in wp.types.float_types:
+                # fill with float scalar
+                fill_value = 13.37
+
+                a1.fill_(fill_value)
+                a2.fill_(fill_value)
+                a3.fill_(fill_value)
+                a4.fill_(fill_value)
+
+                assert_np_equal(a1.numpy(), np.full((*a1.shape, vec_len), fill_value, dtype=nptype))
+                assert_np_equal(a2.numpy(), np.full((*a2.shape, vec_len), fill_value, dtype=nptype))
+                assert_np_equal(a3.numpy(), np.full((*a3.shape, vec_len), fill_value, dtype=nptype))
+                assert_np_equal(a4.numpy(), np.full((*a4.shape, vec_len), fill_value, dtype=nptype))
+
+                # fill with float list of vector length
+                fill_list = [-2.5, -1.25, 1.25, 2.5, 5.0][:vec_len]
+
+                a1.fill_(fill_list)
+                a2.fill_(fill_list)
+                a3.fill_(fill_list)
+                a4.fill_(fill_list)
+
+                expected1 = np.tile(np.array(fill_list, dtype=nptype), a1.size).reshape((*a1.shape, vec_len))
+                expected2 = np.tile(np.array(fill_list, dtype=nptype), a2.size).reshape((*a2.shape, vec_len))
+                expected3 = np.tile(np.array(fill_list, dtype=nptype), a3.size).reshape((*a3.shape, vec_len))
+                expected4 = np.tile(np.array(fill_list, dtype=nptype), a4.size).reshape((*a4.shape, vec_len))
+
+                assert_np_equal(a1.numpy(), expected1)
+                assert_np_equal(a2.numpy(), expected2)
+                assert_np_equal(a3.numpy(), expected3)
+                assert_np_equal(a4.numpy(), expected4)
+
+
+def test_indexedarray_fill_matrix(test, device):
+    # test filling a matrix array with scalar or matrix values (mat_type, nested list, or 2d numpy array)
+
+    dim_x = 4
+
+    for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
+        # matrix types
+        matrix_types = [
+            # square matrices
+            wp.types.matrix((2, 2), wptype),
+            wp.types.matrix((3, 3), wptype),
+            wp.types.matrix((4, 4), wptype),
+            wp.types.matrix((5, 5), wptype),
+            # non-square matrices
+            wp.types.matrix((2, 3), wptype),
+            wp.types.matrix((3, 2), wptype),
+            wp.types.matrix((3, 4), wptype),
+            wp.types.matrix((4, 3), wptype),
+        ]
+
+        for mat_type in matrix_types:
+            mat_len = mat_type._length_
+            mat_shape = mat_type._shape_
+
+            data1 = wp.zeros(dim_x, dtype=mat_type, device=device)
+            data2 = wp.zeros((dim_x, dim_x), dtype=mat_type, device=device)
+            data3 = wp.zeros((dim_x, dim_x, dim_x), dtype=mat_type, device=device)
+            data4 = wp.zeros((dim_x, dim_x, dim_x, dim_x), dtype=mat_type, device=device)
+
+            indices = wp.array(np.arange(0, dim_x, 2, dtype=np.int32), device=device)
+
+            a1 = data1[indices]
+            a2 = data2[indices]
+            a3 = data3[indices]
+            a4 = data4[indices]
+
+            assert_np_equal(a1.numpy(), np.zeros((*a1.shape, *mat_shape), dtype=nptype))
+            assert_np_equal(a2.numpy(), np.zeros((*a2.shape, *mat_shape), dtype=nptype))
+            assert_np_equal(a3.numpy(), np.zeros((*a3.shape, *mat_shape), dtype=nptype))
+            assert_np_equal(a4.numpy(), np.zeros((*a4.shape, *mat_shape), dtype=nptype))
+
+            # fill with scalar
+            fill_value = 42
+
+            a1.fill_(fill_value)
+            a2.fill_(fill_value)
+            a3.fill_(fill_value)
+            a4.fill_(fill_value)
+
+            assert_np_equal(a1.numpy(), np.full((*a1.shape, *mat_shape), fill_value, dtype=nptype))
+            assert_np_equal(a2.numpy(), np.full((*a2.shape, *mat_shape), fill_value, dtype=nptype))
+            assert_np_equal(a3.numpy(), np.full((*a3.shape, *mat_shape), fill_value, dtype=nptype))
+            assert_np_equal(a4.numpy(), np.full((*a4.shape, *mat_shape), fill_value, dtype=nptype))
+
+            # test zeroing
+            a1.zero_()
+            a2.zero_()
+            a3.zero_()
+            a4.zero_()
+
+            assert_np_equal(a1.numpy(), np.zeros((*a1.shape, *mat_shape), dtype=nptype))
+            assert_np_equal(a2.numpy(), np.zeros((*a2.shape, *mat_shape), dtype=nptype))
+            assert_np_equal(a3.numpy(), np.zeros((*a3.shape, *mat_shape), dtype=nptype))
+            assert_np_equal(a4.numpy(), np.zeros((*a4.shape, *mat_shape), dtype=nptype))
+
+            # matrix values can be passed as a 1d numpy array, 2d numpy array, flat list, nested list, or Warp matrix instance
+            if wptype != wp.bool:
+                fill_arr1 = np.arange(mat_len, dtype=nptype)
+            else:
+                fill_arr1 = np.ones(mat_len, dtype=nptype)
+            fill_arr2 = fill_arr1.reshape(mat_shape)
+            fill_list1 = list(fill_arr1)
+            fill_list2 = [list(row) for row in fill_arr2]
+            fill_mat = mat_type(fill_arr1)
+
+            expected1 = np.tile(fill_arr1, a1.size).reshape((*a1.shape, *mat_shape))
+            expected2 = np.tile(fill_arr1, a2.size).reshape((*a2.shape, *mat_shape))
+            expected3 = np.tile(fill_arr1, a3.size).reshape((*a3.shape, *mat_shape))
+            expected4 = np.tile(fill_arr1, a4.size).reshape((*a4.shape, *mat_shape))
+
+            # fill with 1d numpy array
+            a1.fill_(fill_arr1)
+            a2.fill_(fill_arr1)
+            a3.fill_(fill_arr1)
+            a4.fill_(fill_arr1)
+
+            assert_np_equal(a1.numpy(), expected1)
+            assert_np_equal(a2.numpy(), expected2)
+            assert_np_equal(a3.numpy(), expected3)
+            assert_np_equal(a4.numpy(), expected4)
+
+            # clear
+            a1.zero_()
+            a2.zero_()
+            a3.zero_()
+            a4.zero_()
+
+            # fill with 2d numpy array
+            a1.fill_(fill_arr2)
+            a2.fill_(fill_arr2)
+            a3.fill_(fill_arr2)
+            a4.fill_(fill_arr2)
+
+            assert_np_equal(a1.numpy(), expected1)
+            assert_np_equal(a2.numpy(), expected2)
+            assert_np_equal(a3.numpy(), expected3)
+            assert_np_equal(a4.numpy(), expected4)
+
+            # clear
+            a1.zero_()
+            a2.zero_()
+            a3.zero_()
+            a4.zero_()
+
+            # fill with flat list
+            a1.fill_(fill_list1)
+            a2.fill_(fill_list1)
+            a3.fill_(fill_list1)
+            a4.fill_(fill_list1)
+
+            assert_np_equal(a1.numpy(), expected1)
+            assert_np_equal(a2.numpy(), expected2)
+            assert_np_equal(a3.numpy(), expected3)
+            assert_np_equal(a4.numpy(), expected4)
+
+            # clear
+            a1.zero_()
+            a2.zero_()
+            a3.zero_()
+            a4.zero_()
+
+            # fill with nested list
+            a1.fill_(fill_list2)
+            a2.fill_(fill_list2)
+            a3.fill_(fill_list2)
+            a4.fill_(fill_list2)
+
+            assert_np_equal(a1.numpy(), expected1)
+            assert_np_equal(a2.numpy(), expected2)
+            assert_np_equal(a3.numpy(), expected3)
+            assert_np_equal(a4.numpy(), expected4)
+
+            # clear
+            a1.zero_()
+            a2.zero_()
+            a3.zero_()
+            a4.zero_()
+
+            # fill with mat instance
+            a1.fill_(fill_mat)
+            a2.fill_(fill_mat)
+            a3.fill_(fill_mat)
+            a4.fill_(fill_mat)
+
+            assert_np_equal(a1.numpy(), expected1)
+            assert_np_equal(a2.numpy(), expected2)
+            assert_np_equal(a3.numpy(), expected3)
+            assert_np_equal(a4.numpy(), expected4)
+
+
+def test_indexedarray_fill_struct(test, device):
+    dim_x = 8
+
+    nptype = FillStruct.numpy_dtype()
+
+    data1 = wp.zeros(dim_x, dtype=FillStruct, device=device)
+    data2 = wp.zeros((dim_x, dim_x), dtype=FillStruct, device=device)
+    data3 = wp.zeros((dim_x, dim_x, dim_x), dtype=FillStruct, device=device)
+    data4 = wp.zeros((dim_x, dim_x, dim_x, dim_x), dtype=FillStruct, device=device)
+
+    indices = wp.array(np.arange(0, dim_x, 2, dtype=np.int32), device=device)
+
+    a1 = data1[indices]
+    a2 = data2[indices]
+    a3 = data3[indices]
+    a4 = data4[indices]
+
+    assert_np_equal(a1.numpy(), np.zeros(a1.shape, dtype=nptype))
+    assert_np_equal(a2.numpy(), np.zeros(a2.shape, dtype=nptype))
+    assert_np_equal(a3.numpy(), np.zeros(a3.shape, dtype=nptype))
+    assert_np_equal(a4.numpy(), np.zeros(a4.shape, dtype=nptype))
+
+    s = FillStruct()
+
+    # fill with default struct value (should be all zeros)
+    a1.fill_(s)
+    a2.fill_(s)
+    a3.fill_(s)
+    a4.fill_(s)
+
+    assert_np_equal(a1.numpy(), np.zeros(a1.shape, dtype=nptype))
+    assert_np_equal(a2.numpy(), np.zeros(a2.shape, dtype=nptype))
+    assert_np_equal(a3.numpy(), np.zeros(a3.shape, dtype=nptype))
+    assert_np_equal(a4.numpy(), np.zeros(a4.shape, dtype=nptype))
+
+    # scalars
+    s.i1 = -17
+    s.i2 = 42
+    s.i4 = 99
+    s.i8 = 101
+    s.f2 = -1.25
+    s.f4 = 13.37
+    s.f8 = 0.125
+    # vectors
+    s.v2 = [21, 22]
+    s.v3 = [31, 32, 33]
+    s.v4 = [41, 42, 43, 44]
+    s.v5 = [51, 52, 53, 54, 55]
+    # matrices
+    s.m2 = [[61, 62]] * 2
+    s.m3 = [[71, 72, 73]] * 3
+    s.m4 = [[81, 82, 83, 84]] * 4
+    s.m5 = [[91, 92, 93, 94, 95]] * 5
+    # arrays
+    s.a1 = wp.zeros((2,) * 1, dtype=float, device=device)
+    s.a2 = wp.zeros((2,) * 2, dtype=float, device=device)
+    s.a3 = wp.zeros((2,) * 3, dtype=float, device=device)
+    s.a4 = wp.zeros((2,) * 4, dtype=float, device=device)
+
+    # fill with custom struct value
+    a1.fill_(s)
+    a2.fill_(s)
+    a3.fill_(s)
+    a4.fill_(s)
+
+    ns = s.numpy_value()
+
+    expected1 = np.empty(a1.shape, dtype=nptype)
+    expected2 = np.empty(a2.shape, dtype=nptype)
+    expected3 = np.empty(a3.shape, dtype=nptype)
+    expected4 = np.empty(a4.shape, dtype=nptype)
+
+    expected1.fill(ns)
+    expected2.fill(ns)
+    expected3.fill(ns)
+    expected4.fill(ns)
+
+    assert_np_equal(a1.numpy(), expected1)
+    assert_np_equal(a2.numpy(), expected2)
+    assert_np_equal(a3.numpy(), expected3)
+    assert_np_equal(a4.numpy(), expected4)
+
+    # test clearing
+    a1.zero_()
+    a2.zero_()
+    a3.zero_()
+    a4.zero_()
+
+    assert_np_equal(a1.numpy(), np.zeros(a1.shape, dtype=nptype))
+    assert_np_equal(a2.numpy(), np.zeros(a2.shape, dtype=nptype))
+    assert_np_equal(a3.numpy(), np.zeros(a3.shape, dtype=nptype))
+    assert_np_equal(a4.numpy(), np.zeros(a4.shape, dtype=nptype))
+
+
+devices = get_test_devices()
+
+
+class TestIndexedArray(unittest.TestCase):
+    pass
+
+
+add_function_test(TestIndexedArray, "test_indexedarray_1d", test_indexedarray_1d, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_2d", test_indexedarray_2d, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_3d", test_indexedarray_3d, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_4d", test_indexedarray_4d, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_mixed", test_indexedarray_mixed, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_shape", test_indexedarray_shape, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_getitem", test_indexedarray_getitem, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_slicing", test_indexedarray_slicing, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_generics", test_indexedarray_generics, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_empty", test_indexedarray_empty, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_fill_scalar", test_indexedarray_fill_scalar, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_fill_vector", test_indexedarray_fill_vector, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_fill_matrix", test_indexedarray_fill_matrix, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_fill_struct", test_indexedarray_fill_struct, devices=devices)
+
+
+if __name__ == "__main__":
+    wp.build.clear_kernel_cache()
+    unittest.main(verbosity=2)