Store SparseTensors for feeding into batch, etc. If `shared_map_ops` is provided, the underlying `SparseTensorsMap` objects are reused (shared). This argument is useful for, e.g., `batch_join` where multiple enqueue operations write to the same Queue component, and another (dequeue) thread
(tensor_list, enqueue_many, keep_input,
shared_map_ops=None)
| 476 | |
| 477 | |
| 478 | def _store_sparse_tensors(tensor_list, enqueue_many, keep_input, |
| 479 | shared_map_ops=None): |
| 480 | """Store SparseTensors for feeding into batch, etc. |
| 481 | |
| 482 | If `shared_map_ops` is provided, the underlying `SparseTensorsMap` objects |
| 483 | are reused (shared). This argument is useful for, e.g., `batch_join` |
| 484 | where multiple enqueue operations write to the same Queue component, |
| 485 | and another (dequeue) thread reads from that same location and must then |
| 486 | restore the associated `SparseTensor` objects. In this case, the sparse |
| 487 | restore must have a single `SparseTensorMap` from which to read out the |
| 488 | handles; so a single `SparseTensorMap` must be shared for storing |
| 489 | across the multiple enqueue operations. This sharing is performed by |
| 490 | calling `_store_sparse_tensors` the first time with `shared_map_ops=None`, |
| 491 | and then in subsequent times with this value set to the list of `Operation` |
| 492 | objects created in the first call. |
| 493 | |
| 494 | Args: |
| 495 | tensor_list: List of `Tensor` and `SparseTensor` objects. |
| 496 | enqueue_many: Python `Boolean`. |
| 497 | keep_input: Must be a scalar bool Tensor (not a Python bool). If False, |
| 498 | don't store. |
| 499 | shared_map_ops: (optional) List of `Operation` objects from a previous |
| 500 | call to `_store_sparse_tensors`. If not `None`, the op types should be |
| 501 | one of `AddSparseToTensorsMap` or `AddManySparseToTensorsMap` in the |
| 502 | locations corresponding to `SparseTensors` in `tensor_list`. |
| 503 | |
| 504 | Returns: |
| 505 | A tuple `(stored_list, sparse_info_list)` where `stored_list` is a list |
| 506 | of `Tensor` objects (same length as `tensor_list`) and `sparse_info_list` |
| 507 | is a list of the same length of `_SparseMetaData` objects. |
| 508 | """ |
| 509 | maybe_shared_map_ops = shared_map_ops or [None] * len(tensor_list) |
| 510 | |
| 511 | def _sparse_meta_data(t, storing_op, map_op): |
| 512 | if not isinstance(t, sparse_tensor.SparseTensor): |
| 513 | return _SparseMetaData(False, None, None) |
| 514 | rank = t.dense_shape.shape.with_rank(1).dims[0] |
| 515 | if enqueue_many: |
| 516 | rank -= 1 |
| 517 | # If a shared map_op was provided, use that. Otherwise use the name of |
| 518 | # the operation used to store the SparseTensor. |
| 519 | return _SparseMetaData( |
| 520 | sparse=True, map_op=map_op or storing_op, rank=rank) |
| 521 | |
| 522 | def _maybe_store(t, shared_map_op): |
| 523 | """Store Sparse tensor, if necessary.""" |
| 524 | if not isinstance(t, sparse_tensor.SparseTensor): |
| 525 | return t |
| 526 | map_op_name = shared_map_op.name if shared_map_op else None |
| 527 | def _maybe_store_sparse(t, map_op_name, keep_input): |
| 528 | """Conditionally store a single sparse Tensor.""" |
| 529 | return utils.smart_cond( |
| 530 | keep_input, |
| 531 | lambda: _store_sparse(t, shared_name=map_op_name), |
| 532 | lambda: constant_op.constant(-1, dtypes.int64)) |
| 533 | def _maybe_store_many_sparse(t, map_op_name, keep_input): |
| 534 | """Conditionally store multiple sparse Tensors.""" |
| 535 | out_tensor = utils.smart_cond( |
no test coverage detected