Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/tensorflow/python/ops/gen_nccl_ops.py: 18%
129 statements
« 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
24def nccl_all_reduce(input, reduction, num_devices, shared_name, name=None):
25 r"""Outputs a tensor containing the reduction across all input tensors.
27 Outputs a tensor containing the reduction across all input tensors passed to ops
28 within the same `shared_name.
30 The graph should be constructed so if one op runs with shared_name value `c`,
31 then `num_devices` ops will run with shared_name value `c`. Failure to do so
32 will cause the graph execution to fail to complete.
34 input: the input to the reduction
35 data: the value of the reduction across all `num_devices` devices.
36 reduction: the reduction operation to perform.
37 num_devices: The number of devices participating in this reduction.
38 shared_name: Identifier that shared between ops of the same reduction.
40 Args:
41 input: A `Tensor`. Must be one of the following types: `half`, `float32`, `float64`, `int32`, `int64`.
42 reduction: A `string` from: `"min", "max", "prod", "sum"`.
43 num_devices: An `int`.
44 shared_name: A `string`.
45 name: A name for the operation (optional).
47 Returns:
48 A `Tensor`. Has the same type as `input`.
49 """
50 _ctx = _context._context or _context.context()
51 tld = _ctx._thread_local_data
52 if tld.is_eager:
53 try:
54 _result = pywrap_tfe.TFE_Py_FastPathExecute(
55 _ctx, "NcclAllReduce", name, input, "reduction", reduction,
56 "num_devices", num_devices, "shared_name", shared_name)
57 return _result
58 except _core._NotOkStatusException as e:
59 _ops.raise_from_not_ok_status(e, name)
60 except _core._FallbackException:
61 pass
62 try:
63 return nccl_all_reduce_eager_fallback(
64 input, reduction=reduction, num_devices=num_devices,
65 shared_name=shared_name, name=name, ctx=_ctx)
66 except _core._SymbolicException:
67 pass # Add nodes to the TensorFlow graph.
68 # Add nodes to the TensorFlow graph.
69 reduction = _execute.make_str(reduction, "reduction")
70 num_devices = _execute.make_int(num_devices, "num_devices")
71 shared_name = _execute.make_str(shared_name, "shared_name")
72 _, _, _op, _outputs = _op_def_library._apply_op_helper(
73 "NcclAllReduce", input=input, reduction=reduction,
74 num_devices=num_devices, shared_name=shared_name,
75 name=name)
76 _result = _outputs[:]
77 if _execute.must_record_gradient():
78 _attrs = ("reduction", _op.get_attr("reduction"), "T",
79 _op._get_attr_type("T"), "num_devices",
80 _op._get_attr_int("num_devices"), "shared_name",
81 _op.get_attr("shared_name"))
82 _inputs_flat = _op.inputs
83 _execute.record_gradient(
84 "NcclAllReduce", _inputs_flat, _attrs, _result)
85 _result, = _result
86 return _result
88NcclAllReduce = tf_export("raw_ops.NcclAllReduce")(_ops.to_raw_op(nccl_all_reduce))
91def nccl_all_reduce_eager_fallback(input, reduction, num_devices, shared_name, name, ctx):
92 reduction = _execute.make_str(reduction, "reduction")
93 num_devices = _execute.make_int(num_devices, "num_devices")
94 shared_name = _execute.make_str(shared_name, "shared_name")
95 _attr_T, (input,) = _execute.args_to_matching_eager([input], ctx, [_dtypes.half, _dtypes.float32, _dtypes.float64, _dtypes.int32, _dtypes.int64, ])
96 _inputs_flat = [input]
97 _attrs = ("reduction", reduction, "T", _attr_T, "num_devices", num_devices,
98 "shared_name", shared_name)
99 _result = _execute.execute(b"NcclAllReduce", 1, inputs=_inputs_flat,
100 attrs=_attrs, ctx=ctx, name=name)
101 if _execute.must_record_gradient():
102 _execute.record_gradient(
103 "NcclAllReduce", _inputs_flat, _attrs, _result)
104 _result, = _result
105 return _result
108def nccl_broadcast(input, shape, name=None):
109 r"""Sends `input` to all devices that are connected to the output.
111 Sends `input` to all devices that are connected to the output.
113 The graph should be constructed so that all ops connected to the output have a
114 valid device assignment, and the op itself is assigned one of these devices.
116 input: The input to the broadcast.
117 output: The same as input.
118 shape: The shape of the input tensor.
120 Args:
121 input: A `Tensor`. Must be one of the following types: `half`, `float32`, `float64`, `int32`, `int64`.
122 shape: A `tf.TensorShape` or list of `ints`.
123 name: A name for the operation (optional).
125 Returns:
126 A `Tensor`. Has the same type as `input`.
127 """
128 _ctx = _context._context or _context.context()
129 tld = _ctx._thread_local_data
130 if tld.is_eager:
131 try:
132 _result = pywrap_tfe.TFE_Py_FastPathExecute(
133 _ctx, "NcclBroadcast", name, input, "shape", shape)
134 return _result
135 except _core._NotOkStatusException as e:
136 _ops.raise_from_not_ok_status(e, name)
137 except _core._FallbackException:
138 pass
139 try:
140 return nccl_broadcast_eager_fallback(
141 input, shape=shape, name=name, ctx=_ctx)
142 except _core._SymbolicException:
143 pass # Add nodes to the TensorFlow graph.
144 # Add nodes to the TensorFlow graph.
145 shape = _execute.make_shape(shape, "shape")
146 _, _, _op, _outputs = _op_def_library._apply_op_helper(
147 "NcclBroadcast", input=input, shape=shape, name=name)
148 _result = _outputs[:]
149 if _execute.must_record_gradient():
150 _attrs = ("T", _op._get_attr_type("T"), "shape", _op.get_attr("shape"))
151 _inputs_flat = _op.inputs
152 _execute.record_gradient(
153 "NcclBroadcast", _inputs_flat, _attrs, _result)
154 _result, = _result
155 return _result
157NcclBroadcast = tf_export("raw_ops.NcclBroadcast")(_ops.to_raw_op(nccl_broadcast))
160def nccl_broadcast_eager_fallback(input, shape, name, ctx):
161 shape = _execute.make_shape(shape, "shape")
162 _attr_T, (input,) = _execute.args_to_matching_eager([input], ctx, [_dtypes.half, _dtypes.float32, _dtypes.float64, _dtypes.int32, _dtypes.int64, ])
163 _inputs_flat = [input]
164 _attrs = ("T", _attr_T, "shape", shape)
165 _result = _execute.execute(b"NcclBroadcast", 1, inputs=_inputs_flat,
166 attrs=_attrs, ctx=ctx, name=name)
167 if _execute.must_record_gradient():
168 _execute.record_gradient(
169 "NcclBroadcast", _inputs_flat, _attrs, _result)
170 _result, = _result
171 return _result
174def nccl_reduce(input, reduction, name=None):
175 r"""Reduces `input` from `num_devices` using `reduction` to a single device.
177 Reduces `input` from `num_devices` using `reduction` to a single device.
179 The graph should be constructed so that all inputs have a valid device
180 assignment, and the op itself is assigned one of these devices.
182 input: The input to the reduction.
183 data: the value of the reduction across all `num_devices` devices.
184 reduction: the reduction operation to perform.
186 Args:
187 input: A list of at least 1 `Tensor` objects with the same type in: `half`, `float32`, `float64`, `int32`, `int64`.
188 reduction: A `string` from: `"min", "max", "prod", "sum"`.
189 name: A name for the operation (optional).
191 Returns:
192 A `Tensor`. Has the same type as `input`.
193 """
194 _ctx = _context._context or _context.context()
195 tld = _ctx._thread_local_data
196 if tld.is_eager:
197 try:
198 _result = pywrap_tfe.TFE_Py_FastPathExecute(
199 _ctx, "NcclReduce", name, input, "reduction", reduction)
200 return _result
201 except _core._NotOkStatusException as e:
202 _ops.raise_from_not_ok_status(e, name)
203 except _core._FallbackException:
204 pass
205 try:
206 return nccl_reduce_eager_fallback(
207 input, reduction=reduction, name=name, ctx=_ctx)
208 except _core._SymbolicException:
209 pass # Add nodes to the TensorFlow graph.
210 # Add nodes to the TensorFlow graph.
211 if not isinstance(input, (list, tuple)):
212 raise TypeError(
213 "Expected list for 'input' argument to "
214 "'nccl_reduce' Op, not %r." % input)
215 _attr_num_devices = len(input)
216 reduction = _execute.make_str(reduction, "reduction")
217 _, _, _op, _outputs = _op_def_library._apply_op_helper(
218 "NcclReduce", input=input, reduction=reduction, name=name)
219 _result = _outputs[:]
220 if _execute.must_record_gradient():
221 _attrs = ("reduction", _op.get_attr("reduction"), "T",
222 _op._get_attr_type("T"), "num_devices",
223 _op._get_attr_int("num_devices"))
224 _inputs_flat = _op.inputs
225 _execute.record_gradient(
226 "NcclReduce", _inputs_flat, _attrs, _result)
227 _result, = _result
228 return _result
230NcclReduce = tf_export("raw_ops.NcclReduce")(_ops.to_raw_op(nccl_reduce))
233def nccl_reduce_eager_fallback(input, reduction, name, ctx):
234 if not isinstance(input, (list, tuple)):
235 raise TypeError(
236 "Expected list for 'input' argument to "
237 "'nccl_reduce' Op, not %r." % input)
238 _attr_num_devices = len(input)
239 reduction = _execute.make_str(reduction, "reduction")
240 _attr_T, input = _execute.args_to_matching_eager(list(input), ctx, [_dtypes.half, _dtypes.float32, _dtypes.float64, _dtypes.int32, _dtypes.int64, ])
241 _inputs_flat = list(input)
242 _attrs = ("reduction", reduction, "T", _attr_T, "num_devices",
243 _attr_num_devices)
244 _result = _execute.execute(b"NcclReduce", 1, inputs=_inputs_flat,
245 attrs=_attrs, ctx=ctx, name=name)
246 if _execute.must_record_gradient():
247 _execute.record_gradient(
248 "NcclReduce", _inputs_flat, _attrs, _result)
249 _result, = _result
250 return _result