def field(dtype, shape=None, offset=None, needs_grad=False):
_taichi_skip_traceback = 1
dtype = cook_dtype(dtype)
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 not None and shape is None
) == False, f'The shape cannot be None when offset is being set'
if get_runtime().materialized:
raise RuntimeError(
"No new variables can be declared after materialization, i.e. kernel invocations "
"or Python-scope field accesses. I.e., data layouts must be specified before "
"any computation. Try appending ti.init() or ti.reset() "
"right after 'import taichi as ti' if you are using Jupyter notebook or Blender."
)
del _taichi_skip_traceback
# primal
x = Expr(_ti_core.make_id_expr(""))
x.declaration_tb = get_traceback(stacklevel=2)
x.ptr = _ti_core.global_new(x.ptr, dtype)
x.ptr.set_is_primal(True)
pytaichi.global_vars.append(x)
if _ti_core.needs_grad(dtype):
# adjoint
x_grad = Expr(_ti_core.make_id_expr(""))
x_grad.ptr = _ti_core.global_new(x_grad.ptr, dtype)
x_grad.ptr.set_is_primal(False)
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 create_field_member(dtype, name, needs_grad, needs_dual):
dtype = cook_dtype(dtype)
# primal
prog = get_runtime().prog
if prog is None:
raise TaichiRuntimeError(
"Cannont create field, maybe you forgot to call `ti.init()` first?"
)
x = Expr(prog.make_id_expr(""))
x.declaration_tb = get_traceback(stacklevel=4)
x.ptr = _ti_core.global_new(x.ptr, dtype)
x.ptr.set_name(name)
x.ptr.set_is_primal(True)
pytaichi.global_vars.append(x)
x_grad = None
x_dual = None
if _ti_core.is_real(dtype):
# adjoint
x_grad = Expr(get_runtime().prog.make_id_expr(""))
x_grad.declaration_tb = get_traceback(stacklevel=4)
x_grad.ptr = _ti_core.global_new(x_grad.ptr, dtype)
x_grad.ptr.set_name(name + ".grad")
x_grad.ptr.set_is_primal(False)
x.ptr.set_adjoint(x_grad.ptr)
if needs_grad:
pytaichi.grad_vars.append(x_grad)
# dual
x_dual = Expr(get_runtime().prog.make_id_expr(""))
x_dual.ptr = _ti_core.global_new(x_dual.ptr, dtype)
x_dual.ptr.set_name(name + ".dual")
x_dual.ptr.set_is_primal(False)
x.ptr.set_dual(x_dual.ptr)
if needs_dual:
pytaichi.dual_vars.append(x_dual)
elif needs_grad or needs_dual:
raise TaichiRuntimeError(
f'{dtype} is not supported for field with `needs_grad=True` or `needs_dual=True`.'
)
return x, x_grad, x_dual
def field(dtype, shape=None, name="", offset=None, needs_grad=False):
_taichi_skip_traceback = 1
dtype = cook_dtype(dtype)
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 not None and shape is None
) == False, f'The shape cannot be None when offset is being set'
del _taichi_skip_traceback
# primal
x = Expr(_ti_core.make_id_expr(""))
x.declaration_tb = get_traceback(stacklevel=2)
x.ptr = _ti_core.global_new(x.ptr, dtype)
x.ptr.set_name(name)
x.ptr.set_is_primal(True)
pytaichi.global_vars.append(x)
if _ti_core.needs_grad(dtype):
# adjoint
x_grad = Expr(_ti_core.make_id_expr(""))
x_grad.ptr = _ti_core.global_new(x_grad.ptr, dtype)
x_grad.ptr.set_name(name + ".grad")
x_grad.ptr.set_is_primal(False)
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 create_field_member(dtype, name):
dtype = cook_dtype(dtype)
# primal
x = Expr(_ti_core.make_id_expr(""))
x.declaration_tb = get_traceback(stacklevel=4)
x.ptr = _ti_core.global_new(x.ptr, dtype)
x.ptr.set_name(name)
x.ptr.set_is_primal(True)
pytaichi.global_vars.append(x)
x_grad = None
if _ti_core.needs_grad(dtype):
# adjoint
x_grad = Expr(_ti_core.make_id_expr(""))
x_grad.ptr = _ti_core.global_new(x_grad.ptr, dtype)
x_grad.ptr.set_name(name + ".grad")
x_grad.ptr.set_is_primal(False)
x.ptr.set_grad(x_grad.ptr)
return x, x_grad
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/docs/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
dtype = cook_dtype(dtype)
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 not None and shape is None
) == False, f'The shape cannot be None when offset is being set'
del _taichi_skip_traceback
# primal
x = Expr(_ti_core.make_id_expr(""))
x.declaration_tb = get_traceback(stacklevel=2)
x.ptr = _ti_core.global_new(x.ptr, dtype)
x.ptr.set_name(name)
x.ptr.set_is_primal(True)
pytaichi.global_vars.append(x)
if _ti_core.needs_grad(dtype):
# adjoint
x_grad = Expr(_ti_core.make_id_expr(""))
x_grad.ptr = _ti_core.global_new(x_grad.ptr, dtype)
x_grad.ptr.set_name(name + ".grad")
x_grad.ptr.set_is_primal(False)
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