Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/tensorflow/compiler/jit/ops/xla_ops.py: 17%

206 statements  

« prev     ^ index     » next       coverage.py v7.4.0, created at 2024-01-03 07:57 +0000

1"""Python wrappers around TensorFlow ops. 

2 

3This file is MACHINE GENERATED! Do not edit. 

4""" 

5 

6import collections 

7 

8from tensorflow.python import pywrap_tfe as pywrap_tfe 

9from tensorflow.python.eager import context as _context 

10from tensorflow.python.eager import core as _core 

11from tensorflow.python.eager import execute as _execute 

12from tensorflow.python.framework import dtypes as _dtypes 

13from tensorflow.security.fuzzing.py import annotation_types as _atypes 

14 

15from tensorflow.python.framework import op_def_registry as _op_def_registry 

16from tensorflow.python.framework import ops as _ops 

17from tensorflow.python.framework import op_def_library as _op_def_library 

18from tensorflow.python.util.deprecation import deprecated_endpoints 

19from tensorflow.python.util import dispatch as _dispatch 

20from tensorflow.python.util.tf_export import tf_export 

21 

22from typing import TypeVar 

23 

24@_dispatch.add_fallback_dispatch_list 

25@_dispatch.add_type_based_api_dispatcher 

26@tf_export('xla_cluster_output') 

27def xla_cluster_output(input, name=None): 

28 r"""Operator that connects the output of an XLA computation to other consumer graph nodes. 

29 

30 Args: 

31 input: A `Tensor`. 

32 name: A name for the operation (optional). 

33 

34 Returns: 

35 A `Tensor`. Has the same type as `input`. 

36 """ 

37 _ctx = _context._context or _context.context() 

38 tld = _ctx._thread_local_data 

39 if tld.is_eager: 

40 try: 

41 _result = pywrap_tfe.TFE_Py_FastPathExecute( 

42 _ctx, "XlaClusterOutput", name, input) 

43 return _result 

44 except _core._NotOkStatusException as e: 

45 _ops.raise_from_not_ok_status(e, name) 

46 except _core._FallbackException: 

47 pass 

48 try: 

49 _result = _dispatcher_for_xla_cluster_output( 

50 (input, name,), None) 

51 if _result is not NotImplemented: 

52 return _result 

53 return xla_cluster_output_eager_fallback( 

54 input, name=name, ctx=_ctx) 

55 except _core._SymbolicException: 

56 pass # Add nodes to the TensorFlow graph. 

57 except (TypeError, ValueError): 

58 _result = _dispatch.dispatch( 

59 xla_cluster_output, (), dict(input=input, name=name) 

60 ) 

61 if _result is not _dispatch.OpDispatcher.NOT_SUPPORTED: 

62 return _result 

63 raise 

64 else: 

65 _result = _dispatcher_for_xla_cluster_output( 

66 (input, name,), None) 

67 if _result is not NotImplemented: 

68 return _result 

69 # Add nodes to the TensorFlow graph. 

70 try: 

71 _, _, _op, _outputs = _op_def_library._apply_op_helper( 

72 "XlaClusterOutput", input=input, name=name) 

73 except (TypeError, ValueError): 

74 _result = _dispatch.dispatch( 

75 xla_cluster_output, (), dict(input=input, name=name) 

76 ) 

77 if _result is not _dispatch.OpDispatcher.NOT_SUPPORTED: 

78 return _result 

79 raise 

80 _result = _outputs[:] 

81 if _execute.must_record_gradient(): 

82 _attrs = ("T", _op._get_attr_type("T")) 

83 _inputs_flat = _op.inputs 

84 _execute.record_gradient( 

85 "XlaClusterOutput", _inputs_flat, _attrs, _result) 

86 _result, = _result 

87 return _result 

88 

89XlaClusterOutput = tf_export("raw_ops.XlaClusterOutput")(_ops.to_raw_op(xla_cluster_output)) 

90_dispatcher_for_xla_cluster_output = xla_cluster_output._tf_type_based_dispatcher.Dispatch 

91 

92 

93def xla_cluster_output_eager_fallback(input, name, ctx): 

94 _attr_T, (input,) = _execute.args_to_matching_eager([input], ctx, []) 

95 _inputs_flat = [input] 

96 _attrs = ("T", _attr_T) 

97 _result = _execute.execute(b"XlaClusterOutput", 1, inputs=_inputs_flat, 

98 attrs=_attrs, ctx=ctx, name=name) 

99 if _execute.must_record_gradient(): 

100 _execute.record_gradient( 

101 "XlaClusterOutput", _inputs_flat, _attrs, _result) 

102 _result, = _result 

103 return _result 

104 

105 

106@_dispatch.add_fallback_dispatch_list 

107@_dispatch.add_type_based_api_dispatcher 

108@tf_export('xla_launch') 

109def xla_launch(constants, args, resources, Tresults, function, name=None): 

110 r"""XLA Launch Op. For use by the XLA JIT only. 

111 

112 Args: 

113 constants: A list of `Tensor` objects. 

114 args: A list of `Tensor` objects. 

115 resources: A list of `Tensor` objects with type `resource`. 

116 Tresults: A list of `tf.DTypes`. 

117 function: A function decorated with @Defun. 

118 name: A name for the operation (optional). 

119 

120 Returns: 

121 A list of `Tensor` objects of type `Tresults`. 

122 """ 

123 _ctx = _context._context or _context.context() 

124 tld = _ctx._thread_local_data 

125 if tld.is_eager: 

126 try: 

127 _result = pywrap_tfe.TFE_Py_FastPathExecute( 

128 _ctx, "XlaLaunch", name, constants, args, resources, "Tresults", 

129 Tresults, "function", function) 

130 return _result 

131 except _core._NotOkStatusException as e: 

132 _ops.raise_from_not_ok_status(e, name) 

133 except _core._FallbackException: 

134 pass 

135 try: 

136 _result = _dispatcher_for_xla_launch( 

137 (constants, args, resources, Tresults, function, name,), None) 

138 if _result is not NotImplemented: 

139 return _result 

140 return xla_launch_eager_fallback( 

141 constants, args, resources, Tresults=Tresults, function=function, 

142 name=name, ctx=_ctx) 

143 except _core._SymbolicException: 

144 pass # Add nodes to the TensorFlow graph. 

145 except (TypeError, ValueError): 

146 _result = _dispatch.dispatch( 

147 xla_launch, (), dict(constants=constants, args=args, 

148 resources=resources, Tresults=Tresults, 

149 function=function, name=name) 

150 ) 

151 if _result is not _dispatch.OpDispatcher.NOT_SUPPORTED: 

152 return _result 

153 raise 

154 else: 

155 _result = _dispatcher_for_xla_launch( 

156 (constants, args, resources, Tresults, function, name,), None) 

157 if _result is not NotImplemented: 

158 return _result 

159 # Add nodes to the TensorFlow graph. 

160 if not isinstance(resources, (list, tuple)): 

161 raise TypeError( 

162 "Expected list for 'resources' argument to " 

163 "'xla_launch' Op, not %r." % resources) 

164 _attr_Nresources = len(resources) 

165 if not isinstance(Tresults, (list, tuple)): 

166 raise TypeError( 

167 "Expected list for 'Tresults' argument to " 

168 "'xla_launch' Op, not %r." % Tresults) 

169 Tresults = [_execute.make_type(_t, "Tresults") for _t in Tresults] 

170 try: 

171 _, _, _op, _outputs = _op_def_library._apply_op_helper( 

172 "XlaLaunch", constants=constants, args=args, resources=resources, 

173 Tresults=Tresults, function=function, name=name) 

174 except (TypeError, ValueError): 

175 _result = _dispatch.dispatch( 

176 xla_launch, (), dict(constants=constants, args=args, 

177 resources=resources, Tresults=Tresults, 

178 function=function, name=name) 

179 ) 

180 if _result is not _dispatch.OpDispatcher.NOT_SUPPORTED: 

181 return _result 

182 raise 

183 _result = _outputs[:] 

184 if not _result: 

185 return _op 

186 if _execute.must_record_gradient(): 

187 _attrs = ("Tconstants", _op.get_attr("Tconstants"), "Targs", 

188 _op.get_attr("Targs"), "Nresources", 

189 _op._get_attr_int("Nresources"), "Tresults", 

190 _op.get_attr("Tresults"), "function", _op.get_attr("function")) 

191 _inputs_flat = _op.inputs 

192 _execute.record_gradient( 

193 "XlaLaunch", _inputs_flat, _attrs, _result) 

194 return _result 

195 

196XlaLaunch = tf_export("raw_ops.XlaLaunch")(_ops.to_raw_op(xla_launch)) 

197_dispatcher_for_xla_launch = xla_launch._tf_type_based_dispatcher.Dispatch 

198 

199 

200def xla_launch_eager_fallback(constants, args, resources, Tresults, function, name, ctx): 

201 if not isinstance(resources, (list, tuple)): 

202 raise TypeError( 

203 "Expected list for 'resources' argument to " 

204 "'xla_launch' Op, not %r." % resources) 

205 _attr_Nresources = len(resources) 

206 if not isinstance(Tresults, (list, tuple)): 

207 raise TypeError( 

208 "Expected list for 'Tresults' argument to " 

209 "'xla_launch' Op, not %r." % Tresults) 

210 Tresults = [_execute.make_type(_t, "Tresults") for _t in Tresults] 

211 _attr_Tconstants, constants = _execute.convert_to_mixed_eager_tensors(constants, ctx) 

212 _attr_Targs, args = _execute.convert_to_mixed_eager_tensors(args, ctx) 

213 resources = _ops.convert_n_to_tensor(resources, _dtypes.resource) 

214 _inputs_flat = list(constants) + list(args) + list(resources) 

215 _attrs = ("Tconstants", _attr_Tconstants, "Targs", _attr_Targs, 

216 "Nresources", _attr_Nresources, "Tresults", Tresults, "function", function) 

217 _result = _execute.execute(b"XlaLaunch", len(Tresults), inputs=_inputs_flat, 

218 attrs=_attrs, ctx=ctx, name=name) 

219 if _execute.must_record_gradient(): 

220 _execute.record_gradient( 

221 "XlaLaunch", _inputs_flat, _attrs, _result) 

222 return _result 

223 

224 

225@_dispatch.add_fallback_dispatch_list 

226@_dispatch.add_type_based_api_dispatcher 

227@tf_export('xla_launch_v2') 

228def xla_launch_v2(args, Tresults, constants, resources, function, name=None): 

229 r"""XLA Launch Op. For use by the XLA JIT only. 

230 

231 Args: 

232 args: A list of `Tensor` objects. 

233 Tresults: A list of `tf.DTypes`. 

234 constants: A list of `ints`. 

235 resources: A list of `ints`. 

236 function: A function decorated with @Defun. 

237 name: A name for the operation (optional). 

238 

239 Returns: 

240 A list of `Tensor` objects of type `Tresults`. 

241 """ 

242 _ctx = _context._context or _context.context() 

243 tld = _ctx._thread_local_data 

244 if tld.is_eager: 

245 try: 

246 _result = pywrap_tfe.TFE_Py_FastPathExecute( 

247 _ctx, "XlaLaunchV2", name, args, "Tresults", Tresults, "constants", 

248 constants, "resources", resources, "function", function) 

249 return _result 

250 except _core._NotOkStatusException as e: 

251 _ops.raise_from_not_ok_status(e, name) 

252 except _core._FallbackException: 

253 pass 

254 try: 

255 _result = _dispatcher_for_xla_launch_v2( 

256 (args, Tresults, constants, resources, function, name,), None) 

257 if _result is not NotImplemented: 

258 return _result 

259 return xla_launch_v2_eager_fallback( 

260 args, Tresults=Tresults, constants=constants, resources=resources, 

261 function=function, name=name, ctx=_ctx) 

262 except _core._SymbolicException: 

263 pass # Add nodes to the TensorFlow graph. 

264 except (TypeError, ValueError): 

265 _result = _dispatch.dispatch( 

266 xla_launch_v2, (), dict(args=args, Tresults=Tresults, 

267 constants=constants, resources=resources, 

268 function=function, name=name) 

269 ) 

270 if _result is not _dispatch.OpDispatcher.NOT_SUPPORTED: 

271 return _result 

272 raise 

273 else: 

274 _result = _dispatcher_for_xla_launch_v2( 

275 (args, Tresults, constants, resources, function, name,), None) 

276 if _result is not NotImplemented: 

277 return _result 

278 # Add nodes to the TensorFlow graph. 

279 if not isinstance(Tresults, (list, tuple)): 

280 raise TypeError( 

281 "Expected list for 'Tresults' argument to " 

282 "'xla_launch_v2' Op, not %r." % Tresults) 

283 Tresults = [_execute.make_type(_t, "Tresults") for _t in Tresults] 

284 if not isinstance(constants, (list, tuple)): 

285 raise TypeError( 

286 "Expected list for 'constants' argument to " 

287 "'xla_launch_v2' Op, not %r." % constants) 

288 constants = [_execute.make_int(_i, "constants") for _i in constants] 

289 if not isinstance(resources, (list, tuple)): 

290 raise TypeError( 

291 "Expected list for 'resources' argument to " 

292 "'xla_launch_v2' Op, not %r." % resources) 

293 resources = [_execute.make_int(_i, "resources") for _i in resources] 

294 try: 

295 _, _, _op, _outputs = _op_def_library._apply_op_helper( 

296 "XlaLaunchV2", args=args, Tresults=Tresults, constants=constants, 

297 resources=resources, function=function, name=name) 

298 except (TypeError, ValueError): 

299 _result = _dispatch.dispatch( 

300 xla_launch_v2, (), dict(args=args, Tresults=Tresults, 

301 constants=constants, resources=resources, 

302 function=function, name=name) 

303 ) 

304 if _result is not _dispatch.OpDispatcher.NOT_SUPPORTED: 

305 return _result 

306 raise 

307 _result = _outputs[:] 

308 if _execute.must_record_gradient(): 

309 _attrs = ("Targs", _op.get_attr("Targs"), "Tresults", 

310 _op.get_attr("Tresults"), "constants", 

311 _op.get_attr("constants"), "resources", 

312 _op.get_attr("resources"), "function", _op.get_attr("function")) 

313 _inputs_flat = _op.inputs 

314 _execute.record_gradient( 

315 "XlaLaunchV2", _inputs_flat, _attrs, _result) 

316 return _result 

317 

318XlaLaunchV2 = tf_export("raw_ops.XlaLaunchV2")(_ops.to_raw_op(xla_launch_v2)) 

319_dispatcher_for_xla_launch_v2 = xla_launch_v2._tf_type_based_dispatcher.Dispatch 

320 

321 

322def xla_launch_v2_eager_fallback(args, Tresults, constants, resources, function, name, ctx): 

323 if not isinstance(Tresults, (list, tuple)): 

324 raise TypeError( 

325 "Expected list for 'Tresults' argument to " 

326 "'xla_launch_v2' Op, not %r." % Tresults) 

327 Tresults = [_execute.make_type(_t, "Tresults") for _t in Tresults] 

328 if not isinstance(constants, (list, tuple)): 

329 raise TypeError( 

330 "Expected list for 'constants' argument to " 

331 "'xla_launch_v2' Op, not %r." % constants) 

332 constants = [_execute.make_int(_i, "constants") for _i in constants] 

333 if not isinstance(resources, (list, tuple)): 

334 raise TypeError( 

335 "Expected list for 'resources' argument to " 

336 "'xla_launch_v2' Op, not %r." % resources) 

337 resources = [_execute.make_int(_i, "resources") for _i in resources] 

338 _attr_Targs, args = _execute.convert_to_mixed_eager_tensors(args, ctx) 

339 _inputs_flat = list(args) 

340 _attrs = ("Targs", _attr_Targs, "Tresults", Tresults, "constants", 

341 constants, "resources", resources, "function", function) 

342 _result = _execute.execute(b"XlaLaunchV2", len(Tresults), 

343 inputs=_inputs_flat, attrs=_attrs, ctx=ctx, 

344 name=name) 

345 if _execute.must_record_gradient(): 

346 _execute.record_gradient( 

347 "XlaLaunchV2", _inputs_flat, _attrs, _result) 

348 return _result 

349