Merge pull request #14 from ailzhang/v1.1.0-fix

Compatibility changes for upcoming Taichi v1.1.0
ifsheldon · Aug 10, 2022 · 4039bf4 · 4039bf4
2 parents 6736b11 + 26231ad
commit 4039bf4
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/src/stannum/tube.py b/src/stannum/tube.py
@@ -2,8 +2,7 @@
 import torch
 from taichi.lang.impl import axes
 from typing import Optional, Callable, Union, Tuple, List, Iterable, Dict, Any
-from .utils import to_taichi_type, need_auto_clearing_fields
-from taichi._lib.core.taichi_core import DataType as TiDataType
+from .utils import to_taichi_type, need_auto_clearing_fields, TiDataType
 from taichi.lang.field import ScalarField
 from taichi.lang.matrix import MatrixField
 from functools import partial

diff --git a/src/stannum/utils.py b/src/stannum/utils.py
@@ -5,7 +5,11 @@
 from taichi.lang.kernel_impl import Kernel
 from taichi.lang.matrix import MatrixField
 from taichi.lang.field import ScalarField
-from taichi._lib.core.taichi_core import DataType as TiDataType
+
+if __ti_version >= (1, 1, 0):
+    from taichi._lib.core.taichi_python import DataType as TiDataType
+else:
+    from taichi._lib.core.taichi_core import DataType as TiDataType
 import torch
 import taichi as ti