gstaichi 2.1.1rc3__cp313-cp313-macosx_11_0_arm64.whl

gstaichi/lang/matrix_ops.py

@@ -0,0 +1,341 @@
+# type: ignore
+
+import gstaichi.lang.ops as ops_mod
+from gstaichi.lang.impl import static
+from gstaichi.lang.kernel_impl import func, pyfunc
+from gstaichi.lang.matrix import Matrix, Vector
+from gstaichi.lang.matrix_ops_utils import (
+    arg_at,
+    arg_foreach_check,
+    assert_list,
+    assert_tensor,
+    assert_vector,
+    check_matmul,
+    check_transpose,
+    dim_lt,
+    is_int_const,
+    preconditions,
+    same_shapes,
+    square_matrix,
+)
+from gstaichi.types.annotations import template
+
+
+@preconditions(arg_at(0, assert_tensor))
+@pyfunc
+def _reduce(mat, fun: template()):
+    shape = static(mat.get_shape())
+    if static(len(shape) == 1):
+        result = mat[0]
+        for i in static(range(1, shape[0])):
+            result = fun(result, mat[i])
+        return result
+    result = mat[0, 0]
+    for i in static(range(shape[0])):
+        for j in static(range(shape[1])):
+            if static(i != 0 or j != 0):
+                result = fun(result, mat[i, j])
+    return result
+
+
+@pyfunc
+def _filled_vector(n: template(), dtype: template(), val: template()):
+    return Vector([val for _ in static(range(n))], dtype)
+
+
+@pyfunc
+def _filled_matrix(n: template(), m: template(), dtype: template(), val: template()):
+    return Matrix([[val for _ in static(range(m))] for _ in static(range(n))], dtype)
+
+
+@pyfunc
+def _unit_vector(n: template(), i: template(), dtype: template()):
+    return Vector([i == j for j in static(range(n))], dtype)
+
+
+@pyfunc
+def _identity_matrix(n: template(), dtype: template()):
+    return Matrix([[i == j for j in static(range(n))] for i in static(range(n))], dtype)
+
+
+@preconditions(
+    arg_at(0, lambda xs: same_shapes(*xs)),
+    arg_foreach_check(
+        0,
+        fns=[assert_vector(), assert_list],
+        logic="or",
+        msg="Cols/rows must be a list of lists, or a list of vectors",
+    ),
+)
+@pyfunc
+def rows(rows):  # pylint: disable=W0621
+    return Matrix([[x for x in row] for row in rows])
+
+
+@pyfunc
+def cols(cols):  # pylint: disable=W0621
+    return rows(cols).transpose()
+
+
+@pyfunc
+def E(mat: template(), x: template(), y: template(), n: template()):
+    return mat[x % n, y % n]
+
+
+@preconditions(square_matrix, dim_lt(0, 5))
+@pyfunc
+def determinant(mat):
+    shape = static(mat.get_shape())
+    if static(shape[0] == 1):
+        return mat[0, 0]
+    if static(shape[0] == 2):
+        return mat[0, 0] * mat[1, 1] - mat[0, 1] * mat[1, 0]
+    if static(shape[0] == 3):
+        return (
+            mat[0, 0] * (mat[1, 1] * mat[2, 2] - mat[2, 1] * mat[1, 2])
+            - mat[1, 0] * (mat[0, 1] * mat[2, 2] - mat[2, 1] * mat[0, 2])
+            + mat[2, 0] * (mat[0, 1] * mat[1, 2] - mat[1, 1] * mat[0, 2])
+        )
+    if static(shape[0] == 4):
+        det = mat[0, 0] * 0  # keep type
+        for i in static(range(4)):
+            det = det + (-1) ** i * (
+                mat[i, 0]
+                * (
+                    E(mat, i + 1, 1, 4)
+                    * (E(mat, i + 2, 2, 4) * E(mat, i + 3, 3, 4) - E(mat, i + 3, 2, 4) * E(mat, i + 2, 3, 4))
+                    - E(mat, i + 2, 1, 4)
+                    * (E(mat, i + 1, 2, 4) * E(mat, i + 3, 3, 4) - E(mat, i + 3, 2, 4) * E(mat, i + 1, 3, 4))
+                    + E(mat, i + 3, 1, 4)
+                    * (E(mat, i + 1, 2, 4) * E(mat, i + 2, 3, 4) - E(mat, i + 2, 2, 4) * E(mat, i + 1, 3, 4))
+                )
+            )
+        return det
+    # unreachable
+    return None
+
+
+@preconditions(square_matrix, dim_lt(0, 5))
+@pyfunc
+def inverse(mat):
+    shape = static(mat.get_shape())
+    if static(shape[0] == 1):
+        return Matrix([[1.0 / mat[0, 0]]])
+    inv_determinant = 1.0 / determinant(mat)
+    if static(shape[0] == 2):
+        return inv_determinant * Matrix([[mat[1, 1], -mat[0, 1]], [-mat[1, 0], mat[0, 0]]])
+    if static(shape[0] == 3):
+        return inv_determinant * Matrix(
+            [
+                [
+                    E(mat, i + 1, j + 1, 3) * E(mat, i + 2, j + 2, 3)
+                    - E(mat, i + 2, j + 1, 3) * E(mat, i + 1, j + 2, 3)
+                    for i in static(range(3))
+                ]
+                for j in static(range(3))
+            ]
+        )
+    if static(shape[0] == 4):
+        return inv_determinant * Matrix(
+            [
+                [
+                    (-1) ** (i + j)
+                    * (
+                        (
+                            E(mat, i + 1, j + 1, 4)
+                            * (
+                                E(mat, i + 2, j + 2, 4) * E(mat, i + 3, j + 3, 4)
+                                - E(mat, i + 3, j + 2, 4) * E(mat, i + 2, j + 3, 4)
+                            )
+                            - E(mat, i + 2, j + 1, 4)
+                            * (
+                                E(mat, i + 1, j + 2, 4) * E(mat, i + 3, j + 3, 4)
+                                - E(mat, i + 3, j + 2, 4) * E(mat, i + 1, j + 3, 4)
+                            )
+                            + E(mat, i + 3, j + 1, 4)
+                            * (
+                                E(mat, i + 1, j + 2, 4) * E(mat, i + 2, j + 3, 4)
+                                - E(mat, i + 2, j + 2, 4) * E(mat, i + 1, j + 3, 4)
+                            )
+                        )
+                    )
+                    for i in static(range(4))
+                ]
+                for j in static(range(4))
+            ]
+        )
+    # unreachable
+    return None
+
+
+@preconditions(check_transpose)
+@pyfunc
+def transpose(mat):
+    shape = static(mat.get_shape())
+    return Matrix([[mat[i, j] for i in static(range(shape[0]))] for j in static(range(shape[1]))])
+
+
+@preconditions(arg_at(0, is_int_const))
+@pyfunc
+def diag(dim: template(), val: template()):
+    return Matrix([[val if i == j else 0 for j in static(range(dim))] for i in static(range(dim))])
+
+
+@preconditions(assert_tensor)
+@pyfunc
+def sum(mat):  # pylint: disable=W0622
+    return _reduce(mat, ops_mod.add)
+
+
+@preconditions(assert_tensor)
+@pyfunc
+def norm_sqr(mat):
+    return sum(mat * mat)
+
+
+@preconditions(arg_at(0, assert_tensor))
+@pyfunc
+def norm(mat, eps=0.0):
+    return ops_mod.sqrt(norm_sqr(mat) + eps)
+
+
+@preconditions(arg_at(0, assert_tensor))
+@pyfunc
+def norm_inv(mat, eps=0.0):
+    return ops_mod.rsqrt(norm_sqr(mat) + eps)
+
+
+@preconditions(arg_at(0, assert_vector()))
+@pyfunc
+def normalized(vec, eps=0.0):
+    invlen = 1 / (norm(vec) + eps)
+    return invlen * vec
+
+
+@preconditions(assert_tensor)
+@pyfunc
+def any(mat):  # pylint: disable=W0622
+    return _reduce(mat != 0, ops_mod.logical_or) and True
+
+
+@preconditions(assert_tensor)
+@pyfunc
+def all(mat):  # pylint: disable=W0622
+    return _reduce(mat != 0, ops_mod.logical_and) and True
+
+
+@preconditions(assert_tensor)
+@pyfunc
+def max(mat):  # pylint: disable=W0622
+    return _reduce(mat, ops_mod.max_impl)
+
+
+@preconditions(assert_tensor)
+@pyfunc
+def min(mat):  # pylint: disable=W0622
+    return _reduce(mat, ops_mod.min_impl)
+
+
+@preconditions(square_matrix)
+@pyfunc
+def trace(mat):
+    shape = static(mat.get_shape())
+    result = mat[0, 0]
+    # TODO: get rid of static when
+    # CHI IR Tensor repr is ready stable
+    for i in static(range(1, shape[0])):
+        result = result + mat[i, i]
+    return result
+
+
+@preconditions(arg_at(0, assert_tensor))
+@pyfunc
+def fill(mat: template(), val):
+    shape = static(mat.get_shape())
+    if static(len(shape) == 1):
+        for i in static(range(shape[0])):
+            mat[i] = val
+    else:
+        for i in static(range(shape[0])):
+            for j in static(range(shape[1])):
+                mat[i, j] = val
+
+
+@preconditions(check_matmul)
+@pyfunc
+def _matmul_helper(mat_x, mat_y):
+    shape_x = static(mat_x.get_shape())
+    shape_y = static(mat_y.get_shape())
+    if static(len(shape_x) == 1 and len(shape_y) == 1):
+        return dot(mat_x, mat_y)
+    if static(len(shape_y) == 1):
+        zero_elem = mat_x[0, 0] * mat_y[0] * 0  # for correct return type
+        vec_z = _filled_vector(shape_x[0], None, zero_elem)
+        for i in static(range(shape_x[0])):
+            for j in static(range(shape_x[1])):
+                vec_z[i] = vec_z[i] + mat_x[i, j] * mat_y[j]
+        return vec_z
+    zero_elem = mat_x[0, 0] * mat_y[0, 0] * 0  # for correct return type
+    mat_z = _filled_matrix(shape_x[0], shape_y[1], None, zero_elem)
+    for i in static(range(shape_x[0])):
+        for j in static(range(shape_y[1])):
+            for k in static(range(shape_x[1])):
+                mat_z[i, j] = mat_z[i, j] + mat_x[i, k] * mat_y[k, j]
+    return mat_z
+
+
+@pyfunc
+def matmul(mat_x, mat_y):
+    shape_x = static(mat_x.get_shape())
+    shape_y = static(mat_y.get_shape())
+    if static(len(shape_x) == 1 and len(shape_y) == 2):
+        return _matmul_helper(transpose(mat_y), mat_x)
+    return _matmul_helper(mat_x, mat_y)
+
+
+@preconditions(
+    arg_at(0, assert_vector("lhs for dot is not a vector")),
+    arg_at(1, assert_vector("rhs for dot is not a vector")),
+)
+@pyfunc
+def dot(vec_x, vec_y):
+    return sum(vec_x * vec_y)
+
+
+@preconditions(
+    arg_at(0, assert_vector("lhs for cross is not a vector")),
+    arg_at(1, assert_vector("rhs for cross is not a vector")),
+    same_shapes,
+    arg_at(0, dim_lt(0, 4)),
+)
+@pyfunc
+def cross(vec_x, vec_y):
+    shape = static(vec_x.get_shape())
+    if static(shape[0] == 2):
+        return vec_x[0] * vec_y[1] - vec_x[1] * vec_y[0]
+    if static(shape[0] == 3):
+        return Vector(
+            [
+                vec_x[1] * vec_y[2] - vec_x[2] * vec_y[1],
+                vec_x[2] * vec_y[0] - vec_x[0] * vec_y[2],
+                vec_x[0] * vec_y[1] - vec_x[1] * vec_y[0],
+            ]
+        )
+    return None
+
+
+@preconditions(
+    arg_at(0, assert_vector("lhs for outer_product is not a vector")),
+    arg_at(1, assert_vector("rhs for outer_product is not a vector")),
+)
+@pyfunc
+def outer_product(vec_x, vec_y):
+    shape_x = static(vec_x.get_shape())
+    shape_y = static(vec_y.get_shape())
+    return Matrix([[vec_x[i] * vec_y[j] for j in static(range(shape_y[0]))] for i in static(range(shape_x[0]))])
+
+
+@preconditions(assert_tensor)
+@func
+def cast(mat, dtype: template()):
+    return ops_mod.cast(mat, dtype)

gstaichi/lang/matrix_ops_utils.py

@@ -0,0 +1,190 @@
+# type: ignore
+
+import functools
+
+from gstaichi.lang.exception import GsTaichiCompilationError
+from gstaichi.lang.expr import Expr
+from gstaichi.lang.matrix import Matrix
+
+
+def do_check(checker_fns, *args, **kwargs):
+    for f in checker_fns:
+        ok, msg = f(*args, **kwargs)
+        if not ok:
+            return False, msg
+    return True, None
+
+
+def preconditions(*checker_funcs):
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            ok, msg = do_check(checker_funcs, *args, **kwargs)
+            if not ok:
+                raise GsTaichiCompilationError(msg)
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+def arg_at(indices, *fns):
+    def check(*args, **kwargs):
+        nonlocal indices
+        if isinstance(indices, int):
+            indices = [indices]
+        for i in indices:
+            if i in kwargs:
+                arg = kwargs[i]
+            else:
+                arg = args[i]
+            ok, msg = do_check(fns, arg)
+            if not ok:
+                return False, msg
+        return True, None
+
+    return check
+
+
+def assert_tensor(m, msg="not tensor type: {}"):
+    if isinstance(m, Matrix):
+        return True, None
+    if isinstance(m, Expr) and m.is_tensor():
+        return True, None
+    return False, msg.format(type(m))
+
+
+def assert_vector(msg="expected a vector, got {}"):
+    def check(v):
+        if (isinstance(v, Expr) or isinstance(v, Matrix)) and len(v.get_shape()) == 1:
+            return True, None
+        return False, msg.format(type(v))
+
+    return check
+
+
+def assert_list(x, msg="not a list: {}"):
+    if isinstance(x, list):
+        return True, None
+    return False, msg.format(type(x))
+
+
+def arg_foreach_check(*arg_indices, fns=[], logic="or", msg=None):
+    def check(*args, **kwargs):
+        for i in arg_indices:
+            if i in kwargs:
+                arg = kwargs[i]
+            else:
+                arg = args[i]
+            if logic == "or":
+                for a in arg:
+                    passed = False
+                    for fn in fns:
+                        ok, _ = do_check([fn], a)
+                        if ok:
+                            passed = True
+                            break
+                    if not passed:
+                        return False, msg
+            elif logic == "and":
+                for a in arg:
+                    ok, _ = do_check(fns, a)
+                    if not ok:
+                        return False, msg
+            else:
+                raise ValueError(f"Unknown logic: {logic}")
+        return True, None
+
+    return check
+
+
+def get_list_shape(x):
+    outer_shape = [len(x)]
+    inner_shape = None
+    for element in x:
+        if isinstance(element, list):
+            cur_shape = get_list_shape(element)
+        else:
+            cur_shape = []
+
+        if inner_shape:
+            assert curr_shape == inner_shape
+        else:
+            inner_shape = cur_shape
+
+    return outer_shape + inner_shape
+
+
+def same_shapes(*xs):
+    shapes = []
+    for x in xs:
+        if isinstance(x, Matrix):
+            shapes.append(x.get_shape())
+        elif isinstance(x, list):
+            shapes.append(tuple(get_list_shape(x)))
+        elif isinstance(x, Expr):
+            shapes.append(tuple(x.ptr.get_rvalue_type().shape()))
+        else:
+            return False, f"same_shapes() received an unexpected argument of type: {x}"
+
+    if len(set(shapes)) != 1:
+        return False, f"required shapes to be the same, got shapes {shapes}"
+    return True, None
+
+
+def square_matrix(x):
+    assert_tensor(x)
+    shape = x.get_shape()
+    if len(shape) != 2 or shape[0] != shape[1]:
+        return False, f"expected a square matrix, got shape {shape}"
+    return True, None
+
+
+def dim_lt(dim, limit):
+    def check(x):
+        assert_tensor(x)
+        shape = x.get_shape()
+        return shape[dim] < limit, (f"only dimension < {limit} is supported, got shape {shape}")
+
+    return check
+
+
+def is_int_const(x):
+    if isinstance(x, int):
+        return True, None
+    if isinstance(x, Expr) and x.val_int() is not None:
+        return True, None
+    return False, f"not an integer: {x} of type {type(x).__name__}"
+
+
+def check_matmul(x, y):
+    assert_tensor(x, f"left hand side is not a matrix: {type(x)}")
+    assert_tensor(y, f"right hand side is not a matrix: {type(y)}")
+    x_shape = x.get_shape()
+    y_shape = y.get_shape()
+    if len(x_shape) == 1:
+        if len(y_shape) == 1:
+            return True, None
+        if x_shape[0] != y_shape[0]:
+            return (
+                False,
+                f"dimension mismatch between {x_shape} and {y_shape} for left multiplication",
+            )
+    else:
+        if x_shape[1] != y_shape[0]:
+            return (
+                False,
+                f"dimension mismatch between {x_shape} and {y_shape} for matrix multiplication",
+            )
+    return True, None
+
+
+def check_transpose(x):
+    ok, msg = assert_tensor(x)
+    if ok and len(x.get_shape()) == 1:
+        return (
+            False,
+            "`transpose()` cannot apply to a vector. If you want something like `a @ b.transpose()`, write `a.outer_product(b)` instead.",
+        )
+    return ok, msg