Skip to content

Commit

Permalink
[doc] Add info about paddle.Tensor
Browse files Browse the repository at this point in the history
  • Loading branch information
0xzhang committed Apr 29, 2022
1 parent 71e7ffe commit e0b2674
Showing 1 changed file with 36 additions and 27 deletions.
63 changes: 36 additions & 27 deletions docs/lang/articles/basic/external.md
Original file line number Diff line number Diff line change
Expand Up @@ -4,17 +4,12 @@ sidebar_position: 5

# Interacting with external arrays

Although Taichi fields are mainly used in Taichi-scope, in some cases
efficiently manipulating Taichi field data in Python-scope could also be
Although Taichi fields are mainly used in Taichi-scope, in some cases efficiently manipulating Taichi field data in Python-scope could also be
helpful.

We provide various interfaces to copy the data between Taichi fields and
external arrays. External arrays refer to NumPy arrays or PyTorch tensors.
Let's take a look at the most common usage: interacting with NumPy arrays.
We provide various interfaces to copy the data between Taichi fields and external arrays. External arrays refer to NumPy arrays, PyTorch tensors or Paddle Tensors. Let's take a look at the most common usage: interacting with NumPy arrays.

**Export data in Taichi fields to NumPy arrays** via `to_numpy()`. This
allows us to export computation results to other Python packages that
support NumPy, e.g. `matplotlib`.
**Export data in Taichi fields to NumPy arrays** via `to_numpy()`. This allows us to export computation results to other Python packages that support NumPy, e.g. `matplotlib`.

```python {8}
@ti.kernel
Expand All @@ -28,8 +23,7 @@ x_np = x.to_numpy()
print(x_np) # np.array([0, 2, 4, 6])
```

**Import data from NumPy arrays to Taichi fields** via `from_numpy()`.
This allows us to initialize Taichi fields via NumPy arrays:
**Import data from NumPy arrays to Taichi fields** via `from_numpy()`. This allows us to initialize Taichi fields via NumPy arrays:

```python {3}
x = ti.field(ti.f32, 4)
Expand Down Expand Up @@ -59,21 +53,42 @@ print(x[1]) # 7
print(x[2]) # 3
print(x[3]) # 5
```
And Taichi fields also can be **imported from and exported to Paddle tensors**:

```python
@ti.kernel
def my_kernel():
for i in x:
x[i] = i * 2

x = ti.field(ti.f32, 4)
my_kernel()
x_paddle = x.to_paddle()
print(x_paddle) # paddle.Tensor([0, 2, 4, 6])

x.from_numpy(paddle.to_tensor([1, 7, 3, 5]))
print(x[0]) # 1
print(x[1]) # 7
print(x[2]) # 3
print(x[3]) # 5
```

When calling `to_torch()`, specify the PyTorch device where the Taichi field is exported using the `device` argument:

```python
x = ti.field(ti.f32, 4)
x.fill(3.0)
x_torch = x.to_torch(device="cuda:0")
print(x_torch.device) # device(type='cuda', index=0)
```

For Paddle, specify the device by `paddle.CPUPlace()` or `paddle.CUDAPlace(n)` where n is an optional ID, default is 0.

## External array shapes

Shapes of Taichi fields and those of corresponding NumPy arrays or PyTorch tensors are closely
connected via the following rules:
Shapes of Taichi fields and those of corresponding NumPy arrays, PyTorch tensors or Paddle Tensors are closely connected via the following rules:

- For scalar fields, **the shape of NumPy array or PyTorch tensor equals the shape of
the Taichi field**:
- For scalar fields, **the shape of NumPy array, PyTorch tensor or Paddle Tensor equals the shape of the Taichi field**

```python
field = ti.field(ti.i32, shape=(256, 512))
Expand All @@ -85,8 +100,7 @@ array.shape # (256, 512)
field.from_numpy(array) # the input array must be of shape (256, 512)
```

- For vector fields, if the vector is `n`-D, then **the shape of NumPy
array or Pytorch tensor should be** `(*field_shape, vector_n)`:
- For vector fields, if the vector is `n`-D, then **the shape of NumPy array, PyTorch tensor or Paddle Tensor should be** `(*field_shape, vector_n)`:

```python
field = ti.Vector.field(3, ti.i32, shape=(256, 512))
Expand All @@ -99,8 +113,7 @@ array.shape # (256, 512, 3)
field.from_numpy(array) # the input array must be of shape (256, 512, 3)
```

- For matrix fields, if the matrix is `n`-by-`m` (`n x m`), then **the shape of NumPy
array or Pytorch Tensor should be** `(*field_shape, matrix_n, matrix_m)`:
- For matrix fields, if the matrix is `n`-by-`m` (`n x m`), then **the shape of NumPy array, PyTorch tensor or Paddle Tensor should be** `(*field_shape, matrix_n, matrix_m)`:

```python
field = ti.Matrix.field(3, 4, ti.i32, shape=(256, 512))
Expand All @@ -114,8 +127,7 @@ array.shape # (256, 512, 3, 4)
field.from_numpy(array) # the input array must be of shape (256, 512, 3, 4)
```

- For struct fields, the external array will be exported as **a dictionary of NumPy arrays or PyTorch tensors** with keys
being struct member names and values being struct member arrays. Nested structs will be exported as nested dictionaries:
- For struct fields, the external array will be exported 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 will be exported as nested dictionaries:

```python
field = ti.Struct.field({'a': ti.i32, 'b': ti.types.vector(float, 3)} shape=(256, 512))
Expand All @@ -131,8 +143,7 @@ field.from_numpy(array_dict) # the input array must have the same keys as the fi

## Using external arrays as Taichi kernel arguments

Use type hint `ti.types.ndarray()` to pass external arrays as kernel
arguments. For example:
Use type hint `ti.types.ndarray()` to pass external arrays as kernel arguments. For example:

```python {10}
import taichi as ti
Expand Down Expand Up @@ -163,8 +174,7 @@ for i in range(n):
assert a[i, j] == i * j + i + j
```

Note that the elements in an external array must be indexed using a single square bracket.
This contrasts with a Taichi vector or matrix field where field and matrix indices are indexed separately:
Note that the elements in an external array must be indexed using a single square bracket. This contrasts with a Taichi vector or matrix field where field and matrix indices are indexed separately:
```python
@ti.kernel
def copy_vector(x: ti.template(), y: ti.types.ndarray()):
Expand All @@ -174,9 +184,8 @@ def copy_vector(x: ti.template(), y: ti.types.ndarray()):
# y[i][j][k] = x[i, j][k] incorrect
# y[i, j][k] = x[i, j][k] incorrect
```
Also, external arrays in a Taichi kernel are indexed using its **physical memory layout**. For PyTorch users,
this implies that the PyTorch tensor [needs to be made contiguous](https://pytorch.org/docs/stable/generated/torch.Tensor.contiguous.html)
before being passed into a Taichi kernel:
Also, external arrays in a Taichi kernel are indexed using its **physical memory layout**. For PyTorch users, this implies that the PyTorch tensor [needs to be made contiguous](https://pytorch.org/docs/stable/generated/torch.Tensor.contiguous.html) before being passed into a Taichi kernel:

```python
@ti.kernel
def copy_scalar(x: ti.template(), y: ti.types.ndarray()):
Expand Down

0 comments on commit e0b2674

Please sign in to comment.