Returns the list of colocation groups of the op.
(self)
| 1835 | return grouped_inputs |
| 1836 | |
| 1837 | def colocation_groups(self): |
| 1838 | """Returns the list of colocation groups of the op.""" |
| 1839 | default_colocation_group = [compat.as_bytes("loc:@%s" % self.name)] |
| 1840 | try: |
| 1841 | class_attr = self.get_attr("_class") |
| 1842 | except ValueError: |
| 1843 | # This op has no explicit colocation group, so it is itself its |
| 1844 | # own root of a colocation group. |
| 1845 | return default_colocation_group |
| 1846 | |
| 1847 | attr_groups = [ |
| 1848 | class_name for class_name in class_attr |
| 1849 | if class_name.startswith(b"loc:@") |
| 1850 | ] |
| 1851 | |
| 1852 | # If there are no colocation groups in the explicit _class field, |
| 1853 | # return the default colocation group. |
| 1854 | return attr_groups if attr_groups else default_colocation_group |
| 1855 | |
| 1856 | def values(self): |
| 1857 | """DEPRECATED: Use outputs.""" |