Module KaunSource

Sourcetype ('layout, 'dev) tensor = (float, 'layout, 'dev) Rune.t
Sourcetype 'layout dtype = (float, 'layout) Rune.dtype
Sourcetype 'dev device = 'dev Rune.device
Sourcetype ('layout, 'dev) params =
  1. | Tensor of ('layout, 'dev) tensor
  2. | List of ('layout, 'dev) params list
  3. | Record of (string * ('layout, 'dev) params) list
Sourcemodule Rngs : sig ... end
Sourcetype model =
  1. | Model : {
    1. init : 'layout 'dev. rngs:Rngs.t -> ('layout, 'dev) tensor -> ('layout, 'dev) params;
    2. apply : 'layout 'dev. ('layout, 'dev) params -> training:bool -> ?rngs:Rngs.t -> ('layout, 'dev) tensor -> ('layout, 'dev) tensor;
    } -> model
Sourceval init : model -> rngs:Rngs.t -> ('layout, 'dev) tensor -> ('layout, 'dev) params
Sourceval apply : model -> ('layout, 'dev) params -> training:bool -> ?rngs:Rngs.t -> ('layout, 'dev) tensor -> ('layout, 'dev) tensor
Sourceval value_and_grad : (('layout, 'dev) params -> ('layout, 'dev) tensor) -> ('layout, 'dev) params -> ('layout, 'dev) tensor * ('layout, 'dev) params
Sourceval grad : (('layout, 'dev) params -> ('layout, 'dev) tensor) -> ('layout, 'dev) params -> ('layout, 'dev) params
Sourcemodule Metrics : sig ... end
Sourcemodule Dataset : sig ... end
Sourcemodule Loss : sig ... end
Sourcemodule Initializer : sig ... end
Sourcemodule Layer : sig ... end
Sourcemodule Optimizer : sig ... end