EXPERIMENTAL: A context manager for setting attributes on operators. This context manager can be used to add additional attributes to operators within the scope of the context. For example: with ops.Graph().as_default() as g: f_1 = Foo() # No extra attributes
(self, attr_map)
| 4752 | # pylint: disable=g-doc-return-or-yield |
| 4753 | @tf_contextlib.contextmanager |
| 4754 | def _attr_scope(self, attr_map): |
| 4755 | """EXPERIMENTAL: A context manager for setting attributes on operators. |
| 4756 | |
| 4757 | This context manager can be used to add additional |
| 4758 | attributes to operators within the scope of the context. |
| 4759 | |
| 4760 | For example: |
| 4761 | |
| 4762 | with ops.Graph().as_default() as g: |
| 4763 | f_1 = Foo() # No extra attributes |
| 4764 | with g._attr_scope({"_a": tf.attr_value_pb2.AttrValue(b=False)}): |
| 4765 | f_2 = Foo() # Additional attribute _a=False |
| 4766 | with g._attr_scope({"_a": tf.attr_value_pb2.AttrValue(b=True)}): |
| 4767 | f_3 = Foo() # Additional attribute _a=False |
| 4768 | with g._attr_scope({"_a": None}): |
| 4769 | f_4 = Foo() # No additional attributes. |
| 4770 | |
| 4771 | Args: |
| 4772 | attr_map: A dictionary mapping attr name strings to AttrValue protocol |
| 4773 | buffers or None. |
| 4774 | |
| 4775 | Returns: |
| 4776 | A context manager that sets the kernel label to be used for one or more |
| 4777 | ops created in that context. |
| 4778 | |
| 4779 | Raises: |
| 4780 | TypeError: If attr_map is not a dictionary mapping |
| 4781 | strings to AttrValue protobufs. |
| 4782 | """ |
| 4783 | if not isinstance(attr_map, dict): |
| 4784 | raise TypeError("attr_map must be a dictionary mapping " |
| 4785 | "strings to AttrValue protocol buffers") |
| 4786 | # The saved_attrs dictionary stores any currently-set labels that |
| 4787 | # will be overridden by this context manager. |
| 4788 | saved_attrs = {} |
| 4789 | # Install the given attribute |
| 4790 | for name, attr in attr_map.items(): |
| 4791 | if not (isinstance(name, six.string_types) and |
| 4792 | (isinstance(attr, (type(None), attr_value_pb2.AttrValue)) or |
| 4793 | callable(attr))): |
| 4794 | raise TypeError("attr_map must be a dictionary mapping " |
| 4795 | "strings to AttrValue protocol buffers or " |
| 4796 | "callables that emit AttrValue protocol buffers") |
| 4797 | try: |
| 4798 | saved_attrs[name] = self._attr_scope_map[name] |
| 4799 | except KeyError: |
| 4800 | pass |
| 4801 | if attr is None: |
| 4802 | del self._attr_scope_map[name] |
| 4803 | else: |
| 4804 | self._attr_scope_map[name] = attr |
| 4805 | try: |
| 4806 | yield # The code within the context runs here. |
| 4807 | finally: |
| 4808 | # Remove the attributes set for this context, and restore any saved |
| 4809 | # attributes. |
| 4810 | for name, attr in attr_map.items(): |
| 4811 | try: |