Module Tolk_nn.Linear

Affine transformation of the last axis.

type t = {
  1. weight : Tolk_frontend.Tensor.t;
    (*

    Shape (out_features, in_features).

    *)
  2. bias : Tolk_frontend.Tensor.t option;
    (*

    Shape (out_features), when present.

    *)
}

A linear layer.

val create : ?bias:bool -> int -> int -> t

create in_features out_features is a linear layer mapping the last axis from in_features to out_features. The weight and bias are initialised from a uniform distribution over [-b, b) with b = 1 / sqrt in_features, drawing from the Rand stream. bias (default true) selects whether an additive bias is included. Load trained values with State.load_state_dict.

apply l x is x @ weight^T + bias, contracting the last axis of x with in_features.