Package fnn implements a feed-forward neural network for the OGBN-MAG problem.

Package gnn implements a generic GNN modeling based on [TF-GNN MtAlbis].

BuildLayerWiseInferenceModel returns a function that builds the OGBN-MAG GNN inference model, that expects to run inference on the whole dataset in one go.

Download and prepares the tensors with the data into the `baseDir`.

ExcludeOgbnMagVariablesFromSave marks the OGBN-MAG variables as not to be saved by the given `checkpoint`.

ExtractLabelsFromInput create the labels from the input seed indices.

FeaturePreprocessing converts the `spec` and `inputs` given by the dataset into a map of node type name to its initial embeddings.

InitTrainingSchedule initializes custom scheduled training.

LayerWiseEvaluation returns the train, validation and test accuracy of the model, using layer-wise inference.

MagModelGraph builds a OGBN-MAG GNN model that sends [ParamNumGraphUpdates] along its sampling strategy, and then adding a final layer on top of the seeds.

MakeDatasets takes a directory where to store the downloaded data and return 4 datasets: "train", "trainEval", "validEval", "testEval".

NewSampler will create a [sampler.Sampler] and configure it with the OGBN-MAG graph definition.

NewSamplerStrategy creates a sampling strategy given the sampler, batch size and seeds candidates to sample from.

PapersSeedDatasets returns the train, validation and test datasets (`data.InMemoryDataset`) with only the papers seed nodes, to be used with FNN (Feedforward Neural Networks).

Train GNN model based on configuration in `ctx`.

TrainingSchedule is used to control hyperparameters during training.

UploadOgbnMagVariables creates frozen variables with the various static tables of the OGBN-MAG dataset, so it can be used by models.

BatchSize used for the sampler: the value was taken from the TF-GNN OGBN-MAG demo colab, and it was the best found with some hyperparameter tuning.

Counts to the various edge types.

Counts to the various edge types.

EdgesAffiliatedWith `(Int32)[1043998, 2]`, pairs with (author_id, institution_id).

EdgesCites `(Int32)[5416271, 2]`, pairs with (paper_id, paper_id).

EdgesHasTopic `(Int32)[7505078, 2]`, pairs with (paper_id, topic_id).

EdgesWrites `(Int32)[7145660, 2]`, pairs with (author_id, paper_id).

IdentitySubSeeds controls whether to use an IdentitySubSeed, to allow more sharing of the kernel.

KeepDegrees will also make sampler keep the degrees of the edges as separate tensors.

NumLabels is the number of labels for the papers.

OgbnMagVariablesRef maps variable names to a reference to their values.

OgbnMagVariablesScope is the absolute scope where the dataset variables are stored.

PaperEmbeddingsSize is the size of the node features given.

PapersEmbeddings contains the embeddings, shaped `(Float32)[NumPapers, PaperEmbeddingsSize]`.

PapersLabels for each paper, values from 0 to 348 (so 349 in total).

PapersYears for each paper, where year starts in 2000 (so 10 corresponds to 2010).

ParamDType controls the dtype to be used: either "float32" or "float16".

ParamEmbedDropoutRate adds an extra dropout to learning embeddings.

ParamIdentitySubSeeds controls whether to use an IdentitySubSeed, to allow more sharing of the kernel.

ParamNumCheckpoints is the number of past checkpoints to keep.

ParamReuseKernels context parameter configs whether the kernels for similar sampling rules will be reused.

ParamSplitEmbedTablesSize will make embed tables share entries across these many entries.

ReuseShareableKernels will share the kernels across similar messages in the strategy tree.

TrainSplit, ValidSplit, TestSplit splits of the data.

TrainSplit, ValidSplit, TestSplit splits of the data.

TrainSplit, ValidSplit, TestSplit splits of the data.

WithReplacement indicates whether the training dataset is created with replacement.