Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/tensorflow/python/ops/gen_tpu_ops.py: 7%

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 

24def all_to_all(input, group_assignment, concat_dimension, split_dimension, split_count, name=None): 

25 r"""An Op to exchange data across TPU replicas. 

26 

27 On each replica, the input is split into `split_count` blocks along 

28 `split_dimension` and send to the other replicas given group_assignment. After 

29 receiving `split_count` - 1 blocks from other replicas, we concatenate the 

30 blocks along `concat_dimension` as the output. 

31 

32 For example, suppose there are 2 TPU replicas: 

33 replica 0 receives input: `[[A, B]]` 

34 replica 1 receives input: `[[C, D]]` 

35 

36 group_assignment=`[[0, 1]]` 

37 concat_dimension=0 

38 split_dimension=1 

39 split_count=2 

40 

41 replica 0's output: `[[A], [C]]` 

42 replica 1's output: `[[B], [D]]` 

43 

44 Args: 

45 input: A `Tensor`. Must be one of the following types: `float32`, `float64`, `int32`, `uint8`, `int16`, `int8`, `complex64`, `int64`, `qint8`, `quint8`, `qint32`, `bfloat16`, `qint16`, `quint16`, `uint16`, `complex128`, `half`, `uint32`, `uint64`, `bool`. 

46 The local input to the sum. 

47 group_assignment: A `Tensor` of type `int32`. An int32 tensor with shape 

48 [num_groups, num_replicas_per_group]. `group_assignment[i]` represents the 

49 replica ids in the ith subgroup. 

50 concat_dimension: An `int`. The dimension number to concatenate. 

51 split_dimension: An `int`. The dimension number to split. 

52 split_count: An `int`. 

53 The number of splits, this number must equal to the sub-group 

54 size(group_assignment.get_shape()[1]) 

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

56 

57 Returns: 

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

59 """ 

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

61 tld = _ctx._thread_local_data 

62 if tld.is_eager: 

63 try: 

64 _result = pywrap_tfe.TFE_Py_FastPathExecute( 

65 _ctx, "AllToAll", name, input, group_assignment, "concat_dimension", 

66 concat_dimension, "split_dimension", split_dimension, "split_count", 

67 split_count) 

68 return _result 

69 except _core._NotOkStatusException as e: 

70 _ops.raise_from_not_ok_status(e, name) 

71 except _core._FallbackException: 

72 pass 

73 try: 

74 return all_to_all_eager_fallback( 

75 input, group_assignment, concat_dimension=concat_dimension, 

76 split_dimension=split_dimension, split_count=split_count, name=name, 

77 ctx=_ctx) 

78 except _core._SymbolicException: 

79 pass # Add nodes to the TensorFlow graph. 

80 # Add nodes to the TensorFlow graph. 

81 concat_dimension = _execute.make_int(concat_dimension, "concat_dimension") 

82 split_dimension = _execute.make_int(split_dimension, "split_dimension") 

83 split_count = _execute.make_int(split_count, "split_count") 

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

85 "AllToAll", input=input, group_assignment=group_assignment, 

86 concat_dimension=concat_dimension, 

87 split_dimension=split_dimension, split_count=split_count, 

88 name=name) 

89 _result = _outputs[:] 

90 if _execute.must_record_gradient(): 

91 _attrs = ("T", _op._get_attr_type("T"), "concat_dimension", 

92 _op._get_attr_int("concat_dimension"), "split_dimension", 

93 _op._get_attr_int("split_dimension"), "split_count", 

94 _op._get_attr_int("split_count")) 

95 _inputs_flat = _op.inputs 

96 _execute.record_gradient( 

97 "AllToAll", _inputs_flat, _attrs, _result) 

98 _result, = _result 

99 return _result 

100 

101AllToAll = tf_export("raw_ops.AllToAll")(_ops.to_raw_op(all_to_all)) 

102 

103 

104def all_to_all_eager_fallback(input, group_assignment, concat_dimension, split_dimension, split_count, name, ctx): 

105 concat_dimension = _execute.make_int(concat_dimension, "concat_dimension") 

106 split_dimension = _execute.make_int(split_dimension, "split_dimension") 

107 split_count = _execute.make_int(split_count, "split_count") 

108 _attr_T, (input,) = _execute.args_to_matching_eager([input], ctx, [_dtypes.float32, _dtypes.float64, _dtypes.int32, _dtypes.uint8, _dtypes.int16, _dtypes.int8, _dtypes.complex64, _dtypes.int64, _dtypes.qint8, _dtypes.quint8, _dtypes.qint32, _dtypes.bfloat16, _dtypes.qint16, _dtypes.quint16, _dtypes.uint16, _dtypes.complex128, _dtypes.half, _dtypes.uint32, _dtypes.uint64, _dtypes.bool, ]) 

109 group_assignment = _ops.convert_to_tensor(group_assignment, _dtypes.int32) 

110 _inputs_flat = [input, group_assignment] 

111 _attrs = ("T", _attr_T, "concat_dimension", concat_dimension, 

112 "split_dimension", split_dimension, "split_count", split_count) 

113 _result = _execute.execute(b"AllToAll", 1, inputs=_inputs_flat, 

114 attrs=_attrs, ctx=ctx, name=name) 

115 if _execute.must_record_gradient(): 

116 _execute.record_gradient( 

117 "AllToAll", _inputs_flat, _attrs, _result) 

118 _result, = _result 

119 return _result 

120 

121 

122def assign_variable_xla_concat_nd(resource, inputs, num_concats, paddings=[], name=None): 

123 r"""Concats input tensor across all dimensions. 

124 

125 An op which merges slices the input tensor based on the given num_splits 

126 attribute, strips paddings optionally, and writes the merged tensor without 

127 paddings to the resource variable. 

128 

129 This op may be generated via the TPU bridge. 

130 

131 For example, with `input` tensor: 

132 ``` 

133 [[0, 1], 

134 [4, 5]] 

135 [[2, 3], 

136 [6, 7]] 

137 [[8, 9], 

138 [12, 13]] 

139 [[10, 11], 

140 [14, 15]] 

141 ``` 

142 `num_splits`: 

143 ``` 

144 [2, 2] 

145 ``` 

146 and `paddings`: 

147 ``` 

148 [1, 1] 

149 ``` 

150 the expected `outputs` is: 

151 ``` 

152 [[0, 1, 2], 

153 [4, 5, 6], 

154 [8, 9, 10]] 

155 ``` 

156 

157 Args: 

158 resource: A `Tensor` of type `resource`. 

159 Resource variable for concatenated input tensors across all dimensions. 

160 } 

161 in_arg { 

162 name: "inputs" 

163 description: <<END 

164 Input tensor slices in row-major order to merge across all dimensions. All 

165 inputs must have the same shape. 

166 } 

167 out_arg { 

168 name: "output" 

169 description: <<END 

170 Output tensor formed from merging input slices based on num_concats defined. 

171 inputs: A list of at least 1 `Tensor` objects with the same type. 

172 num_concats: A list of `ints`. Number of ways to merge per dimension. 

173 paddings: An optional list of `ints`. Defaults to `[]`. 

174 Optional list of right paddings per dimension to strip from the final merged 

175 tensor. These paddings must not exceed the dimension size of the merged result 

176 prior to stripping paddings. 

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

178 

179 Returns: 

180 The created Operation. 

181 """ 

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

183 tld = _ctx._thread_local_data 

184 if tld.is_eager: 

185 try: 

186 _result = pywrap_tfe.TFE_Py_FastPathExecute( 

187 _ctx, "AssignVariableXlaConcatND", name, resource, inputs, 

188 "num_concats", num_concats, "paddings", paddings) 

189 return _result 

190 except _core._NotOkStatusException as e: 

191 _ops.raise_from_not_ok_status(e, name) 

192 except _core._FallbackException: 

193 pass 

194 try: 

195 return assign_variable_xla_concat_nd_eager_fallback( 

196 resource, inputs, num_concats=num_concats, paddings=paddings, 

197 name=name, ctx=_ctx) 

198 except _core._SymbolicException: 

199 pass # Add nodes to the TensorFlow graph. 

200 # Add nodes to the TensorFlow graph. 

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

202 raise TypeError( 

203 "Expected list for 'inputs' argument to " 

204 "'assign_variable_xla_concat_nd' Op, not %r." % inputs) 

205 _attr_N = len(inputs) 

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

207 raise TypeError( 

208 "Expected list for 'num_concats' argument to " 

209 "'assign_variable_xla_concat_nd' Op, not %r." % num_concats) 

210 num_concats = [_execute.make_int(_i, "num_concats") for _i in num_concats] 

211 if paddings is None: 

212 paddings = [] 

213 if not isinstance(paddings, (list, tuple)): 

214 raise TypeError( 

215 "Expected list for 'paddings' argument to " 

216 "'assign_variable_xla_concat_nd' Op, not %r." % paddings) 

217 paddings = [_execute.make_int(_i, "paddings") for _i in paddings] 

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

219 "AssignVariableXlaConcatND", resource=resource, inputs=inputs, 

220 num_concats=num_concats, 

221 paddings=paddings, name=name) 

222 return _op 

223AssignVariableXlaConcatND = tf_export("raw_ops.AssignVariableXlaConcatND")(_ops.to_raw_op(assign_variable_xla_concat_nd)) 

224 

225 

226def assign_variable_xla_concat_nd_eager_fallback(resource, inputs, num_concats, paddings, name, ctx): 

227 if not isinstance(inputs, (list, tuple)): 

228 raise TypeError( 

229 "Expected list for 'inputs' argument to " 

230 "'assign_variable_xla_concat_nd' Op, not %r." % inputs) 

231 _attr_N = len(inputs) 

232 if not isinstance(num_concats, (list, tuple)): 

233 raise TypeError( 

234 "Expected list for 'num_concats' argument to " 

235 "'assign_variable_xla_concat_nd' Op, not %r." % num_concats) 

236 num_concats = [_execute.make_int(_i, "num_concats") for _i in num_concats] 

237 if paddings is None: 

238 paddings = [] 

239 if not isinstance(paddings, (list, tuple)): 

240 raise TypeError( 

241 "Expected list for 'paddings' argument to " 

242 "'assign_variable_xla_concat_nd' Op, not %r." % paddings) 

243 paddings = [_execute.make_int(_i, "paddings") for _i in paddings] 

244 _attr_T, inputs = _execute.args_to_matching_eager(list(inputs), ctx, []) 

245 resource = _ops.convert_to_tensor(resource, _dtypes.resource) 

246 _inputs_flat = [resource] + list(inputs) 

247 _attrs = ("T", _attr_T, "N", _attr_N, "num_concats", num_concats, 

248 "paddings", paddings) 

249 _result = _execute.execute(b"AssignVariableXlaConcatND", 0, 

250 inputs=_inputs_flat, attrs=_attrs, ctx=ctx, 

251 name=name) 

252 _result = None 

253 return _result 

254 

255 

256def collective_permute(input, source_target_pairs, name=None): 

257 r"""An Op to permute tensors across replicated TPU instances. 

258 

259 Each instance supplies its own input. 

260 

261 For example, suppose there are 4 TPU instances: `[A, B, C, D]`. Passing 

262 source_target_pairs=`[[0,1],[1,2],[2,3],[3,0]]` gets the outputs: 

263 `[D, A, B, C]`. 

264 

265 Args: 

266 input: A `Tensor`. Must be one of the following types: `float32`, `float64`, `int32`, `uint8`, `int16`, `int8`, `complex64`, `int64`, `qint8`, `quint8`, `qint32`, `bfloat16`, `qint16`, `quint16`, `uint16`, `complex128`, `half`, `uint32`, `uint64`. 

267 The local input to be permuted. Currently only supports float and 

268 bfloat16. 

269 source_target_pairs: A `Tensor` of type `int32`. 

270 A tensor with shape [num_pairs, 2]. 

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

272 

273 Returns: 

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

275 """ 

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

277 tld = _ctx._thread_local_data 

278 if tld.is_eager: 

279 try: 

280 _result = pywrap_tfe.TFE_Py_FastPathExecute( 

281 _ctx, "CollectivePermute", name, input, source_target_pairs) 

282 return _result 

283 except _core._NotOkStatusException as e: 

284 _ops.raise_from_not_ok_status(e, name) 

285 except _core._FallbackException: 

286 pass 

287 try: 

288 return collective_permute_eager_fallback( 

289 input, source_target_pairs, name=name, ctx=_ctx) 

290 except _core._SymbolicException: 

291 pass # Add nodes to the TensorFlow graph. 

292 # Add nodes to the TensorFlow graph. 

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

294 "CollectivePermute", input=input, 

295 source_target_pairs=source_target_pairs, 

296 name=name) 

297 _result = _outputs[:] 

298 if _execute.must_record_gradient(): 

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

300 _inputs_flat = _op.inputs 

301 _execute.record_gradient( 

302 "CollectivePermute", _inputs_flat, _attrs, _result) 

303 _result, = _result 

304 return _result 

305 

306CollectivePermute = tf_export("raw_ops.CollectivePermute")(_ops.to_raw_op(collective_permute)) 

307 

308 

309def collective_permute_eager_fallback(input, source_target_pairs, name, ctx): 

310 _attr_T, (input,) = _execute.args_to_matching_eager([input], ctx, [_dtypes.float32, _dtypes.float64, _dtypes.int32, _dtypes.uint8, _dtypes.int16, _dtypes.int8, _dtypes.complex64, _dtypes.int64, _dtypes.qint8, _dtypes.quint8, _dtypes.qint32, _dtypes.bfloat16, _dtypes.qint16, _dtypes.quint16, _dtypes.uint16, _dtypes.complex128, _dtypes.half, _dtypes.uint32, _dtypes.uint64, ]) 

311 source_target_pairs = _ops.convert_to_tensor(source_target_pairs, _dtypes.int32) 

312 _inputs_flat = [input, source_target_pairs] 

313 _attrs = ("T", _attr_T) 

314 _result = _execute.execute(b"CollectivePermute", 1, inputs=_inputs_flat, 

315 attrs=_attrs, ctx=ctx, name=name) 

316 if _execute.must_record_gradient(): 

317 _execute.record_gradient( 

318 "CollectivePermute", _inputs_flat, _attrs, _result) 

319 _result, = _result 

320 return _result 

321 

322 

323def configure_distributed_tpu(embedding_config="", tpu_embedding_config="", is_global_init=False, enable_whole_mesh_compilations=False, compilation_failure_closes_chips=True, tpu_cancellation_closes_chips=0, name=None): 

324 r"""Sets up the centralized structures for a distributed TPU system. 

325 

326 Args: 

327 embedding_config: An optional `string`. Defaults to `""`. 

328 Reserved. Do not use. 

329 tpu_embedding_config: An optional `string`. Defaults to `""`. 

330 Serialized tensorflow.tpu.TPUEmbeddingConfiguration that 

331 describes the embedding lookups of the program. 

332 is_global_init: An optional `bool`. Defaults to `False`. 

333 Reserved. Do not use. 

334 enable_whole_mesh_compilations: An optional `bool`. Defaults to `False`. 

335 compilation_failure_closes_chips: An optional `bool`. Defaults to `True`. 

336 tpu_cancellation_closes_chips: An optional `int`. Defaults to `0`. 

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

338 

339 Returns: 

340 A `Tensor` of type `string`. 

341 """ 

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

343 tld = _ctx._thread_local_data 

344 if tld.is_eager: 

345 try: 

346 _result = pywrap_tfe.TFE_Py_FastPathExecute( 

347 _ctx, "ConfigureDistributedTPU", name, "embedding_config", 

348 embedding_config, "tpu_embedding_config", tpu_embedding_config, 

349 "is_global_init", is_global_init, "enable_whole_mesh_compilations", 

350 enable_whole_mesh_compilations, "compilation_failure_closes_chips", 

351 compilation_failure_closes_chips, "tpu_cancellation_closes_chips", 

352 tpu_cancellation_closes_chips) 

353 return _result 

354 except _core._NotOkStatusException as e: 

355 _ops.raise_from_not_ok_status(e, name) 

356 except _core._FallbackException: 

357 pass 

358 try: 

359 return configure_distributed_tpu_eager_fallback( 

360 embedding_config=embedding_config, 

361 tpu_embedding_config=tpu_embedding_config, 

362 is_global_init=is_global_init, 

363 enable_whole_mesh_compilations=enable_whole_mesh_compilations, 

364 compilation_failure_closes_chips=compilation_failure_closes_chips, 

365 tpu_cancellation_closes_chips=tpu_cancellation_closes_chips, 

366 name=name, ctx=_ctx) 

367 except _core._SymbolicException: 

368 pass # Add nodes to the TensorFlow graph. 

369 # Add nodes to the TensorFlow graph. 

370 if embedding_config is None: 

371 embedding_config = "" 

372 embedding_config = _execute.make_str(embedding_config, "embedding_config") 

373 if tpu_embedding_config is None: 

374 tpu_embedding_config = "" 

375 tpu_embedding_config = _execute.make_str(tpu_embedding_config, "tpu_embedding_config") 

376 if is_global_init is None: 

377 is_global_init = False 

378 is_global_init = _execute.make_bool(is_global_init, "is_global_init") 

379 if enable_whole_mesh_compilations is None: 

380 enable_whole_mesh_compilations = False 

381 enable_whole_mesh_compilations = _execute.make_bool(enable_whole_mesh_compilations, "enable_whole_mesh_compilations") 

382 if compilation_failure_closes_chips is None: 

383 compilation_failure_closes_chips = True 

384 compilation_failure_closes_chips = _execute.make_bool(compilation_failure_closes_chips, "compilation_failure_closes_chips") 

385 if tpu_cancellation_closes_chips is None: 

386 tpu_cancellation_closes_chips = 0 

387 tpu_cancellation_closes_chips = _execute.make_int(tpu_cancellation_closes_chips, "tpu_cancellation_closes_chips") 

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

389 "ConfigureDistributedTPU", embedding_config=embedding_config, 

390 tpu_embedding_config=tpu_embedding_config, 

391 is_global_init=is_global_init, 

392 enable_whole_mesh_compilations=enable_whole_mesh_compilations, 

393 compilation_failure_closes_chips=compilation_failure_closes_chips, 

394 tpu_cancellation_closes_chips=tpu_cancellation_closes_chips, 

395 name=name) 

396 _result = _outputs[:] 

397 if _execute.must_record_gradient(): 

398 _attrs = ("embedding_config", _op.get_attr("embedding_config"), 

399 "tpu_embedding_config", _op.get_attr("tpu_embedding_config"), 

400 "is_global_init", _op._get_attr_bool("is_global_init"), 

401 "enable_whole_mesh_compilations", 

402 _op._get_attr_bool("enable_whole_mesh_compilations"), 

403 "compilation_failure_closes_chips", 

404 _op._get_attr_bool("compilation_failure_closes_chips"), 

405 "tpu_cancellation_closes_chips", 

406 _op._get_attr_int("tpu_cancellation_closes_chips")) 

407 _inputs_flat = _op.inputs 

408 _execute.record_gradient( 

409 "ConfigureDistributedTPU", _inputs_flat, _attrs, _result) 

410 _result, = _result 

411 return _result 

412 

413ConfigureDistributedTPU = tf_export("raw_ops.ConfigureDistributedTPU")(_ops.to_raw_op(configure_distributed_tpu)) 

414 

415 

416def configure_distributed_tpu_eager_fallback(embedding_config, tpu_embedding_config, is_global_init, enable_whole_mesh_compilations, compilation_failure_closes_chips, tpu_cancellation_closes_chips, name, ctx): 

417 if embedding_config is None: 

418 embedding_config = "" 

419 embedding_config = _execute.make_str(embedding_config, "embedding_config") 

420 if tpu_embedding_config is None: 

421 tpu_embedding_config = "" 

422 tpu_embedding_config = _execute.make_str(tpu_embedding_config, "tpu_embedding_config") 

423 if is_global_init is None: 

424 is_global_init = False 

425 is_global_init = _execute.make_bool(is_global_init, "is_global_init") 

426 if enable_whole_mesh_compilations is None: 

427 enable_whole_mesh_compilations = False 

428 enable_whole_mesh_compilations = _execute.make_bool(enable_whole_mesh_compilations, "enable_whole_mesh_compilations") 

429 if compilation_failure_closes_chips is None: 

430 compilation_failure_closes_chips = True 

431 compilation_failure_closes_chips = _execute.make_bool(compilation_failure_closes_chips, "compilation_failure_closes_chips") 

432 if tpu_cancellation_closes_chips is None: 

433 tpu_cancellation_closes_chips = 0 

434 tpu_cancellation_closes_chips = _execute.make_int(tpu_cancellation_closes_chips, "tpu_cancellation_closes_chips") 

435 _inputs_flat = [] 

436 _attrs = ("embedding_config", embedding_config, "tpu_embedding_config", 

437 tpu_embedding_config, "is_global_init", is_global_init, 

438 "enable_whole_mesh_compilations", enable_whole_mesh_compilations, 

439 "compilation_failure_closes_chips", compilation_failure_closes_chips, 

440 "tpu_cancellation_closes_chips", tpu_cancellation_closes_chips) 

441 _result = _execute.execute(b"ConfigureDistributedTPU", 1, 

442 inputs=_inputs_flat, attrs=_attrs, ctx=ctx, 

443 name=name) 

444 if _execute.must_record_gradient(): 

445 _execute.record_gradient( 

446 "ConfigureDistributedTPU", _inputs_flat, _attrs, _result) 

447 _result, = _result 

448 return _result 

449 

450 

451def configure_tpu_embedding(config, name=None): 

452 r"""Sets up TPUEmbedding in a distributed TPU system. 

453 

454 Args: 

455 config: A `string`. 

456 Serialized tensorflow.tpu.TPUEmbeddingConfiguration that 

457 describes the embedding lookups of the program. 

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

459 

460 Returns: 

461 The created Operation. 

462 """ 

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

464 tld = _ctx._thread_local_data 

465 if tld.is_eager: 

466 try: 

467 _result = pywrap_tfe.TFE_Py_FastPathExecute( 

468 _ctx, "ConfigureTPUEmbedding", name, "config", config) 

469 return _result 

470 except _core._NotOkStatusException as e: 

471 _ops.raise_from_not_ok_status(e, name) 

472 except _core._FallbackException: 

473 pass 

474 try: 

475 return configure_tpu_embedding_eager_fallback( 

476 config=config, name=name, ctx=_ctx) 

477 except _core._SymbolicException: 

478 pass # Add nodes to the TensorFlow graph. 

479 # Add nodes to the TensorFlow graph. 

480 config = _execute.make_str(config, "config") 

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

482 "ConfigureTPUEmbedding", config=config, name=name) 

483 return _op 

484ConfigureTPUEmbedding = tf_export("raw_ops.ConfigureTPUEmbedding")(_ops.to_raw_op(configure_tpu_embedding)) 

485 

486 

487def configure_tpu_embedding_eager_fallback(config, name, ctx): 

488 config = _execute.make_str(config, "config") 

489 _inputs_flat = [] 

490 _attrs = ("config", config) 

491 _result = _execute.execute(b"ConfigureTPUEmbedding", 0, inputs=_inputs_flat, 

492 attrs=_attrs, ctx=ctx, name=name) 

493 _result = None 

494 return _result 

495 

496 

497def cross_replica_sum(input, group_assignment, name=None): 

498 r"""An Op to sum inputs across replicated TPU instances. 

499 

500 Each instance supplies its own input. 

501 

502 For example, suppose there are 8 TPU instances: `[A, B, C, D, E, F, G, H]`. 

503 Passing group_assignment=`[[0,2,4,6],[1,3,5,7]]` sets `A, C, E, G` as group 0, 

504 and `B, D, F, H` as group 1. Thus we get the outputs: 

505 `[A+C+E+G, B+D+F+H, A+C+E+G, B+D+F+H, A+C+E+G, B+D+F+H, A+C+E+G, B+D+F+H]`. 

506 

507 Args: 

508 input: A `Tensor`. Must be one of the following types: `half`, `bfloat16`, `float32`, `float64`, `int32`, `uint32`. 

509 The local input to the sum. 

510 group_assignment: A `Tensor` of type `int32`. An int32 tensor with shape 

511 [num_groups, num_replicas_per_group]. `group_assignment[i]` represents the 

512 replica ids in the ith subgroup. 

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

514 

515 Returns: 

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

517 """ 

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

519 tld = _ctx._thread_local_data 

520 if tld.is_eager: 

521 try: 

522 _result = pywrap_tfe.TFE_Py_FastPathExecute( 

523 _ctx, "CrossReplicaSum", name, input, group_assignment) 

524 return _result 

525 except _core._NotOkStatusException as e: 

526 _ops.raise_from_not_ok_status(e, name) 

527 except _core._FallbackException: 

528 pass 

529 try: 

530 return cross_replica_sum_eager_fallback( 

531 input, group_assignment, name=name, ctx=_ctx) 

532 except _core._SymbolicException: 

533 pass # Add nodes to the TensorFlow graph. 

534 # Add nodes to the TensorFlow graph. 

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

536 "CrossReplicaSum", input=input, group_assignment=group_assignment, 

537 name=name) 

538 _result = _outputs[:] 

539 if _execute.must_record_gradient(): 

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

541 _inputs_flat = _op.inputs 

542 _execute.record_gradient( 

543 "CrossReplicaSum", _inputs_flat, _attrs, _result) 

544 _result, = _result 

545 return _result 

546 

547CrossReplicaSum = tf_export("raw_ops.CrossReplicaSum")(_ops.to_raw_op(cross_replica_sum)) 

548 

549 

550def cross_replica_sum_eager_fallback(input, group_assignment, name, ctx): 

551 _attr_T, (input,) = _execute.args_to_matching_eager([input], ctx, [_dtypes.half, _dtypes.bfloat16, _dtypes.float32, _dtypes.float64, _dtypes.int32, _dtypes.uint32, ]) 

552 group_assignment = _ops.convert_to_tensor(group_assignment, _dtypes.int32) 

553 _inputs_flat = [input, group_assignment] 

554 _attrs = ("T", _attr_T) 

555 _result = _execute.execute(b"CrossReplicaSum", 1, inputs=_inputs_flat, 

556 attrs=_attrs, ctx=ctx, name=name) 

557 if _execute.must_record_gradient(): 

558 _execute.record_gradient( 

559 "CrossReplicaSum", _inputs_flat, _attrs, _result) 

560 _result, = _result 

561 return _result 

562 

563 

564def dynamic_enqueue_tpu_embedding_arbitrary_tensor_batch(sample_indices_or_row_splits, embedding_indices, aggregation_weights, mode_override, device_ordinal, combiners=[], name=None): 

565 r"""Eases the porting of code that uses tf.nn.embedding_lookup_sparse(). 

566 

567 embedding_indices[i] and aggregation_weights[i] correspond 

568 to the ith feature. 

569 

570 The tensors at corresponding positions in the three input lists (sample_indices, 

571 embedding_indices and aggregation_weights) must have the same shape, i.e. rank 1 

572 with dim_size() equal to the total number of lookups into the table described by 

573 the corresponding feature. 

574 

575 Args: 

576 sample_indices_or_row_splits: A list of at least 1 `Tensor` objects with the same type in: `int32`, `int64`. 

577 A list of rank 2 Tensors specifying the training example to which the 

578 corresponding embedding_indices and aggregation_weights values belong. 

579 If the size of its first dimension is 0, we assume each embedding_indices 

580 belongs to a different sample. Both int32 and int64 are allowed and will 

581 be converted to int32 internally. 

582 

583 Or a list of rank 1 Tensors specifying the row splits for splitting 

584 embedding_indices and aggregation_weights into rows. It corresponds to 

585 ids.row_splits in embedding_lookup(), when ids is a RaggedTensor. When 

586 enqueuing N-D ragged tensor, only the last dimension is allowed to be ragged. 

587 the row splits is 1-D dense tensor. When empty, we assume a dense tensor is 

588 passed to the op Both int32 and int64 are allowed and will be converted to 

589 int32 internally. 

590 embedding_indices: A list with the same length as `sample_indices_or_row_splits` of `Tensor` objects with the same type in: `int32`, `int64`. 

591 A list of rank 1 Tensors, indices into the embedding 

592 tables. Both int32 and int64 are allowed and will be converted to 

593 int32 internally. 

594 aggregation_weights: A list with the same length as `sample_indices_or_row_splits` of `Tensor` objects with the same type in: `float32`, `float64`. 

595 A list of rank 1 Tensors containing per training 

596 example aggregation weights. Both float32 and float64 are allowed and will 

597 be converted to float32 internally. 

598 mode_override: A `Tensor` of type `string`. 

599 A string input that overrides the mode specified in the 

600 TPUEmbeddingConfiguration. Supported values are {'unspecified', 'inference', 

601 'training', 'backward_pass_only'}. When set to 'unspecified', the mode set 

602 in TPUEmbeddingConfiguration is used, otherwise mode_override is used. 

603 device_ordinal: A `Tensor` of type `int32`. 

604 The TPU device to use. Should be >= 0 and less than the number 

605 of TPU cores in the task on which the node is placed. 

606 combiners: An optional list of `strings`. Defaults to `[]`. 

607 A list of string scalars, one for each embedding table that specify 

608 how to normalize the embedding activations after weighted summation. 

609 Supported combiners are 'mean', 'sum', or 'sqrtn'. It is invalid to have 

610 the sum of the weights be 0 for 'mean' or the sum of the squared weights be 

611 0 for 'sqrtn'. If combiners isn't passed, the default is to use 'sum' for 

612 all tables. 

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

614 

615 Returns: 

616 The created Operation. 

617 """ 

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

619 tld = _ctx._thread_local_data 

620 if tld.is_eager: 

621 try: 

622 _result = pywrap_tfe.TFE_Py_FastPathExecute( 

623 _ctx, "DynamicEnqueueTPUEmbeddingArbitraryTensorBatch", name, 

624 sample_indices_or_row_splits, embedding_indices, aggregation_weights, 

625 mode_override, device_ordinal, "combiners", combiners) 

626 return _result 

627 except _core._NotOkStatusException as e: 

628 _ops.raise_from_not_ok_status(e, name) 

629 except _core._FallbackException: 

630 pass 

631 try: 

632 return dynamic_enqueue_tpu_embedding_arbitrary_tensor_batch_eager_fallback( 

633 sample_indices_or_row_splits, embedding_indices, 

634 aggregation_weights, mode_override, device_ordinal, 

635 combiners=combiners, name=name, ctx=_ctx) 

636 except _core._SymbolicException: 

637 pass # Add nodes to the TensorFlow graph. 

638 # Add nodes to the TensorFlow graph. 

639 if not isinstance(sample_indices_or_row_splits, (list, tuple)): 

640 raise TypeError( 

641 "Expected list for 'sample_indices_or_row_splits' argument to " 

642 "'dynamic_enqueue_tpu_embedding_arbitrary_tensor_batch' Op, not %r." % sample_indices_or_row_splits) 

643 _attr_N = len(sample_indices_or_row_splits) 

644 if not isinstance(embedding_indices, (list, tuple)): 

645 raise TypeError( 

646 "Expected list for 'embedding_indices' argument to " 

647 "'dynamic_enqueue_tpu_embedding_arbitrary_tensor_batch' Op, not %r." % embedding_indices) 

648 if len(embedding_indices) != _attr_N: 

649 raise ValueError( 

650 "List argument 'embedding_indices' to 'dynamic_enqueue_tpu_embedding_arbitrary_tensor_batch' Op with length %d " 

651 "must match length %d of argument 'sample_indices_or_row_splits'." % 

652 (len(embedding_indices), _attr_N)) 

653 if not isinstance(aggregation_weights, (list, tuple)): 

654 raise TypeError( 

655 "Expected list for 'aggregation_weights' argument to " 

656 "'dynamic_enqueue_tpu_embedding_arbitrary_tensor_batch' Op, not %r." % aggregation_weights) 

657 if len(aggregation_weights) != _attr_N: 

658 raise ValueError( 

659 "List argument 'aggregation_weights' to 'dynamic_enqueue_tpu_embedding_arbitrary_tensor_batch' Op with length %d " 

660 "must match length %d of argument 'sample_indices_or_row_splits'." % 

661 (len(aggregation_weights), _attr_N)) 

662 if combiners is None: 

663 combiners = [] 

664 if not isinstance(combiners, (list, tuple)): 

665 raise TypeError( 

666 "Expected list for 'combiners' argument to " 

667 "'dynamic_enqueue_tpu_embedding_arbitrary_tensor_batch' Op, not %r." % combiners) 

668 combiners = [_execute.make_str(_s, "combiners") for _s in combiners] 

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

670 "DynamicEnqueueTPUEmbeddingArbitraryTensorBatch", sample_indices_or_row_splits=sample_indices_or_row_splits, 

671 embedding_indices=embedding_indices, 

672 aggregation_weights=aggregation_weights, 

673 mode_override=mode_override, 

674 device_ordinal=device_ordinal, 

675 combiners=combiners, 

676 name=name) 

677 return _op 

678DynamicEnqueueTPUEmbeddingArbitraryTensorBatch = tf_export("raw_ops.DynamicEnqueueTPUEmbeddingArbitraryTensorBatch")(_ops.to_raw_op(dynamic_enqueue_tpu_embedding_arbitrary_tensor_batch)) 

679 

680 

681def dynamic_enqueue_tpu_embedding_arbitrary_tensor_batch_eager_fallback(sample_indices_or_row_splits, embedding_indices, aggregation_weights, mode_override, device_ordinal, combiners, name, ctx): 

682 if not isinstance(sample_indices_or_row_splits, (list, tuple)): 

683 raise TypeError( 

684 "Expected list for 'sample_indices_or_row_splits' argument to " 

685 "'dynamic_enqueue_tpu_embedding_arbitrary_tensor_batch' Op, not %r." % sample_indices_or_row_splits) 

686 _attr_N = len(sample_indices_or_row_splits) 

687 if not isinstance(embedding_indices, (list, tuple)): 

688 raise TypeError( 

689 "Expected list for 'embedding_indices' argument to " 

690 "'dynamic_enqueue_tpu_embedding_arbitrary_tensor_batch' Op, not %r." % embedding_indices) 

691 if len(embedding_indices) != _attr_N: 

692 raise ValueError( 

693 "List argument 'embedding_indices' to 'dynamic_enqueue_tpu_embedding_arbitrary_tensor_batch' Op with length %d " 

694 "must match length %d of argument 'sample_indices_or_row_splits'." % 

695 (len(embedding_indices), _attr_N)) 

696 if not isinstance(aggregation_weights, (list, tuple)): 

697 raise TypeError( 

698 "Expected list for 'aggregation_weights' argument to " 

699 "'dynamic_enqueue_tpu_embedding_arbitrary_tensor_batch' Op, not %r." % aggregation_weights) 

700 if len(aggregation_weights) != _attr_N: 

701 raise ValueError( 

702 "List argument 'aggregation_weights' to 'dynamic_enqueue_tpu_embedding_arbitrary_tensor_batch' Op with length %d " 

703 "must match length %d of argument 'sample_indices_or_row_splits'." % 

704 (len(aggregation_weights), _attr_N)) 

705 if combiners is None: 

706 combiners = [] 

707 if not isinstance(combiners, (list, tuple)): 

708 raise TypeError( 

709 "Expected list for 'combiners' argument to " 

710 "'dynamic_enqueue_tpu_embedding_arbitrary_tensor_batch' Op, not %r." % combiners) 

711 combiners = [_execute.make_str(_s, "combiners") for _s in combiners] 

712 _attr_T1, sample_indices_or_row_splits = _execute.args_to_matching_eager(list(sample_indices_or_row_splits), ctx, [_dtypes.int32, _dtypes.int64, ], _dtypes.int32) 

713 _attr_T2, embedding_indices = _execute.args_to_matching_eager(list(embedding_indices), ctx, [_dtypes.int32, _dtypes.int64, ], _dtypes.int32) 

714 _attr_T3, aggregation_weights = _execute.args_to_matching_eager(list(aggregation_weights), ctx, [_dtypes.float32, _dtypes.float64, ], _dtypes.float32) 

715 mode_override = _ops.convert_to_tensor(mode_override, _dtypes.string) 

716 device_ordinal = _ops.convert_to_tensor(device_ordinal, _dtypes.int32) 

717 _inputs_flat = list(sample_indices_or_row_splits) + list(embedding_indices) + list(aggregation_weights) + [mode_override, device_ordinal] 

718 _attrs = ("T1", _attr_T1, "T2", _attr_T2, "T3", _attr_T3, "N", _attr_N, 

719 "combiners", combiners) 

720 _result = _execute.execute(b"DynamicEnqueueTPUEmbeddingArbitraryTensorBatch", 

721 0, inputs=_inputs_flat, attrs=_attrs, ctx=ctx, 

722 name=name) 

723 _result = None 

724 return _result 

725 

726 

727def enqueue_tpu_embedding_arbitrary_tensor_batch(sample_indices_or_row_splits, embedding_indices, aggregation_weights, mode_override, device_ordinal=-1, combiners=[], name=None): 

728 r"""Eases the porting of code that uses tf.nn.embedding_lookup_sparse(). 

729 

730 embedding_indices[i] and aggregation_weights[i] correspond 

731 to the ith feature. 

732 

733 The tensors at corresponding positions in the three input lists (sample_indices, 

734 embedding_indices and aggregation_weights) must have the same shape, i.e. rank 1 

735 with dim_size() equal to the total number of lookups into the table described by 

736 the corresponding feature. 

737 

738 Args: 

739 sample_indices_or_row_splits: A list of at least 1 `Tensor` objects with the same type in: `int32`, `int64`. 

740 A list of rank 2 Tensors specifying the training example to which the 

741 corresponding embedding_indices and aggregation_weights values belong. 

742 If the size of its first dimension is 0, we assume each embedding_indices 

743 belongs to a different sample. Both int32 and int64 are allowed and will 

744 be converted to int32 internally. 

745 

746 Or a list of rank 1 Tensors specifying the row splits for splitting 

747 embedding_indices and aggregation_weights into rows. It corresponds to 

748 ids.row_splits in embedding_lookup(), when ids is a RaggedTensor. When 

749 enqueuing N-D ragged tensor, only the last dimension is allowed to be ragged. 

750 the row splits is 1-D dense tensor. When empty, we assume a dense tensor is 

751 passed to the op Both int32 and int64 are allowed and will be converted to 

752 int32 internally. 

753 embedding_indices: A list with the same length as `sample_indices_or_row_splits` of `Tensor` objects with the same type in: `int32`, `int64`. 

754 A list of rank 1 Tensors, indices into the embedding 

755 tables. Both int32 and int64 are allowed and will be converted to 

756 int32 internally. 

757 aggregation_weights: A list with the same length as `sample_indices_or_row_splits` of `Tensor` objects with the same type in: `float32`, `float64`. 

758 A list of rank 1 Tensors containing per training 

759 example aggregation weights. Both float32 and float64 are allowed and will 

760 be converted to float32 internally. 

761 mode_override: A `Tensor` of type `string`. 

762 A string input that overrides the mode specified in the 

763 TPUEmbeddingConfiguration. Supported values are {'unspecified', 'inference', 

764 'training', 'backward_pass_only'}. When set to 'unspecified', the mode set 

765 in TPUEmbeddingConfiguration is used, otherwise mode_override is used. 

766 device_ordinal: An optional `int`. Defaults to `-1`. 

767 The TPU device to use. Should be >= 0 and less than the number 

768 of TPU cores in the task on which the node is placed. 

769 combiners: An optional list of `strings`. Defaults to `[]`. 

770 A list of string scalars, one for each embedding table that specify 

771 how to normalize the embedding activations after weighted summation. 

772 Supported combiners are 'mean', 'sum', or 'sqrtn'. It is invalid to have 

773 the sum of the weights be 0 for 'mean' or the sum of the squared weights be 

774 0 for 'sqrtn'. If combiners isn't passed, the default is to use 'sum' for 

775 all tables. 

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

777 

778 Returns: 

779 The created Operation. 

780 """ 

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

782 tld = _ctx._thread_local_data 

783 if tld.is_eager: 

784 try: 

785 _result = pywrap_tfe.TFE_Py_FastPathExecute( 

786 _ctx, "EnqueueTPUEmbeddingArbitraryTensorBatch", name, 

787 sample_indices_or_row_splits, embedding_indices, aggregation_weights, 

788 mode_override, "device_ordinal", device_ordinal, "combiners", 

789 combiners) 

790 return _result 

791 except _core._NotOkStatusException as e: 

792 _ops.raise_from_not_ok_status(e, name) 

793 except _core._FallbackException: 

794 pass 

795 try: 

796 return enqueue_tpu_embedding_arbitrary_tensor_batch_eager_fallback( 

797 sample_indices_or_row_splits, embedding_indices, 

798 aggregation_weights, mode_override, device_ordinal=device_ordinal, 

799 combiners=combiners, name=name, ctx=_ctx) 

800 except _core._SymbolicException: 

801 pass # Add nodes to the TensorFlow graph. 

802 # Add nodes to the TensorFlow graph. 

803 if not isinstance(sample_indices_or_row_splits, (list, tuple)): 

804 raise TypeError( 

805 "Expected list for 'sample_indices_or_row_splits' argument to " 

806 "'enqueue_tpu_embedding_arbitrary_tensor_batch' Op, not %r." % sample_indices_or_row_splits) 

807 _attr_N = len(sample_indices_or_row_splits) 

808 if not isinstance(embedding_indices, (list, tuple)): 

809 raise TypeError( 

810 "Expected list for 'embedding_indices' argument to " 

811 "'enqueue_tpu_embedding_arbitrary_tensor_batch' Op, not %r." % embedding_indices) 

812 if len(embedding_indices) != _attr_N: 

813 raise ValueError( 

814 "List argument 'embedding_indices' to 'enqueue_tpu_embedding_arbitrary_tensor_batch' Op with length %d " 

815 "must match length %d of argument 'sample_indices_or_row_splits'." % 

816 (len(embedding_indices), _attr_N)) 

817 if not isinstance(aggregation_weights, (list, tuple)): 

818 raise TypeError( 

819 "Expected list for 'aggregation_weights' argument to " 

820 "'enqueue_tpu_embedding_arbitrary_tensor_batch' Op, not %r." % aggregation_weights) 

821 if len(aggregation_weights) != _attr_N: 

822 raise ValueError( 

823 "List argument 'aggregation_weights' to 'enqueue_tpu_embedding_arbitrary_tensor_batch' Op with length %d " 

824 "must match length %d of argument 'sample_indices_or_row_splits'." % 

825 (len(aggregation_weights), _attr_N))