Helper to return (possibly cached) zero tensors in eager mode.
(shape, dtype)
| 624 | |
| 625 | |
| 626 | def _zeros(shape, dtype): |
| 627 | """Helper to return (possibly cached) zero tensors in eager mode.""" |
| 628 | if (dtype == dtypes.variant |
| 629 | or dtype == dtypes.string |
| 630 | or dtype == dtypes.resource): |
| 631 | # TODO(apassos): need to save enough information about variant tensors to do |
| 632 | # a zeros |
| 633 | return None |
| 634 | |
| 635 | ctx = context.context() |
| 636 | if not ctx.executing_eagerly(): |
| 637 | return array_ops.zeros(shape, dtype) |
| 638 | |
| 639 | device = ctx.device_name |
| 640 | |
| 641 | if tensor_util.is_tensor(shape): |
| 642 | shape_key = shape.experimental_ref() |
| 643 | else: |
| 644 | shape_key = shape |
| 645 | cache_key = shape_key, dtype, device |
| 646 | cached = ctx.zeros_cache().get(cache_key) |
| 647 | if cached is None: |
| 648 | if dtypes.as_dtype(dtype).is_bool: |
| 649 | value = False |
| 650 | else: |
| 651 | value = 0 |
| 652 | cached = _fast_fill(value, shape, dtype) |
| 653 | ctx.zeros_cache().put(cache_key, cached) |
| 654 | return cached |
| 655 | |
| 656 | |
| 657 | def _ones(shape, dtype): |
nothing calls this directly
no test coverage detected