def field(dtype, shape=None, name="", offset=None, needs_grad=False):
"""Defines a Taichi field
A Taichi field can be viewed as an abstract N-dimensional array, hiding away
the complexity of how its underlying :class:`~taichi.lang.snode.SNode` are
actually defined. The data in a Taichi field can be directly accessed by
a Taichi :func:`~taichi.lang.kernel_impl.kernel`.
See also https://docs.taichi.graphics/lang/articles/basic/field
Args:
dtype (DataType): data type of the field.
shape (Union[int, tuple[int]], optional): shape of the field
name (str, optional): name of the field
offset (Union[int, tuple[int]], optional): offset of the field domain
needs_grad (bool, optional): whether this field participates in autodiff
and thus needs an adjoint field to store the gradients.
Example:
The code below shows how a Taichi field can be declared and defined::
>>> x1 = ti.field(ti.f32, shape=(16, 8))
>>>
>>> # Equivalently
>>> x2 = ti.field(ti.f32)
>>> ti.root.dense(ti.ij, shape=(16, 8)).place(x2)
"""
_taichi_skip_traceback = 1
if isinstance(shape, numbers.Number):
shape = (shape, )
if isinstance(offset, numbers.Number):
offset = (offset, )
if shape is not None and offset is not None:
assert len(shape) == len(
offset
), f'The dimensionality of shape and offset must be the same ({len(shape)} != {len(offset)})'
assert (offset is None or shape
is not None), 'The shape cannot be None when offset is being set'
del _taichi_skip_traceback
x, x_grad = create_field_member(dtype, name)
x, x_grad = ScalarField(x), ScalarField(x_grad)
x.set_grad(x_grad)
if shape is not None:
dim = len(shape)
root.dense(index_nd(dim), shape).place(x, offset=offset)
if needs_grad:
root.dense(index_nd(dim), shape).place(x_grad)
return x
def get_scalar_field(self, *indices):
"""Creates a ScalarField using a specific field member. Only used for quant.
Args:
indices (Tuple[Int]): Specified indices of the field member.
Returns:
ScalarField: The result ScalarField.
"""
assert len(indices) in [1, 2]
i = indices[0]
j = 0 if len(indices) == 1 else indices[1]
return ScalarField(self.vars[i * self.m + j])
def clear_gradients(_vars: template()):
for I in grouped(ScalarField(Expr(_vars[0]))):
for s in static(_vars):
ScalarField(Expr(s))[I] = 0
def field(cls,
n,
m,
dtype,
shape=None,
name="",
offset=None,
needs_grad=False,
layout=Layout.AOS):
"""Construct a data container to hold all elements of the Matrix.
Args:
n (int): The desired number of rows of the Matrix.
m (int): The desired number of columns of the Matrix.
dtype (DataType, optional): The desired data type of the Matrix.
shape (Union[int, tuple of int], optional): The desired shape of the Matrix.
name (string, optional): The custom name of the field.
offset (Union[int, tuple of int], optional): The coordinate offset of all elements in a field.
needs_grad (bool, optional): Whether the Matrix need gradients.
layout (Layout, optional): The field layout, i.e., Array Of Structure (AOS) or Structure Of Array (SOA).
Returns:
:class:`~taichi.lang.matrix.Matrix`: A :class:`~taichi.lang.matrix.Matrix` instance serves as the data container.
"""
entries = []
if isinstance(dtype, (list, tuple, np.ndarray)):
# set different dtype for each element in Matrix
# see #2135
if m == 1:
assert len(np.shape(dtype)) == 1 and len(
dtype
) == n, f'Please set correct dtype list for Vector. The shape of dtype list should be ({n}, ) instead of {np.shape(dtype)}'
for i in range(n):
entries.append(
impl.create_field_member(dtype[i], name=name))
else:
assert len(np.shape(dtype)) == 2 and len(dtype) == n and len(
dtype[0]
) == m, f'Please set correct dtype list for Matrix. The shape of dtype list should be ({n}, {m}) instead of {np.shape(dtype)}'
for i in range(n):
for j in range(m):
entries.append(
impl.create_field_member(dtype[i][j], name=name))
else:
for _ in range(n * m):
entries.append(impl.create_field_member(dtype, name=name))
entries, entries_grad = zip(*entries)
entries, entries_grad = MatrixField(entries, n, m), MatrixField(
entries_grad, n, m)
entries.set_grad(entries_grad)
if shape is None:
assert offset is None, "shape cannot be None when offset is being set"
if shape is not None:
if isinstance(shape, numbers.Number):
shape = (shape, )
if isinstance(offset, numbers.Number):
offset = (offset, )
if offset is not None:
assert len(shape) == len(
offset
), f'The dimensionality of shape and offset must be the same ({len(shape)} != {len(offset)})'
dim = len(shape)
if layout == Layout.SOA:
for e in entries.get_field_members():
ti.root.dense(impl.index_nd(dim),
shape).place(ScalarField(e), offset=offset)
if needs_grad:
for e in entries_grad.get_field_members():
ti.root.dense(impl.index_nd(dim),
shape).place(ScalarField(e),
offset=offset)
else:
ti.root.dense(impl.index_nd(dim), shape).place(entries,
offset=offset)
if needs_grad:
ti.root.dense(impl.index_nd(dim),
shape).place(entries_grad, offset=offset)
return entries