decagon-pytorch/src/icosagon/compile.py

#
# The goal of this module is to make Icosagon more efficient.
# It takes the nice Icosagon model architecture and tries to
# formulate it in terms of batch matrix multiplications instead
# of using Python for loops.
#

from .weights import init_glorot
from .input
import torch


class EncodeLayer(object):
    def __init__(self, num_relation_types, input_dim, output_dim):
        weights = [ init_glorot(input_dim, output_dim) \
            for _ in range(num_relation_types) ]
        weights = torch.cat(weights)


class Compiler(object):
    def __init__(self, data: Data, layer_dimensions: List[int] = [32, 64]) -> None:
        self.data = data
        self.layer_dimensions = layer_dimensions
        self.build()

    def build(self) -> None:
        for fam in data.relation_families:
            init_glorot(in_channels, out_channels)