Module KaunSource

Sourcetype 'layout tensor = (float, 'layout) Rune.t
Sourcetype 'layout dtype = (float, 'layout) Rune.dtype
Sourcetype params = Ptree.t =
  1. | Tensor of Ptree.tensor
  2. | List of params list
  3. | Dict of (string * params) list
Sourcetype module_ = Layer.module_ = {
  1. init : 'layout. rngs:Rune.Rng.key -> dtype:(float, 'layout) Rune.dtype -> Ptree.t;
  2. apply : 'layout. Ptree.t -> training:bool -> ?rngs:Rune.Rng.key -> (float, 'layout) Rune.t -> (float, 'layout) Rune.t;
}
Sourceval init : module_ -> rngs:Rune.Rng.key -> dtype:'layout dtype -> params
Sourceval apply : module_ -> params -> training:bool -> ?rngs:Rune.Rng.key -> 'layout tensor -> 'layout tensor
Sourceval value_and_grad : (params -> 'layout tensor) -> params -> 'layout tensor * params
Sourceval grad : (params -> 'layout tensor) -> params -> params
Sourcemodule Metrics : sig ... end

Performance metrics for neural network training and evaluation.

Sourcemodule Dataset : sig ... end

Efficient dataset handling for machine learning pipelines This module provides composable dataset transformations with support for:

Sourcemodule Loss : sig ... end

Loss functions for neural network training.

Sourcemodule Initializers : sig ... end

Flax-compatible weight initializers for Kaun

Sourcemodule Attention : sig ... end
Sourcemodule Layer : sig ... end

Neural network layer constructors.

module Checkpoint : sig ... end
Sourcemodule Train_state : sig ... end

Unified training state for Kaun.

Sourcemodule Ptree : sig ... end

Heterogeneous parameter tree structure.

Sourcemodule Optimizer : sig ... end

Optax-inspired optimisation algorithms with explicit checkpoint support.

Sourcemodule Activations : sig ... end

Activation functions for neural networks.

Sourcemodule Training : sig ... end

High-level training utilities operating on Train_state.