Store Sparse tensor, if necessary.
(t, shared_map_op)
| 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( |
| 536 | keep_input, |
| 537 | lambda: _store_many_sparse(t, shared_name=map_op_name), |
| 538 | lambda: -1 * array_ops.ones(array_ops.shape(t)[0:1], dtypes.int64)) |
| 539 | out_tensor.set_shape([None]) # necessary when t.ndims is unknown |
| 540 | return out_tensor |
| 541 | def _sparse_values_to_keep(t, keep_input): |
| 542 | """Convert a per-row `keep_input` vector to a per-value one.""" |
| 543 | # Get the rows of every value in the sparse Tensor. |
| 544 | row_values = t.indices[:, 0] |
| 545 | # The value should be kept iff the row should be kept. |
| 546 | return array_ops.gather(keep_input, row_values) |
| 547 | if keep_input.shape.ndims == 1: |
| 548 | t = sparse_ops.sparse_retain(t, _sparse_values_to_keep(t, keep_input)) |
| 549 | store_f = lambda t, name, _: _store_many_sparse(t, shared_name=name) |
| 550 | elif enqueue_many: |
| 551 | store_f = _maybe_store_many_sparse |
| 552 | else: |
| 553 | store_f = _maybe_store_sparse |
| 554 | return store_f(t, map_op_name, keep_input) |
| 555 | |
| 556 | stored_list = [ |
| 557 | _maybe_store(t, shared_map_op) for t, shared_map_op |
no test coverage detected