Interacting with External Arrays
This document provides instructions on how to transfer data from external arrays to the Taichi scope and vice versa. For now, the external arrays supported by Taichi are NumPy arrays, PyTorch tensors, and Paddle tensors.
We use NumPy arrays as an example to illustrate the data transfer process because NumPy arrays are the most commonly used external arrays in Taichi. The same steps apply to PyTorch tensors and Paddle tensors.
There are two ways to import a NumPy array arr to the Taichi scope:
Create a Taichi field
f, whose shape and dtype match the shape and dtype ofarr, and callf.from_numpy(arr)to copy the data inarrintof. This approach is preferred when the original array is visited frequently from elsewhere in the Taichi scope (for example, in the case of texture sampling).Pass
arras an argument to a kernel or a Taichi function usingti.types.ndarray()as type hint. The argument is passed by reference without creating a copy ofarr. Thus, any modification to this argument from inside a kernel or Taichi function also changes the original arrayarr. This approach is preferred when the kernel or Taichi function that takes in the argument needs to process the original array (for storage or filtering, for example).
Data transfer between NumPy arrays and Taichi fields
To import data from a NumPy array to a Taichi field, first make sure that the field and the array have the same shape:
x = ti.field(float, shape=(3, 3))
a = np.arange(9).reshape(3, 3).astype(np.int32)
x.from_numpy(a)
print(x)
#[[0 1 2]
# [3 4 5]
# [6 7 8]]
In the example above, the scalar field x and the array a have the same shape (3, 3). This operation would fail if their shapes did not match. Shape matching of a vector or matrix field with a NumPy array is slightly different, which will be discussed in a later section.
The field should also have the same dtype as the array; otherwise, an implicit type casting would occur - see type system.
Conversely, to export the data in x to a NumPy array, call to_numpy():
arr = x.to_numpy()
#array([[0, 1, 2],
# [3, 4, 5],
# [6, 7, 8]], dtype=int32)
Data transfer between PyTorch/Paddle tensors and Taichi fields
Data transfer between a PyTorch tensor and a Taichi field is similar to the NumPy case above: Call from_torch() for data import and to_torch() for data export. But note that to_torch() requires one more argument device, which specifies the PyTorch device:
tensor = x.to_torch(device="cuda:0")
print(tensor.device) # device(type='cuda', index=0)
For Paddle, you need to specify the device by calling paddle.CPUPlace() or paddle.CUDAPlace(n), where n is an optional ID set to 0 by default.
External array shapes
As mentioned before, when transferring data between a ti.field/ti.Vector.field/ti.Matrix.field and a NumPy array, you need to make sure that the shapes of both sides are in alignment. The shape matching rules are summarized as below:
When importing data to or exporting data from a scalar field, ensure that the shape of the corresponding NumPy array, PyTorch tensor, or Paddle tensor equals the shape of the scalar field
field = ti.field(int, shape=(256, 512))
field.shape # (256, 512)
array = field.to_numpy()
array.shape # (256, 512)
field.from_numpy(array) # the input array must be of shape (256, 512)An illustration is shown below:
field.shape[1]=array.shape[1]
(=512)
┌───────────────────────┐
┌ ┌───┬───┬───┬───┬───┬───┐ ┐
│ │ │ │ │ │ │ │ │
│ ├───┼───┼───┼───┼───┼───┤ │
field.shape[0]=array.shape[0] │ │ │ │ │ │ │ │ │
(=256) │ ├───┼───┼───┼───┼───┼───┤ │
│ │ │ │ │ │ │ │ │
└ └───┴───┴───┴───┴───┴───┘ ┘When importing data to or exporting data from an
n-dimensional vector field, ensure that the shape of the corresponding NumPy array, PyTorch tensor, or Paddle tensor is set to(*field_shape, n):field = ti.Vector.field(3, int, shape=(256, 512))
field.shape # (256, 512)
field.n # 3
array = field.to_numpy()
array.shape # (256, 512, 3)
field.from_numpy(array) # the input array must in the shape (256, 512, 3)An illustration is shown below:
field.shape[1]=array.shape[1]
(=512)
┌─────────────────────────────┐
┌ ┌─────────┬─────────┬─────────┐ ┐
│ │[*, *, *]│[*, *, *]│[*, *, *]│ │
│ ├─────────┼─────────┼─────────┤ │
field.shape[0]=array.shape[0] │ │[*, *, *]│[*, *, *]│[*, *, *]│ │ [*, *, *]
(=256) │ ├─────────┼─────────┼─────────┤ │ └───────┘
│ │[*, *, *]│[*, *, *]│[*, *, *]│ │ n=array.shape[2]=3
└ └─────────┴─────────┴─────────┘ ┘When importing data to or exporting data from an
n-by-m(n x m) matrix field, ensure that the shape of the corresponding NumPy array, PyTorch tensor, or Paddle tensor is set to(*field_shape, n, m):field = ti.Matrix.field(3, 4, ti.i32, shape=(256, 512))
field.shape # (256, 512)
field.n # 3
field.m # 4
array = field.to_numpy()
array.shape # (256, 512, 3, 4)
field.from_numpy(array) # the input array must be of shape (256, 512, 3, 4)When importing data to a struct field, export the data of the corresponding external array as a dictionary of NumPy arrays, PyTorch tensors, or Paddle tensors with keys being struct member names and values being struct member arrays. Nested structs are exported as nested dictionaries:
field = ti.Struct.field({'a': ti.i32, 'b': ti.types.vector(3, float)}, shape=(256, 512))
field.shape # (256, 512)
array_dict = field.to_numpy()
array_dict.keys() # dict_keys(['a', 'b'])
array_dict['a'].shape # (256, 512)
array_dict['b'].shape # (256, 512, 3)
field.from_numpy(array_dict) # the input array must have the same keys as the field
Using external arrays as Taichi kernel arguments
Use type hint ti.types.ndarray() to pass external arrays as kernel arguments. For example:
import taichi as ti
import numpy as np
ti.init()
a = np.zeros((5, 5))
@ti.kernel
def test(a: ti.types.ndarray()):
for i in range(a.shape[0]): # a parallel for loop
for j in range(a.shape[1]):
a[i, j] = i + j
test()
print(a)
This is an entry-level example to show you how to call ti.types.ndarray(). We now illustrate a more advanced usage of this method.
Assume that a and b are both 2D arrays of the same shape and dtype. For each cell (i, j) in a, we want to calculate the difference between its value and the average of its four neighboring cells while storing the result in the corresponding cell in b. In this case, cells on the boundary, which are cells with fewer than four neighbors, are ruled out for simplicity. This operation is usually denoted as the discrete Laplace operator:
b[i, j] = a[i, j] - (a[i-1, j] + a[i, j-1] + a[i+1, j] + a[i, j+1]) / 4
Such an operation is usually very slow even with NumPy's vectorization, as the following code snippet shows:
b[1:-1, 1:-1] += ( a[ :-2, 1:-1] +
a[1:-1, :-2] + a[1:-1, 2:] +
a[2: , 1:-1])
But Taichi can meet the same purpose in one parallel for loop only:
@ti.kernel
def test(a: ti.types.ndarray(), b: ti.types.ndarray()): # assume a, b have the same shape
H, W = a.shape[0], a.shape[1]
for i, j in ti.ndrange(H, W): # one parallel for loop
if 0 < i < H - 1 and 0 < j < W - 1:
b[i, j] = a[i, j] - (a[i-1, j] + a[i, j-1] + a[i+1, j] + a[i, j+1]) / 4
This code snippet is more readable than the NumPy version above and runs way faster even on the CPU backend.
Note that the elements in an external array must be indexed using a single square bracket. This contrasts with a Taichi vector field or matrix field where field members and elements are indexed separately:
x = ti.Vector.field(3, float, shape=(5, 5))
y = np.random.random((5, 5, 3))
@ti.kernel
def copy_vector(x: ti.template(), y: ti.types.ndarray()):
for i, j in ti.ndrange(5, 5):
for k in ti.static(range(3)):
y[i, j, k] = x[i, j][k] # correct
# y[i][j][k] = x[i, j][k] incorrect
# y[i, j][k] = x[i, j][k] incorrect
In addition, external arrays in a Taichi kernel are indexed using their physical memory layout. For PyTorch users, this means that a PyTorch tensor needs to be made contiguous before being passed into a Taichi kernel:
x = ti.field(dtype=int, shape=(3, 3))
y = torch.Tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
y = y.T # Transposing the tensor returns a view of the tensor which is not contiguous
@ti.kernel
def copy_scalar(x: ti.template(), y: ti.types.ndarray()):
for i, j in x:
y[i, j] = x[i, j]
copy(x, y) # error!
copy(x, y.clone()) # correct
copy(x, y.contiguous()) # correct