def augassign(self, value, op):
from taichi.lang.impl import call_internal, ti_float
if op == 'Add':
call_internal("insert_triplet", self.ptr, self.i, self.j,
ti_float(value))
elif op == 'Sub':
call_internal("insert_triplet", self.ptr, self.i, self.j,
-ti_float(value))
else:
assert False, f"Only operations '+=' and '-=' are supported on sparse matrices."
def foo(x: ti.i32, y: ti.i32, a: ti.template()):
a[0] = x + y
a[1] = x - y
a[2] = x * y
a[3] = impl.ti_float(x) / y
a[4] = x // y
a[5] = x % y
a[6] = x**y
a[7] = x << y
a[8] = x >> y
a[9] = x | y
a[10] = x ^ y
a[11] = x & y
def build_Call(ctx, node):
is_in_static_scope_prev = ctx.is_in_static_scope
if ASTTransformer.get_decorator(ctx, node) == 'static':
ctx.is_in_static_scope = True
build_stmt(ctx, node.func)
build_stmts(ctx, node.args)
build_stmts(ctx, node.keywords)
ctx.is_in_static_scope = is_in_static_scope_prev
args = []
for arg in node.args:
if isinstance(arg, ast.Starred):
args += arg.ptr
else:
args.append(arg.ptr)
keywords = dict(ChainMap(*[keyword.ptr for keyword in node.keywords]))
if isinstance(node.func, ast.Attribute):
attr_name = node.func.attr
if attr_name == 'format' and isinstance(node.func.value.ptr, str):
args.insert(0, node.func.value.ptr)
node.ptr = impl.ti_format(*args, **keywords)
else:
node.ptr = node.func.ptr(*args, **keywords)
elif isinstance(node.func, ast.Name):
func_name = node.func.id
if func_name == 'print':
node.ptr = impl.ti_print(*args, **keywords)
elif func_name == 'min':
node.ptr = ti_ops.ti_min(*args, **keywords)
elif func_name == 'max':
node.ptr = ti_ops.ti_max(*args, **keywords)
elif func_name == 'int':
node.ptr = impl.ti_int(*args, **keywords)
elif func_name == 'float':
node.ptr = impl.ti_float(*args, **keywords)
elif func_name == 'any':
node.ptr = ti_ops.ti_any(*args, **keywords)
elif func_name == 'all':
node.ptr = ti_ops.ti_all(*args, **keywords)
else:
node.ptr = node.func.ptr(*args, **keywords)
else:
node.ptr = node.func.ptr(*args, **keywords)
return node.ptr