Encoder

class Encoder(*args, **kwargs)

Bases: Module, ABC

Return type:

Any

Implementations

class IdentityEncoder(*args, **kwargs)

Bases: Encoder

Identity encoder that returns the input graph unchanged.

\[\tilde{x} = x\]

Return type:

Any

class MLPEncoder(*args, **kwargs)

Bases: Encoder

Encoder that applies class-specific Multi Layer Perceptrons.

\[\begin{align} &\forall c \in \mathcal{C}, \forall e \in \mathcal{E}^c_x, & \tilde{x}_e = \phi_\theta^c(x_e), \end{align}\]

where \(({\phi}_{\theta}^c)_{c\in C}\) is a set of class-specific MLPs.

Parameters:

• input_structure – Input graph structure.

• hidden_sizes – Hidden sizes of MLPs \(({\phi}_{\theta}^c)_{c\in C}\).

• activation – Activation functions of MLPs \(({\phi}_{\theta}^c)_{c\in C}\).

• out_size – Output size of MLPs \(({\phi}_{\theta}^c)_{c\in C}\).

• use_bias – Whether to use bias in MLPs \(({\phi}_{\theta}^c)_{c\in C}\).

• kernel_init – Kernel initializer for MLPs \(({\phi}_{\theta}^c)_{c\in C}\).

• bias_init – Bias initializer for MLPs \(({\phi}_{\theta}^c)_{c\in C}\).

• final_activation – Final activation function for MLPs \(({\phi}_{\theta}^c)_{c\in C}\).

• dtype – Computation dtype of the MLPs (e.g. jnp.bfloat16 for mixed precision). Parameters are stored in float32 regardless, and the output is cast back to the input features dtype. None (default) keeps the computation in the input dtype, i.e. full float32.

• seed – Seed for RNG streams for weight initialization.

Return type:

Any