def ext_arr_to_matrix(arr: ext_arr(), mat: template(), as_vector: template()): for I in grouped(mat): for p in static(range(mat.n)): for q in static(range(mat.m)): if static(as_vector): mat[I][p] = arr[I, p] else: mat[I][p, q] = arr[I, p, q]
def matrix_to_ext_arr(mat: template(), arr: ndarray_type.ndarray(), as_vector: template()): for I in grouped(mat): for p in static(range(mat.n)): for q in static(range(mat.m)): if static(as_vector): arr[I, p] = mat[I][p] else: arr[I, p, q] = mat[I][p, q]
def ndarray_matrix_to_ext_arr(ndarray: any_arr(), arr: ext_arr(), as_vector: template()): for I in grouped(ndarray): for p in static(range(ndarray[I].n)): for q in static(range(ndarray[I].m)): if static(as_vector): arr[I, p] = ndarray[I][p] else: arr[I, p, q] = ndarray[I][p, q]
def ndarray_matrix_to_ext_arr(ndarray: ndarray_type.ndarray(), arr: ndarray_type.ndarray(), layout_is_aos: template(), as_vector: template()): for I in grouped(ndarray): for p in static(range(ndarray[I].n)): for q in static(range(ndarray[I].m)): if static(as_vector): if static(layout_is_aos): arr[I, p] = ndarray[I][p] else: arr[p, I] = ndarray[I][p] else: if static(layout_is_aos): arr[I, p, q] = ndarray[I][p, q] else: arr[p, q, I] = ndarray[I][p, q]
def build_static_for(ctx, node, is_grouped): if is_grouped: assert len(node.iter.args[0].args) == 1 ndrange_arg = build_stmt(ctx, node.iter.args[0].args[0]) if not isinstance(ndrange_arg, ndrange): raise TaichiSyntaxError( "Only 'ti.ndrange' is allowed in 'ti.static(ti.grouped(...))'." ) targets = ASTTransformer.get_for_loop_targets(node) if len(targets) != 1: raise TaichiSyntaxError( f"Group for should have 1 loop target, found {len(targets)}" ) target = targets[0] for value in impl.grouped(ndrange_arg): with ctx.variable_scope_guard(): ctx.create_variable(target, value) build_stmts(ctx, node.body) status = ctx.loop_status() if status == LoopStatus.Break: break elif status == LoopStatus.Continue: ctx.set_loop_status(LoopStatus.Normal) else: build_stmt(ctx, node.iter) targets = ASTTransformer.get_for_loop_targets(node) for target_values in node.iter.ptr: if not isinstance( target_values, collections.abc.Sequence) or len(targets) == 1: target_values = [target_values] with ctx.variable_scope_guard(): for target, target_value in zip(targets, target_values): ctx.create_variable(target, target_value) build_stmts(ctx, node.body) status = ctx.loop_status() if status == LoopStatus.Break: break elif status == LoopStatus.Continue: ctx.set_loop_status(LoopStatus.Normal) return None
def fill_tensor(tensor: template(), val: template()): for I in grouped(tensor): tensor[I] = val
def ext_arr_to_ndarray(arr: ndarray_type.ndarray(), ndarray: ndarray_type.ndarray()): for I in grouped(ndarray): ndarray[I] = arr[I]
def ndarray_to_ndarray(ndarray: any_arr(), other: any_arr()): for I in grouped(ndarray): ndarray[I] = other[I]
def ext_arr_to_tensor(arr: ext_arr(), tensor: template()): for I in grouped(tensor): tensor[I] = arr[I]
def ndarray_to_ndarray(ndarray: ndarray_type.ndarray(), other: ndarray_type.ndarray()): for I in grouped(ndarray): ndarray[I] = other[I]
def tensor_to_ext_arr(tensor: template(), arr: ndarray_type.ndarray()): for I in grouped(tensor): arr[I] = tensor[I]
def ext_arr_to_ndarray(arr: ext_arr(), ndarray: any_arr()): for I in grouped(ndarray): ndarray[I] = arr[I]
def snode_deactivate(b: template()): for I in grouped(b): deactivate(b, I)
def fill_ndarray(ndarray: ndarray_type.ndarray(), val: template()): for I in grouped(ndarray): ndarray[I] = val
def clear_gradients(_vars: template()): for I in grouped(ScalarField(Expr(_vars[0]))): for s in static(_vars): ScalarField(Expr(s))[I] = 0
def ndarray_to_ext_arr(ndarray: any_arr(), arr: ext_arr()): for I in grouped(ndarray): arr[I] = ndarray[I]
def tensor_to_ext_arr(tensor: template(), arr: ext_arr()): for I in grouped(tensor): arr[I] = tensor[I]
def fill_ndarray_matrix(ndarray: any_arr(), val: template()): for I in grouped(ndarray): ndarray[I].fill(val)
def fill_ndarray(ndarray: any_arr(), val: template()): for I in grouped(ndarray): ndarray[I] = val
def fill_matrix(mat: template(), vals: template()): for I in grouped(mat): for p in static(range(mat.n)): for q in static(range(mat.m)): mat[I][p, q] = vals[p][q]
def tensor_to_image(tensor: template(), arr: ext_arr()): for I in grouped(tensor): t = ops.cast(tensor[I], f32) arr[I, 0] = t arr[I, 1] = t arr[I, 2] = t
def snode_deactivate_dynamic(b: template()): for I in grouped(b.parent()): deactivate(b, I)
def vector_to_image(mat: template(), arr: ndarray_type.ndarray()): for I in grouped(mat): for p in static(range(mat.n)): arr[I, p] = ops.cast(mat[I][p], f32) if static(mat.n <= 2): arr[I, 2] = 0
def fill_ndarray_matrix(ndarray: ndarray_type.ndarray(), val: template()): for I in grouped(ndarray): ndarray[I].fill(val)
def tensor_to_tensor(tensor: template(), other: template()): for I in grouped(tensor): tensor[I] = other[I]
def ndarray_to_ext_arr(ndarray: ndarray_type.ndarray(), arr: ndarray_type.ndarray()): for I in grouped(ndarray): arr[I] = ndarray[I]
def ext_arr_to_tensor(arr: ndarray_type.ndarray(), tensor: template()): for I in grouped(tensor): tensor[I] = arr[I]
def tensor_to_image(tensor: template(), arr: ndarray_type.ndarray()): for I in grouped(tensor): t = ops.cast(tensor[I], f32) arr[I, 0] = t arr[I, 1] = t arr[I, 2] = t
def field_fill_taichi_scope(F: template(), val: template()): for I in grouped(F): F[I] = val