mlx_graphs.nn.Linear

mlx_graphs.nn.Linear#

class mlx_graphs.nn.Linear(input_dims: int, output_dims: int, bias: bool = True) → None[source]#

Bases: Linear

Linear layer with Xavier Glorot weight inititalization.

This Linear class inherits from mx.nn.Linear, but uses glorot initialization instead of the default initialization in mlx’s Linear.

Parameters:

  • input_dims (int) – Dimensionality of the input features

  • output_dims (int) – Dimensionality of the output features

  • bias (bool) – If set to False then the layer will not use a bias. Default is True.

__call__(*args: Any, **kwargs: Any) → Any#

Call self as a function.

Return type:

Any

Methods

Contents