Exemplo n.º 1
0
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
Exemplo n.º 2
0
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
Exemplo n.º 3
0
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