diff --git a/python/taichi/lang/__init__.py b/python/taichi/lang/__init__.py index 301a49cf56261..0049f9b67f221 100644 --- a/python/taichi/lang/__init__.py +++ b/python/taichi/lang/__init__.py @@ -23,13 +23,12 @@ from taichi.lang.ops import * from taichi.lang.quant_impl import quant from taichi.lang.runtime_ops import async_flush, sync -from taichi.lang.sparse_matrix import SparseMatrix, SparseMatrixBuilder -from taichi.lang.sparse_solver import SparseSolver from taichi.lang.struct import Struct from taichi.lang.type_factory_impl import type_factory from taichi.lang.util import (has_pytorch, is_taichi_class, python_scope, taichi_scope, to_numpy_type, to_pytorch_type, to_taichi_type) +from taichi.linalg import SparseMatrix, SparseMatrixBuilder, SparseSolver from taichi.misc.util import deprecated from taichi.profiler import KernelProfiler, get_default_kernel_profiler from taichi.snode.fields_builder import FieldsBuilder diff --git a/python/taichi/lang/kernel_arguments.py b/python/taichi/lang/kernel_arguments.py index 016f42cee2981..864d821f21a12 100644 --- a/python/taichi/lang/kernel_arguments.py +++ b/python/taichi/lang/kernel_arguments.py @@ -2,8 +2,8 @@ from taichi.lang.any_array import AnyArray from taichi.lang.enums import Layout from taichi.lang.expr import Expr -from taichi.lang.sparse_matrix import SparseMatrixBuilder from taichi.lang.util import cook_dtype +from taichi.linalg import SparseMatrixBuilder from taichi.type.primitive_types import u64 diff --git a/python/taichi/linalg/__init__.py b/python/taichi/linalg/__init__.py new file mode 100644 index 0000000000000..2bf35da5974fc --- /dev/null +++ b/python/taichi/linalg/__init__.py @@ -0,0 +1,2 @@ +from taichi.linalg.sparse_matrix import SparseMatrix, SparseMatrixBuilder +from taichi.linalg.sparse_solver import SparseSolver diff --git a/python/taichi/lang/sparse_matrix.py b/python/taichi/linalg/sparse_matrix.py similarity index 100% rename from python/taichi/lang/sparse_matrix.py rename to python/taichi/linalg/sparse_matrix.py diff --git a/python/taichi/lang/sparse_solver.py b/python/taichi/linalg/sparse_solver.py similarity index 93% rename from python/taichi/lang/sparse_solver.py rename to python/taichi/linalg/sparse_solver.py index 2ef9cdce75483..bcad5e4553a0a 100644 --- a/python/taichi/lang/sparse_solver.py +++ b/python/taichi/linalg/sparse_solver.py @@ -1,12 +1,11 @@ import numpy as np from taichi.core.util import ti_core as _ti_core -from taichi.lang import Field -from taichi.lang.impl import get_runtime -from taichi.lang.sparse_matrix import SparseMatrix +from taichi.linalg import SparseMatrix class SparseSolver: def __init__(self, solver_type="LLT", ordering="AMD"): + from taichi.lang.impl import get_runtime solver_type_list = ["LLT", "LDLT", "LU"] solver_ordering = ['AMD', 'COLAMD'] if solver_type in solver_type_list and ordering in solver_ordering: @@ -39,6 +38,7 @@ def factorize(self, sparse_matrix): self.type_assert(sparse_matrix) def solve(self, b): + from taichi.lang import Field if isinstance(b, Field): return self.solver.solve(b.to_numpy()) elif isinstance(b, np.ndarray):