From 9db8af2be12c70c0537303ff25430fc7f1de4ac9 Mon Sep 17 00:00:00 2001 From: Tianyi Liu Date: Mon, 8 May 2023 21:31:43 +0800 Subject: [PATCH] Fix issues in matrix from_numpy() --- python/taichi/_kernels.py | 13 +++++++++---- tests/python/test_matrix.py | 26 ++++++++++++++++++++++++++ 2 files changed, 35 insertions(+), 4 deletions(-) diff --git a/python/taichi/_kernels.py b/python/taichi/_kernels.py index d98da79cfbbb6..e6f7c116dfeab 100644 --- a/python/taichi/_kernels.py +++ b/python/taichi/_kernels.py @@ -214,19 +214,24 @@ def matrix_to_ext_arr(mat: template(), arr: ndarray_type.ndarray(), as_vector: t @kernel def ext_arr_to_matrix(arr: ndarray_type.ndarray(), mat: template(), as_vector: template()): + offset = static(mat.snode.ptr.offset) + shape = static(mat.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(mat): for p in static(range(mat.n)): for q in static(range(mat.m)): if static(getattr(mat, "ndim", 2) == 1): if static(as_vector): - mat[I][p] = arr[I, p] + mat[I][p] = arr[I - offset_new, p] else: - mat[I][p] = arr[I, p, q] + mat[I][p] = arr[I - offset_new, p, q] else: if static(as_vector): - mat[I][p, q] = arr[I, p] + mat[I][p, q] = arr[I - offset_new, p] else: - mat[I][p, q] = arr[I, p, q] + mat[I][p, q] = arr[I - offset_new, p, q] # extract ndarray of raw vulkan memory layout to normal memory layout. diff --git a/tests/python/test_matrix.py b/tests/python/test_matrix.py index a146d15baa44c..15128e45e10c6 100644 --- a/tests/python/test_matrix.py +++ b/tests/python/test_matrix.py @@ -1303,3 +1303,29 @@ def access_mat(i: ti.i32, j: ti.i32): # access_mat(-1, 10) # with pytest.raises(AssertionError, match=r"Out of bound access"): # access_mat(3, -1) + + +@pytest.mark.parametrize("dtype", [ti.i32, ti.f32, ti.i64, ti.f64]) +@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_matrix_from_numpy_with_offset(dtype, shape, offset): + import numpy as np + m = 3 + n = 4 + x = ti.Matrix.field(dtype=dtype,m=m,n=n, 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, + } + numpy_shape = ((shape,) if isinstance(shape, int) else shape) + (n, m) + arr = np.ones(numpy_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)