Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

GraphStore: Data, HeteroData respect is_sorted #4922

Merged
merged 7 commits into from
Jul 7, 2022
Merged

GraphStore: Data, HeteroData respect is_sorted #4922

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
36f0651
init
mananshah99 Jul 6, 2022
cc9fe1f
update
mananshah99 Jul 6, 2022
cdde48f
Merge branch 'master' into gs_respect_is_sorted
mananshah99 Jul 6, 2022
b125c97
update
mananshah99 Jul 6, 2022
2f67fb8
Merge branch 'gs_respect_is_sorted' of github.com:pyg-team/pytorch_ge…
mananshah99 Jul 6, 2022
d51ee88
updates
mananshah99 Jul 7, 2022
0317cbc
Merge branch 'master' into gs_respect_is_sorted
mananshah99 Jul 7, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
updates
  • Loading branch information
@mananshah99
mananshah99 committed Jul 7, 2022
commit d51ee887eed2cd8227ebaf3b6952a9b764eb2a6c
21 changes: 11 additions & 10 deletions torch_geometric/data/data.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -843,7 +843,10 @@ def _put_edge_index(self, edge_index: EdgeTensorType,
setattr(self, attr_name, attr_val)

# Set edge attributes:
rusty1s marked this conversation as resolved.
Show resolved Hide resolved
setattr(self, f'{attr_name}_edge_attr', edge_attr)
if not hasattr(self, '_edge_attrs'):
self._edge_attrs = {}

self._edge_attrs[edge_attr.layout.value] = edge_attr

# Set size, if possible:
size = edge_attr.size
Expand All @@ -869,15 +872,13 @@ def _get_edge_index(self, edge_attr: EdgeAttr) -> Optional[EdgeTensorType]:
def get_all_edge_attrs(self) -> List[EdgeAttr]:
r"""Returns `EdgeAttr` objects corresponding to the edge indices stored
in `Data` and their layouts"""
out = []
for layout, attr_name in EDGE_LAYOUT_TO_ATTR_NAME.items():
if attr_name in self:
attr_val = self[f'{attr_name}_edge_attr']
out.append(
EdgeAttr(edge_type=None, layout=layout,
is_sorted=attr_val.is_sorted,
size=(self.num_nodes, self.num_nodes)))
return out
if not hasattr(self, '_edge_attrs'):
return []

edge_attrs = self._edge_attrs.values()
for attr in edge_attrs:
attr.size = (self.num_nodes, self.num_nodes)
return edge_attrs


###############################################################################
Expand Down
1 change: 0 additions & 1 deletion torch_geometric/data/graph_store.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -139,7 +139,6 @@ def _edge_to_layout(
attr: EdgeAttr,
layout: EdgeLayout,
) -> Tuple[Tensor, Tensor, OptTensor]:
print(attr)
from_tuple = self.get_edge_index(attr)

if layout == EdgeLayout.COO:
Expand Down
20 changes: 11 additions & 9 deletions torch_geometric/data/hetero_data.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -755,7 +755,11 @@ def _put_edge_index(self, edge_index: EdgeTensorType,
setattr(self[edge_attr.edge_type], attr_name, attr_val)

# Set edge attributes:
setattr(self[edge_attr.edge_type], f'{attr_name}_edge_attr', edge_attr)
if not hasattr(self[edge_attr.edge_type], '_edge_attrs'):
self[edge_attr.edge_type]._edge_attrs = {}

self[edge_attr.edge_type]._edge_attrs[
edge_attr.layout.value] = edge_attr

key = self._to_canonical(edge_attr.edge_type)
src, _, dst = key
Expand Down Expand Up @@ -784,14 +788,12 @@ def get_all_edge_attrs(self) -> List[EdgeAttr]:
indices stored in `HeteroData` and their layouts."""
out = []
for edge_type, edge_store in self.edge_items():
for layout, attr_name in EDGE_LAYOUT_TO_ATTR_NAME.items():
if attr_name in edge_store:
attr_val = getattr(self[edge_type],
f'{attr_name}_edge_attr')
out.append(
EdgeAttr(edge_type=edge_type, layout=layout,
is_sorted=attr_val.is_sorted,
size=self[edge_type].size()))
if not hasattr(self[edge_type], '_edge_attrs'):
continue
edge_attrs = self[edge_type]._edge_attrs.values()
for attr in edge_attrs:
attr.size = self[edge_type].size()
out.extend(edge_attrs)
return out


Expand Down