fix TensorPy for mindspore 2.5 (#1961)

mindspore-lab · Feb 19, 2025 · d859a1b · d859a1b
1 parent ab20e2e
commit d859a1b
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 4 deletions.
diff --git a/mindnlp/accelerate/big_modeling.py b/mindnlp/accelerate/big_modeling.py
@@ -1,6 +1,10 @@
 """big modeling"""
 from contextlib import contextmanager
-from mindspore._c_expression import Tensor as Tensor_ # pylint: disable=no-name-in-module
+try:
+    from mindspore._c_expression import TensorPy as Tensor_ # pylint: disable=no-name-in-module
+except:
+    from mindspore._c_expression import Tensor as Tensor_ # pylint: disable=no-name-in-module
+
 from mindnlp.utils.testing_utils import parse_flag_from_env
 from mindnlp.core import nn
 

diff --git a/mindnlp/core/ops/creation.py b/mindnlp/core/ops/creation.py
@@ -1,7 +1,12 @@
 """creation ops"""
 import numpy as np
 import mindspore
-from mindspore._c_expression import Tensor as CTensor # pylint: disable=no-name-in-module, import-error
+try:
+    from mindspore._c_expression import TensorPy as CTensor # pylint: disable=no-name-in-module
+except:
+    from mindspore._c_expression import Tensor as CTensor # pylint: disable=no-name-in-module
+
+
 from mindspore import ops
 from mindspore.ops._primitive_cache import _get_cache_prim
 from mindnlp.configs import use_pyboost, ON_ORANGE_PI

diff --git a/mindnlp/parallel/comm_func.py b/mindnlp/parallel/comm_func.py
@@ -1,9 +1,13 @@
 """communication functional api."""
 from mindspore import ops, Tensor
-from mindspore._c_expression import Tensor as Tensor_ # pylint: disable=no-name-in-module
 from mindspore.ops.operations._inner_ops import Send, Receive
 from mindspore.communication import GlobalComm, get_group_rank_from_world_rank
 from mindspore.ops._primitive_cache import _get_cache_prim
+try:
+    from mindspore._c_expression import TensorPy as Tensor_ # pylint: disable=no-name-in-module
+except:
+    from mindspore._c_expression import Tensor as Tensor_ # pylint: disable=no-name-in-module
+
 
 def isend(tensor, dst=0, group=GlobalComm.WORLD_COMM_GROUP, tag=0):
     """

diff --git a/mindnlp/transformers/models/jamba/modeling_jamba.py b/mindnlp/transformers/models/jamba/modeling_jamba.py
@@ -23,7 +23,12 @@
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import mindspore
-from mindspore._c_expression import Tensor as RawTensor # pylint: disable=no-name-in-module
+try:
+    from mindspore._c_expression import TensorPy as RawTensor # pylint: disable=no-name-in-module
+except:
+    from mindspore._c_expression import Tensor as RawTensor # pylint: disable=no-name-in-module
+
+
 import mindnlp.core.nn.functional as F
 from mindnlp.core import nn, ops
 from mindnlp.core.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss