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
« prev ^ index » next coverage.py v7.4.0, created at 2024-01-03 07:57 +0000
1"""Python wrappers around TensorFlow ops.
3This file is MACHINE GENERATED! Do not edit.
4"""
6import collections
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
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
22from typing import TypeVar
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.
30 Args:
31 input: A `Tensor`.
32 name: A name for the operation (optional).
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
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
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
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.
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).
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
196XlaLaunch = tf_export("raw_ops.XlaLaunch")(_ops.to_raw_op(xla_launch))
197_dispatcher_for_xla_launch = xla_launch._tf_type_based_dispatcher.Dispatch
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
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.
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).
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
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
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