0 | module NN.Architectures.Softargmax
 1 |
 2 | import Data.Tensor
 3 | import Data.Para
 4 |
 5 | import public Data.Tensor.Softargmax
 6 |
 7 | ||| Softargmax as a parametric map, with temperature as a parameter
 8 | ||| The underlying implementation lives in `Data.Tensor.Softargmax`
 9 | ||| TODO since distribution is an applicative functor (https://glaive-research.org/2025/02/11/Generalized-Transformers-from-Applicative-Functors.html)
10 | ||| is there a meaningful notion of the "distribution container"?
11 | ||| Is there a sense in which `Dist` is a functor on containers?
12 | public export
13 | softargmax : {i : Axis} ->
14 |   {a : Type} -> Fractional a => Exp a => Ord a => Neg a =>
15 |   IsFoldable i.cont =>
16 |   (allAlg : AllAlgebra [i] a) =>
17 |   Tensor [i] a -\-> Tensor [i] a
18 | softargmax = MkPara 
19 |   (\_ => a) -- temperature is the parameter
20 |   (\(t ** temperature) => softargmaxImpl {temperature} t)
21 |