How is __cat_dim__ calculated?

[davidireland3]

Following on from my question the other day, I have made a custom data class:

class DynamicsData(Data):
    def __init__(self, x=None, edge_index=None, edge_index_blocks=None, edge_index_obstacles=None, neg_block_ids=None, neg_obstacle_ids=None, **kwargs):
        super().__init__(**kwargs)
        self.x = x
        self.edge_index = edge_index
        self.edge_index_blocks = edge_index_blocks
        self.edge_index_obstacles = edge_index_obstacles
        self.neg_block_ids = neg_block_ids
        self.neg_obstacle_ids = neg_obstacle_ids

    def __inc__(self, key, value, *args, **kwargs):
        if key in ['edge_index', 'edge_index_blocks', 'edge_index_obstacles', 'neg_block_ids', 'neg_obstacle_ids']:
            return self.x.size(0)
        else:
            return super().__inc__(key, value, *args, **kwargs)

The inputs to all the new attributes will have shape [2, N], as with the standard edge index. I initially tested with a DataLoader using only edge_index_blocks as the new attribute along with the original x and edge_index, and everything works fine. However, when using neg_block_ids I get a concatenation error. I tracked the problem down to the _collate function in collate.py, in particular this line. For some reason, it returns 0, when I want it to return -1 as with edge_index, edge_index_blocks, etc.

I can hard code a fix for this, but I was wondering for my own understanding why this is happening?

If it helps, I can attach/copy minimal code to reproduce the error.

[rusty1s]

class DynamicsData(Data):
    def __init__(self, x=None, edge_index=None, edge_index_blocks=None, edge_index_obstacles=None, neg_block_ids=None, neg_obstacle_ids=None, **kwargs):
        super().__init__(**kwargs)
        self.x = x
        self.edge_index = edge_index
        self.edge_index_blocks = edge_index_blocks
        self.edge_index_obstacles = edge_index_obstacles
        self.neg_block_ids = neg_block_ids
        self.neg_obstacle_ids = neg_obstacle_ids

    def __inc__(self, key, value, *args, **kwargs):
        if key in ['edge_index', 'edge_index_blocks', 'edge_index_obstacles', 'neg_block_ids', 'neg_obstacle_ids']:
            return self.x.size(0)
        else:
            return super().__inc__(key, value, *args, **kwargs)

    def __cat_dim__(self, key, value, *args, **kwargs):
        if key in ['edge_index', 'edge_index_blocks', 'edge_index_obstacles', 'neg_block_ids', 'neg_obstacle_ids']:
            return -1
        else:
            return super().__cat_dim__(key, value, *args, **kwargs)

should fix this.

[davidireland3]

Many thanks!