diff --git a/python/taichi/_kernels.py b/python/taichi/_kernels.py index 801f869e73876..d98da79cfbbb6 100644 --- a/python/taichi/_kernels.py +++ b/python/taichi/_kernels.py @@ -148,8 +148,12 @@ def tensor_to_tensor(tensor: template(), other: template()): @kernel def ext_arr_to_tensor(arr: ndarray_type.ndarray(), tensor: template()): + offset = static(tensor.snode.ptr.offset) + shape = static(tensor.shape) + # default value of offset is [], replace it with [0] * len + offset_new = static([0] * len(shape) if len(offset) == 0 else offset) for I in grouped(tensor): - tensor[I] = arr[I] + tensor[I] = arr[I - offset_new] @kernel diff --git a/tests/python/test_field.py b/tests/python/test_field.py index b8da9a72a317b..2d36d7b284ecc 100644 --- a/tests/python/test_field.py +++ b/tests/python/test_field.py @@ -84,6 +84,30 @@ def test_scalr_field_from_numpy(dtype, shape): x.from_numpy(arr) +@pytest.mark.parametrize("dtype", data_types) +@pytest.mark.parametrize("shape, offset", [((), ()), (8, 0), (8, 8), (8, -4), ((6, 12), (-4, -4)), ((6, 12), (-4, 4)), ((6, 12), (4, -4)), ((6, 12), (8, 8))]) +@test_utils.test(arch=get_host_arch_list()) +def test_scalr_field_from_numpy_with_offset(dtype, shape, offset): + import numpy as np + + x = ti.field(dtype=dtype, shape=shape, offset=offset) + # use the corresponding dtype for the numpy array. + numpy_dtypes = { + ti.i32: np.int32, + ti.f32: np.float32, + ti.f64: np.float64, + ti.i64: np.int64, + } + arr = np.ones(shape, dtype=numpy_dtypes[dtype]) + x.from_numpy(arr) + + def mat_equal(A, B, tol=1e-6): + return np.max(np.abs(A - B)) < tol + + tol = 1e-5 if dtype == ti.f32 else 1e-12 + assert mat_equal(x.to_numpy(), arr, tol=tol) + + @pytest.mark.parametrize("dtype", data_types) @pytest.mark.parametrize("shape", field_shapes) @test_utils.test(arch=get_host_arch_list())