add conditional import to ensure backward compatibility

ifsheldon · Aug 9, 2022 · 26231ad · 26231ad
1 parent 248b00b
commit 26231ad
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_python 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_python 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