dgf.train

dgf.train.EmbedNodesetFeaturesModule

Bases: Module

A FLAX module to transform a set of features into a fixed-size embedding.

This module can be used before the message passing stage.

This method is best applied on normalized features value computed with dgf.transform.auto_normalize. It only supports and applies the following transformations depending on the semantic+format of the feature: - semantic=EMBEDDING format=float32 => Identity - semantic=CATEGORICAL format=int64 => Create embedding table

Usage example:

# Your GNN model
class GNNModel(nn.Module):
  schema: dgf.data.GraphSchema
  @nn.compact
  def __call__(
      self,
      graph: dgf.data.JaxInMemoryGraph,
      training: bool,
      seed_node_idxs: jax.Array,
  ) -> jax.Array:
    # Ingest the features
    node_embeddings =
    EmbedNodesetFeaturesModule(
          schema=schema,ignore_features=[("target_nodeset","label")])(
              graph,training=training)

    # Message passing
    for _ in range(num_rounds):
      node_embeddings = message_passing(node_embeddings, ...)

    # Result
    return node_embeddings["target_nodeset"]

If using Sparse Deferred (or another JAX GNN library), you can combine this method with dgf.convert.sparse_deferred_struct_to_graph.

Attributes:

Name	Type	Description
schema	GraphSchema	The graph schema.
ignore_features	Sequence[Tuple[str, str]]	A sequence of (nodeset_name, feature_name) tuples specifying features to ignore during ingestion.
categorical_feature_embedding_dim	int	The embedding dimension used for categorical features.