dicee.models.trix

Pure PyTorch TRIX entity and relation prediction.

Reimplementation of Zhang et al., https://arxiv.org/abs/2502.19512, verified against https://github.com/yuchengz99/TRIX at UPSTREAM_COMMIT. The module names, update schedule, binary entity-labelled relation edges, and fused convolution direction follow the released code/checkpoints (see docs/trix.md).

Attributes

UPSTREAM_COMMIT
INTERACTIONS

Classes

EntityReasoner	Base class for all neural network modules.
RelationReasoner	Base class for all neural network modules.
TRIXBase	Shared graph lifecycle and entity-scoring interfaces for graph foundation models.
TRIX	Entity predictor compatible with official entity_prediction.pth.
TRIXRelation	Relation predictor compatible with official relation_prediction.pth.

Functions

build_relation_graph(edge_index, edge_type, ...)

Sparse (relation, relation, shared entity) edges in official role order.

Module Contents

dicee.models.trix.UPSTREAM_COMMIT = '7596e14eefefe89e61396205a0550172cadeddb0'

dicee.models.trix.INTERACTIONS = ('hh', 'ht', 'th', 'tt')

dicee.models.trix.build_relation_graph(edge_index, edge_type, num_entities, num_relations)

Sparse (relation, relation, shared entity) edges in official role order.

Repeated incidences have binary support. hh/tt exclude equal relations; ht/th include them. Distinct shared entities remain distinct edges. No dense entity-by-relation-by-relation tensor is materialized.

class dicee.models.trix.EntityReasoner(dim, num_layers, output_dim=1)

Bases: torch.nn.Module

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:

training (bool) – Boolean represents whether this module is in training or evaluation mode.

layers

mlp

features(edges, num_entities, relations, heads, rels, states=None, tails=None, candidates=None, split=False)

class dicee.models.trix.RelationReasoner(dim, num_layers, entity_feedback=False)

Bases: torch.nn.Module

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:

training (bool) – Boolean represents whether this module is in training or evaluation mode.

step(index, states, boundary, graph, entities, constant_relations=None)

class dicee.models.trix.TRIXBase(args)

Bases: dicee.models.graph_model.GraphKGE

Shared graph lifecycle and entity-scoring interfaces for graph foundation models.

config_prefix = 'trix'

graph_filename = 'trix_graph.pt'

checkpoint_hint = 'use the matching entity/relation checkpoint with trix_dim=32; all keys and shapes must match'

dim

property relation_graph

class dicee.models.trix.TRIX(args)

Bases: TRIXBase

Entity predictor compatible with official entity_prediction.pth.

DICE order is (head, relation, tail). Head corruption retains the original relation seed, then uses the inverse relation for both entity reasoners. Explicit reciprocal DICE queries use that same convention.

name = 'TRIX'

deterministic_inference = True

relation_model

entity_model_1

entity_model_2

relation_cache_mb

clear_inference_cache()

Drop derived representations when graph, weights or device change.

class dicee.models.trix.TRIXRelation(args)

Bases: TRIXBase, dicee.models.graph_model.RelationGraphKGE

Relation predictor compatible with official relation_prediction.pth.

Pairs are (head, tail); triples and grouped relation corruptions use DICE’s (head, relation, tail) order. One reasoning pass scores every relation for a given pair, including external inverse IDs when present in the vocabulary.

name = 'TRIXRelation'

relation_model

entity_model

mlp