Attribute name when batching along new dimension

[AlexNiko47]

Hi,

I'm following the tutorial to overwrite the __cat_dim__() method so that I can batch along a new dimension. I get the expected behaviour however when I try to check whether my node attributes are recognized by get_all_tensor_attrs() I get the RuntimeError: Please look up dimensions by name, got: name = None. Here's my implementation:

class Overwrite(Data):
    def __cat_dim__(self, key: str, value: Any, *args, **kwargs) -> Any:
        if key == "ow":
            return None
        else:
            return super().__cat_dim__(key, value, *args, **kwargs)
    def __inc__(self, key: str, value: Any, *args, **kwargs) -> Any:
        if "num_nodes" in key:
            return self.num_nodes
        else:
            return 0

If I try creating a Data object using the base class I get the expected batching where the nodes are stacked in the 1st dimension:

x1 = torch.tensor(([0],[1],[2]), dtype=torch.float)
x2 = torch.tensor(([3],[4],[5]), dtype=torch.float)

edge_index = torch.tensor(([0,1,1,2],[1,0,2,1]), dtype=torch.long)

graph1 = Data(x=x1, edge_index=edge_index)
graph2 = Data(x=x2, edge_index=edge_index)

graph_list = [graph1,graph2]

loader = DataLoader(graph_list, batch_size=2)

print(next(iter(loader)))
print(graph1.get_all_tensor_attrs())

which returns DataBatch(x=[6, 1], edge_index=[2, 8], batch=[6], ptr=[3]) and [TensorAttr(group_name=<FieldStatus.UNSET: 1>, attr_name='x', index=<FieldStatus.UNSET: 1>)] respectively. However when I use the Overwrite class, ow is not recognized as a node attribute anymore:

graph1 = Overwrite(edge_index=edge_index, ow = x1, num_nodes = 3)
graph2 = Overwrite(edge_index=edge_index, ow = x2, num_nodes = 3)

graph_list = [graph1,graph2]

loader = DataLoader(graph_list, batch_size=2)

print(next(iter(loader)))
print(graph2.get_all_tensor_attrs())

this prints OverwriteBatch(edge_index=[2, 8], ow=[2, 3, 1], num_nodes=6, batch=[6], ptr=[3]) but throws the RuntimeError: Please look up dimensions by name, got: name = None. because of how storage.py checks for node attributes: if value.dim() == 0 or value.size(cat_dim) != num_nodes:, but here my cat_dim is None because of how I've implemented the class.

Is there a way to bypass this and make my class recognize the ow node features as node attributes?

[rusty1s]

cc @mananshah99. This might be a bug with inheritance in FeatureStore. Can you try to access data.keys() instead of get_all_tensor_attrs() to unblock yourself.

[AlexNiko47]

I can access the ow variable as a key:

print(graph2.ow) returns tensor([[3.], [4.], [5.]]) which is what I'm expecting to see. The ow key is also present if I call print(graph2.keys) (which returns ['num_nodes', 'edge_index', 'ow'] ) or print(graph2.__dict__) which returns {'_edge_attr_cls': <class 'torch_geometric.data.data.DataEdgeAttr'>, '_tensor_attr_cls': <class 'torch_geometric.data.data.DataTensorAttr'>, '_store': {'edge_index': tensor([[0, 1, 1, 2], [1, 0, 2, 1]]), 'ow': tensor([[3.], [4.], [5.]]), 'num_nodes': 3}}