Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / mindspore / autograd_compat.py: 84%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15 

16"""MindSpore backward-style autograd compatibility helpers.""" 

17# pylint: disable=protected-access,import-outside-toplevel 

18 

19from __future__ import annotations 

20 

21import warnings 

22 

23from mindspore import ops 

24from mindspore._c_expression import TensorPy, pyboost_detach, run_backward 

25from mindspore._c_expression import typing 

26from mindspore.graph.api import _pynative_executor 

27 

28_BACKWARD_COMPAT_ENABLED = False 

29 

30 

31@property 

32def requires_grad(self): 

33 """Return whether the tensor requires gradient.""" 

34 return self._requires_grad 

35 

36 

37@requires_grad.setter 

38def requires_grad(self, value=True): 

39 if not isinstance(value, bool): 

40 raise TypeError("The argument `requires_grad` must be bool type") 

41 self._requires_grad = value 

42 

43 

44@property 

45def grad(self): 

46 """Return the current accumulated gradient.""" 

47 if not self.is_leaf and self.requires_grad: 

48 warnings.warn( 

49 "The .grad attribute of a Tensor that is not a leaf Tensor is being accessed. " 

50 "Its .grad attribute won't be populated during autograd.backward(). " 

51 "If you indeed want the .grad field to be populated for a non-leaf Tensor, " 

52 "use .retain_grad() on the non-leaf Tensor.", 

53 stacklevel=2, 

54 ) 

55 dtensor_grad = getattr(self, "_dtensor_grad", None) 

56 if dtensor_grad is not None: 

57 return dtensor_grad 

58 return self._grad 

59 

60 

61@grad.setter 

62def grad(self, value): 

63 try: 

64 from hyper_parallel.core.dtensor.dtensor import DTensor 

65 except ImportError: 

66 DTensor = () 

67 

68 if value is None: 

69 self._dtensor_grad = None 

70 self._grad = None 

71 return 

72 

73 if DTensor and isinstance(value, DTensor): 

74 self._dtensor_grad = value 

75 self._grad = value._local_tensor 

76 return 

77 

78 self._dtensor_grad = None 

79 self._grad = value 

80 

81 

82@property 

83def is_leaf(self): 

84 """Return whether the tensor is a leaf.""" 

85 return self._is_leaf 

86 

87 

88@property 

89def retains_grad(self): 

90 """Return whether the tensor retains gradient.""" 

91 return self._retains_grad 

92 

93 

94@property 

95def grad_fn(self): 

96 """Return the gradient function node, or None for leaf tensors.""" 

97 if self._grad_node and self._grad_node.is_leaf(): 

98 return None 

99 return self._grad_node 

100 

101 

102@property 

103def output_nr(self): 

104 """Return the output index of this tensor in the autograd graph.""" 

105 return self._output_index 

106 

107 

108def retain_grad(self): 

109 """Set the tensor retains gradient.""" 

110 return self._retain_grad() 

111 

112 

113def detach(self): 

114 """Detach the tensor.""" 

115 detached = pyboost_detach(self) 

116 detached._dtensor_grad = None 

117 return detached 

118 

119 

120def _is_same_size(output, grad_tensor): 

121 return tuple(output.shape) == tuple(grad_tensor.shape) 

122 

123 

124def _calculate_shape(output, grad_tensor): 

125 return output.shape, grad_tensor.shape 

126 

127 

128def _tensor_or_tensors_to_tuple(tensors, length): 

129 if tensors is None: 

130 return (None,) * length 

131 if isinstance(tensors, TensorPy): 

132 return (tensors,) 

133 return tuple(tensors) 

134 

135 

136def _make_grads(outputs, grads): 

137 """Validate backward gradients and materialize implicit scalar grads.""" 

138 new_grads = [] 

139 for index, (out, grad_tensor) in enumerate(zip(outputs, grads)): 

140 if isinstance(grad_tensor, TensorPy): 

141 if not _is_same_size(out, grad_tensor): 

142 out_shape, grad_shape = _calculate_shape(out, grad_tensor) 

143 raise RuntimeError( 

144 "Mismatch in shape: grad_output[" 

145 + str(index) 

146 + "] has a shape of " 

147 + str(grad_shape) 

148 + " and output[" 

149 + str(index) 

150 + "] has a shape of " 

151 + str(out_shape) 

152 + "." 

153 ) 

154 if out.dtype.is_complex != grad_tensor.dtype.is_complex: 

155 raise RuntimeError( 

156 "For complex Tensors, both grad_output and output" 

157 " are required to have the same dtype." 

158 " Mismatch in dtype: grad_output[" 

159 + str(index) 

160 + "] has a dtype of " 

161 + str(grad_tensor.dtype) 

162 + " and output[" 

163 + str(index) 

164 + "] has a dtype of " 

165 + str(out.dtype) 

166 + "." 

167 ) 

168 new_grads.append(grad_tensor) 

169 elif grad_tensor is None: 

170 if out.numel() != 1: 

171 raise RuntimeError("grad can be implicitly created only for scalar outputs") 

172 if not isinstance(out.dtype, (typing.Float, typing.BFloat)): 

173 raise RuntimeError( 

174 f"grad can be implicitly created only for real scalar outputs but got {out.dtype}" 

175 ) 

176 new_grads.append(ops.ones_like(out)) 

177 else: 

178 raise TypeError( 

179 "gradients can be either Tensors or None, but got " + type(grad_tensor).__name__ 

180 ) 

181 return tuple(new_grads) 

182 

183 

184def backward(self, gradient=None, retain_graph=None, create_graph=False, inputs=None): 

185 """Run torch-style backward on a MindSpore tensor.""" 

186 outputs = (self,) 

187 has_explicit_inputs = inputs is not None 

188 if isinstance(inputs, list): 

189 inputs = tuple(inputs) 

190 elif isinstance(inputs, TensorPy): 

191 inputs = (inputs,) 

192 elif inputs is None: 

193 inputs = () 

194 else: 

195 inputs = tuple(inputs) 

196 if has_explicit_inputs and len(inputs) == 0: 

197 raise RuntimeError("'inputs' argument to backward() cannot be empty.") 

198 

199 grad_tensors = _tensor_or_tensors_to_tuple(gradient, len(outputs)) 

200 grad_tensors = _make_grads(outputs, grad_tensors) 

201 if retain_graph is None: 

202 retain_graph = create_graph 

203 

204 return run_backward( 

205 outputs, 

206 grad_tensors, 

207 retain_graph, 

208 create_graph, 

209 inputs, 

210 allow_unreachable=True, 

211 accumulate_grad=True, 

212 ) 

213 

214 

215def enable_mindspore_backward_compat() -> None: 

216 """Enable torch-like ``Tensor.backward()`` semantics for MindSpore PyNative.""" 

217 global _BACKWARD_COMPAT_ENABLED 

218 if _BACKWARD_COMPAT_ENABLED: 

219 return 

220 

221 _pynative_executor.set_grad_flag(True) 

222 TensorPy.requires_grad = requires_grad 

223 TensorPy.grad = grad 

224 TensorPy.backward = backward 

225 TensorPy.is_leaf = is_leaf 

226 TensorPy.retains_grad = retains_grad 

227 TensorPy.retain_grad = retain_grad 

228 TensorPy.grad_fn = grad_fn 

229 TensorPy.output_nr = output_nr 

230 TensorPy.detach = detach 

231 _BACKWARD_COMPAT_ENABLED = True