-- SPDX-License-Identifier: MPL-2.0

{- | Explicit INT8 quantization and software FP16/BF16 matrix multiplication.
Integer GEMM accumulates exact integers, then applies scales; it rejects
signed-32-bit accumulator overflow. Half inputs round ties to even, with FP32
accumulation and one final conversion to Double. Mixed precision retains
Double master weights; unscale/overflow decisions happen before any update.
-}
module HMath.DeepLearning.Precision (
    Quantization (..),
    QuantizedTensor,
    quantize,
    dequantize,
    quantizedValues,
    quantizationScales,
    quantizationZeroPoints,
    int8Gemm,
    Precision (..),
    reducedGemm,
    LossScaler (..),
    initialLossScaler,
    scaleLoss,
    unscaleGradients,
    updateLossScaler,
    gradientNorm,
    clipGradients,
) where

import qualified Data.Foldable as Foldable (foldl')
import Data.Int (Int32, Int8)
import qualified HMath.Autodiff as AD
import HMath.Error
import qualified HMath.Numeric as N
import qualified HMath.Tensor as T

{- | Integer quantization policy: Symmetric around zero or Asymmetric with a learned zero
point.
-}
data Quantization = Symmetric | Asymmetric deriving (Quantization -> Quantization -> Bool
(Quantization -> Quantization -> Bool)
-> (Quantization -> Quantization -> Bool) -> Eq Quantization
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Quantization -> Quantization -> Bool
== :: Quantization -> Quantization -> Bool
$c/= :: Quantization -> Quantization -> Bool
/= :: Quantization -> Quantization -> Bool
Eq, Int -> Quantization -> ShowS
[Quantization] -> ShowS
Quantization -> String
(Int -> Quantization -> ShowS)
-> (Quantization -> String)
-> ([Quantization] -> ShowS)
-> Show Quantization
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Quantization -> ShowS
showsPrec :: Int -> Quantization -> ShowS
$cshow :: Quantization -> String
show :: Quantization -> String
$cshowList :: [Quantization] -> ShowS
showList :: [Quantization] -> ShowS
Show)

{- | Opaque Int8 tensor with shape,optional channel axis,scales,and zero points. Construct
through the checked quantization API.
-}
data QuantizedTensor = QuantizedTensor [Int] (Maybe Int) [Double] [Int] [Int8] deriving (QuantizedTensor -> QuantizedTensor -> Bool
(QuantizedTensor -> QuantizedTensor -> Bool)
-> (QuantizedTensor -> QuantizedTensor -> Bool)
-> Eq QuantizedTensor
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: QuantizedTensor -> QuantizedTensor -> Bool
== :: QuantizedTensor -> QuantizedTensor -> Bool
$c/= :: QuantizedTensor -> QuantizedTensor -> Bool
/= :: QuantizedTensor -> QuantizedTensor -> Bool
Eq, Int -> QuantizedTensor -> ShowS
[QuantizedTensor] -> ShowS
QuantizedTensor -> String
(Int -> QuantizedTensor -> ShowS)
-> (QuantizedTensor -> String)
-> ([QuantizedTensor] -> ShowS)
-> Show QuantizedTensor
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> QuantizedTensor -> ShowS
showsPrec :: Int -> QuantizedTensor -> ShowS
$cshow :: QuantizedTensor -> String
show :: QuantizedTensor -> String
$cshowList :: [QuantizedTensor] -> ShowS
showList :: [QuantizedTensor] -> ShowS
Show)

-- | Return signed INT8 storage in logical row-major order.
quantizedValues :: QuantizedTensor -> [Int8]
quantizedValues :: QuantizedTensor -> [Int8]
quantizedValues (QuantizedTensor [Int]
_ Maybe Int
_ [Double]
_ [Int]
_ [Int8]
xs) = [Int8]
xs

-- | Return one positive scale per calibrated channel, or a single per-tensor scale.
quantizationScales :: QuantizedTensor -> [Double]
quantizationScales :: QuantizedTensor -> [Double]
quantizationScales (QuantizedTensor [Int]
_ Maybe Int
_ [Double]
xs [Int]
_ [Int8]
_) = [Double]
xs

-- | Return integer zero offsets under the selected symmetric/asymmetric calibration.
quantizationZeroPoints :: QuantizedTensor -> [Int]
quantizationZeroPoints :: QuantizedTensor -> [Int]
quantizationZeroPoints (QuantizedTensor [Int]
_ Maybe Int
_ [Double]
_ [Int]
xs [Int8]
_) = [Int]
xs

finite :: [Double] -> Either MathError ()
finite :: [Double] -> Either MathError ()
finite [Double]
xs = if (Double -> Bool) -> [Double] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\Double
x -> Bool -> Bool
not (Double -> Bool
forall a. RealFloat a => a -> Bool
isNaN Double
x Bool -> Bool -> Bool
|| Double -> Bool
forall a. RealFloat a => a -> Bool
isInfinite Double
x)) [Double]
xs then () -> Either MathError ()
forall a b. b -> Either a b
Right () else MathError -> Either MathError ()
forall a b. a -> Either a b
Left MathError
NonFiniteResult
indices :: [Int] -> [[Int]]
indices :: [Int] -> [[Int]]
indices [] = [[]]
indices (Int
d : [Int]
ds) = [Int
i Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
is | Int
i <- [Int
0 .. Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1], [Int]
is <- [Int] -> [[Int]]
indices [Int]
ds]

-- | Calibrate finite nonempty values per tensor or selected axis. Symmetric uses [-127,127], asymmetric [-128,127], with ties-to-even rounding.
quantize :: Quantization -> Maybe Int -> T.Tensor Double -> Either MathError QuantizedTensor
quantize :: Quantization
-> Maybe Int -> Tensor Double -> Either MathError QuantizedTensor
quantize Quantization
scheme Maybe Int
axis Tensor Double
input = do
    let shape :: [Int]
shape = Tensor Double -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor Double
input; values :: [Double]
values = Tensor Double -> [Double]
forall a. Tensor a -> [a]
T.tensorValues Tensor Double
input
    [Double] -> Either MathError ()
finite [Double]
values
    if [Double] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Double]
values Bool -> Bool -> Bool
|| Bool -> (Int -> Bool) -> Maybe Int -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (\Int
a -> Int
a Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
a Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
shape) Maybe Int
axis
        then MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"quantization empty tensor/axis")
        else () -> Either MathError ()
forall a b. b -> Either a b
Right ()
    let channels :: Int
channels = Int -> (Int -> Int) -> Maybe Int -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
1 ([Int]
shape [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!!) Maybe Int
axis
        group :: [b] -> b
group [b]
ix = b -> (Int -> b) -> Maybe Int -> b
forall b a. b -> (a -> b) -> Maybe a -> b
maybe b
0 ([b]
ix [b] -> Int -> b
forall a. HasCallStack => [a] -> Int -> a
!!) Maybe Int
axis
        samples :: Int -> [Double]
samples Int
c = [Double
v | ([Int]
ix, Double
v) <- [[Int]] -> [Double] -> [([Int], Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip ([Int] -> [[Int]]
indices [Int]
shape) [Double]
values, [Int] -> Int
forall {b}. Num b => [b] -> b
group [Int]
ix Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
c]
        calibrate :: t a -> (a, b)
calibrate t a
xs =
            let lo :: a
lo = a -> a -> a
forall a. Ord a => a -> a -> a
min a
0 (t a -> a
forall a. Ord a => t a -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum t a
xs); hi :: a
hi = a -> a -> a
forall a. Ord a => a -> a -> a
max a
0 (t a -> a
forall a. Ord a => t a -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum t a
xs)
             in case Quantization
scheme of
                    Quantization
Symmetric -> (if a -> a -> a
forall a. Ord a => a -> a -> a
max (a -> a
forall a. Num a => a -> a
abs a
lo) (a -> a
forall a. Num a => a -> a
abs a
hi) a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0 then a
1 else a -> a -> a
forall a. Ord a => a -> a -> a
max (a -> a
forall a. Num a => a -> a
abs a
lo) (a -> a
forall a. Num a => a -> a
abs a
hi) a -> a -> a
forall a. Fractional a => a -> a -> a
/ a
127, b
0)
                    Quantization
Asymmetric ->
                        let s :: a
s = if a
hi a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
lo then a
1 else (a
hi a -> a -> a
forall a. Num a => a -> a -> a
- a
lo) a -> a -> a
forall a. Fractional a => a -> a -> a
/ a
255; z :: b
z = b -> b -> b
forall a. Ord a => a -> a -> a
max (-b
128) (b -> b -> b
forall a. Ord a => a -> a -> a
min b
127 (a -> b
forall b. Integral b => a -> b
forall a b. (RealFrac a, Integral b) => a -> b
round (-a
128 a -> a -> a
forall a. Num a => a -> a -> a
- a
lo a -> a -> a
forall a. Fractional a => a -> a -> a
/ a
s)))
                         in (a
s, b
z)
        params :: [(Double, Int)]
params = (Int -> (Double, Int)) -> [Int] -> [(Double, Int)]
forall a b. (a -> b) -> [a] -> [b]
map ([Double] -> (Double, Int)
forall {t :: * -> *} {a} {b}.
(Foldable t, RealFrac a, Integral b) =>
t a -> (a, b)
calibrate ([Double] -> (Double, Int))
-> (Int -> [Double]) -> Int -> (Double, Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [Double]
samples) [Int
0 .. Int
channels Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
        quantized :: [Int8]
quantized =
            [ let (Double
scale, Int
zero) = [(Double, Int)]
params [(Double, Int)] -> Int -> (Double, Int)
forall a. HasCallStack => [a] -> Int -> a
!! [Int] -> Int
forall {b}. Num b => [b] -> b
group [Int]
ix; low :: Int
low = if Quantization
scheme Quantization -> Quantization -> Bool
forall a. Eq a => a -> a -> Bool
== Quantization
Symmetric then -Int
127 else -Int
128
               in Int -> Int8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
low (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
127 (Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
round (Double
v Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
scale) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
zero)))
            | ([Int]
ix, Double
v) <- [[Int]] -> [Double] -> [([Int], Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip ([Int] -> [[Int]]
indices [Int]
shape) [Double]
values
            ]
    [Double] -> Either MathError ()
finite (((Double, Int) -> Double) -> [(Double, Int)] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double, Int) -> Double
forall a b. (a, b) -> a
fst [(Double, Int)]
params)
    if ((Double, Int) -> Bool) -> [(Double, Int)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any ((Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0) (Double -> Bool)
-> ((Double, Int) -> Double) -> (Double, Int) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Double, Int) -> Double
forall a b. (a, b) -> a
fst) [(Double, Int)]
params
        then MathError -> Either MathError QuantizedTensor
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"quantization scale underflow")
        else QuantizedTensor -> Either MathError QuantizedTensor
forall a b. b -> Either a b
Right ([Int]
-> Maybe Int -> [Double] -> [Int] -> [Int8] -> QuantizedTensor
QuantizedTensor [Int]
shape Maybe Int
axis (((Double, Int) -> Double) -> [(Double, Int)] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double, Int) -> Double
forall a b. (a, b) -> a
fst [(Double, Int)]
params) (((Double, Int) -> Int) -> [(Double, Int)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Double, Int) -> Int
forall a b. (a, b) -> b
snd [(Double, Int)]
params) [Int8]
quantized)

-- | Decode owned INT8 values using scale*(q-zeroPoint), rejecting floating arithmetic overflow.
dequantize :: QuantizedTensor -> Either MathError (T.Tensor Double)
dequantize :: QuantizedTensor -> Either MathError (Tensor Double)
dequantize (QuantizedTensor [Int]
shape Maybe Int
axis [Double]
scales [Int]
zeros [Int8]
values) = do
    let decoded :: [Double]
decoded =
            [ let channel :: Int
channel = Int -> (Int -> Int) -> Maybe Int -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
0 ([Int]
ix [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!!) Maybe Int
axis
               in [Double]
scales [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
channel Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int8
v Int -> Int -> Int
forall a. Num a => a -> a -> a
- ([Int]
zeros [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
channel) :: Int)
            | ([Int]
ix, Int8
v) <- [[Int]] -> [Int8] -> [([Int], Int8)]
forall a b. [a] -> [b] -> [(a, b)]
zip ([Int] -> [[Int]]
indices [Int]
shape) [Int8]
values
            ]
    [Double] -> Either MathError ()
finite [Double]
decoded
    [Int] -> [Double] -> Either MathError (Tensor Double)
forall a. [Int] -> [a] -> Either MathError (Tensor a)
T.tensor [Int]
shape [Double]
decoded

-- | Multiply two INT8 matrices with exact integer accumulation; supports per-row A/per-column B or per-tensor scales.
int8Gemm :: QuantizedTensor -> QuantizedTensor -> Either MathError (T.Tensor Double)
int8Gemm :: QuantizedTensor
-> QuantizedTensor -> Either MathError (Tensor Double)
int8Gemm (QuantizedTensor [Int
m, Int
k] Maybe Int
aa [Double]
as [Int]
az [Int8]
av) (QuantizedTensor [Int
inner, Int
n] Maybe Int
ba [Double]
bs [Int]
bz [Int8]
bv)
    | Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
inner Bool -> Bool -> Bool
|| Bool -> (Int -> Bool) -> Maybe Int -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0) Maybe Int
aa Bool -> Bool -> Bool
|| Bool -> (Int -> Bool) -> Maybe Int -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
1) Maybe Int
ba =
        MathError -> Either MathError (Tensor Double)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"INT8 GEMM requires A per-row/B per-column or per-tensor scales")
    | Bool
otherwise = do
        vals <-
            ((Int, Int) -> Either MathError Double)
-> [(Int, Int)] -> Either MathError [Double]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM
                ( \(Int
i, Int
j) -> do
                    let ai :: Int
ai = Int -> (Int -> Int) -> Maybe Int -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
0 (Int -> Int -> Int
forall a b. a -> b -> a
const Int
i) Maybe Int
aa
                        bj :: Int
bj = Int -> (Int -> Int) -> Maybe Int -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
0 (Int -> Int -> Int
forall a b. a -> b -> a
const Int
j) Maybe Int
ba
                        accum :: Integer
accum =
                            [Integer] -> Integer
forall {b}. Num b => [b] -> b
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum
                                [ (Int8 -> Integer
forall a. Integral a => a -> Integer
toInteger ([Int8]
av [Int8] -> Int -> Int8
forall a. HasCallStack => [a] -> Int -> a
!! (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
t)) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Int -> Integer
forall a. Integral a => a -> Integer
toInteger ([Int]
az [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
ai))
                                    Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* (Int8 -> Integer
forall a. Integral a => a -> Integer
toInteger ([Int8]
bv [Int8] -> Int -> Int8
forall a. HasCallStack => [a] -> Int -> a
!! (Int
t Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j)) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Int -> Integer
forall a. Integral a => a -> Integer
toInteger ([Int]
bz [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
bj))
                                | Int
t <- [Int
0 .. Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                                ]
                    if Integer
accum Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Int32 -> Integer
forall a. Integral a => a -> Integer
toInteger (Int32
forall a. Bounded a => a
minBound :: Int32) Bool -> Bool -> Bool
|| Integer
accum Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Int32 -> Integer
forall a. Integral a => a -> Integer
toInteger (Int32
forall a. Bounded a => a
maxBound :: Int32)
                        then MathError -> Either MathError Double
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"INT8 GEMM accumulator overflow")
                        else Double -> Either MathError Double
forall a b. b -> Either a b
Right (Integer -> Double
forall a. Num a => Integer -> a
fromInteger Integer
accum Double -> Double -> Double
forall a. Num a => a -> a -> a
* ([Double]
as [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
ai) Double -> Double -> Double
forall a. Num a => a -> a -> a
* ([Double]
bs [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
bj))
                )
                [(Int
i, Int
j) | Int
i <- [Int
0 .. Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1], Int
j <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
        finite vals
        T.tensor [m, n] vals
int8Gemm QuantizedTensor
_ QuantizedTensor
_ = MathError -> Either MathError (Tensor Double)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"INT8 GEMM expects matrices")

{- | Mixed-precision storage/rounding selector FP16 or BF16; operations report
finite/overflow conditions through the checked precision API.
-}
data Precision = FP16 | BF16 deriving (Precision -> Precision -> Bool
(Precision -> Precision -> Bool)
-> (Precision -> Precision -> Bool) -> Eq Precision
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Precision -> Precision -> Bool
== :: Precision -> Precision -> Bool
$c/= :: Precision -> Precision -> Bool
/= :: Precision -> Precision -> Bool
Eq, Int -> Precision -> ShowS
[Precision] -> ShowS
Precision -> String
(Int -> Precision -> ShowS)
-> (Precision -> String)
-> ([Precision] -> ShowS)
-> Show Precision
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Precision -> ShowS
showsPrec :: Int -> Precision -> ShowS
$cshow :: Precision -> String
show :: Precision -> String
$cshowList :: [Precision] -> ShowS
showList :: [Precision] -> ShowS
Show)

-- | Round operands to the chosen 16-bit format, accumulate in Float32, and return checked Double values.
reducedGemm :: Precision -> T.Tensor Double -> T.Tensor Double -> Either MathError (T.Tensor Double)
reducedGemm :: Precision
-> Tensor Double
-> Tensor Double
-> Either MathError (Tensor Double)
reducedGemm Precision
precision Tensor Double
a Tensor Double
b = do
    [Double] -> Either MathError ()
finite (Tensor Double -> [Double]
forall a. Tensor a -> [a]
T.tensorValues Tensor Double
a [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ Tensor Double -> [Double]
forall a. Tensor a -> [a]
T.tensorValues Tensor Double
b)
    case (Tensor Double -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor Double
a, Tensor Double -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor Double
b) of
        ([Int
m, Int
k], [Int
inner, Int
n]) | Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
inner -> do
            let convert :: Double -> Double
convert Double
x = case Precision
precision of
                    Precision
FP16 -> let N.Rounded Float16
y [NumericException]
_ = RoundingMode -> Double -> Rounded Float16
N.float16 RoundingMode
N.NearestEven Double
x in Float16 -> Double
N.fromFloat16 Float16
y
                    Precision
BF16 -> let N.Rounded BFloat16
y [NumericException]
_ = RoundingMode -> Double -> Rounded BFloat16
N.bfloat16 RoundingMode
N.NearestEven Double
x in BFloat16 -> Double
N.fromBFloat16 BFloat16
y
                av :: [Double]
av = (Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Double -> Double
convert (Tensor Double -> [Double]
forall a. Tensor a -> [a]
T.tensorValues Tensor Double
a)
                bv :: [Double]
bv = (Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Double -> Double
convert (Tensor Double -> [Double]
forall a. Tensor a -> [a]
T.tensorValues Tensor Double
b)
                fp32 :: a -> Float
fp32 a
x = a -> Float
forall a b. (Real a, Fractional b) => a -> b
realToFrac a
x :: Float
                vals :: [Double]
vals =
                    [ Float -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac
                        ((Float -> Float -> Float) -> Float -> [Float] -> Float
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl' Float -> Float -> Float
forall a. Num a => a -> a -> a
(+) Float
0 [Double -> Float
forall {a}. Real a => a -> Float
fp32 ([Double]
av [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
t)) Float -> Float -> Float
forall a. Num a => a -> a -> a
* Double -> Float
forall {a}. Real a => a -> Float
fp32 ([Double]
bv [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! (Int
t Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j)) | Int
t <- [Int
0 .. Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]])
                    | Int
i <- [Int
0 .. Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                    , Int
j <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
                    ]
            [Double] -> Either MathError ()
finite ([Double]
av [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ [Double]
bv [Double] -> [Double] -> [Double]
forall a. [a] -> [a] -> [a]
++ [Double]
vals)
            [Int] -> [Double] -> Either MathError (Tensor Double)
forall a. [Int] -> [a] -> Either MathError (Tensor a)
T.tensor [Int
m, Int
n] [Double]
vals
        ([Int], [Int])
_ -> MathError -> Either MathError (Tensor Double)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"reduced GEMM expects conformable matrices")

-- | Dynamic mixed-precision loss-scale state with finite-step growth tracking.
data LossScaler = LossScaler
    { LossScaler -> Double
lossScale :: Double
    -- ^ Positive scaling factor applied to loss/gradients in mixed-precision training.
    , LossScaler -> Int
scaleGrowthInterval :: Int
    -- ^ Number of consecutive finite updates before loss-scale growth.
    , LossScaler -> Int
finiteSteps :: Int
    -- ^ Current count of consecutive finite updates since the last reset/growth.
    }
    deriving (LossScaler -> LossScaler -> Bool
(LossScaler -> LossScaler -> Bool)
-> (LossScaler -> LossScaler -> Bool) -> Eq LossScaler
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: LossScaler -> LossScaler -> Bool
== :: LossScaler -> LossScaler -> Bool
$c/= :: LossScaler -> LossScaler -> Bool
/= :: LossScaler -> LossScaler -> Bool
Eq, Int -> LossScaler -> ShowS
[LossScaler] -> ShowS
LossScaler -> String
(Int -> LossScaler -> ShowS)
-> (LossScaler -> String)
-> ([LossScaler] -> ShowS)
-> Show LossScaler
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> LossScaler -> ShowS
showsPrec :: Int -> LossScaler -> ShowS
$cshow :: LossScaler -> String
show :: LossScaler -> String
$cshowList :: [LossScaler] -> ShowS
showList :: [LossScaler] -> ShowS
Show)

-- | Create a positive finite loss scale and positive growth interval; zero accumulated finite steps.
initialLossScaler :: Double -> Int -> Either MathError LossScaler
initialLossScaler :: Double -> Int -> Either MathError LossScaler
initialLossScaler Double
scale Int
interval = do
    [Double] -> Either MathError ()
finite [Double
scale]
    if Double
scale Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 Bool -> Bool -> Bool
|| Int
interval Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1
        then MathError -> Either MathError LossScaler
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"loss scaler scale/interval")
        else LossScaler -> Either MathError LossScaler
forall a b. b -> Either a b
Right (Double -> Int -> Int -> LossScaler
LossScaler Double
scale Int
interval Int
0)

-- | Multiply a finite loss by the checked current scale; overflow is a structured failure.
scaleLoss :: LossScaler -> Double -> Either MathError Double
scaleLoss :: LossScaler -> Double -> Either MathError Double
scaleLoss LossScaler
state Double
loss = do
    _ <- Double -> Int -> Either MathError LossScaler
initialLossScaler (LossScaler -> Double
lossScale LossScaler
state) (LossScaler -> Int
scaleGrowthInterval LossScaler
state)
    finite [loss, loss * lossScale state]
    Right (loss * lossScale state)

{- | False means overflow: every caller must skip the entire optimizer update.
Gradients are returned only when every scaled gradient is finite.
-}
unscaleGradients :: LossScaler -> [Double] -> Either MathError (Bool, [Double])
unscaleGradients :: LossScaler -> [Double] -> Either MathError (Bool, [Double])
unscaleGradients LossScaler
state [Double]
gradients = do
    _ <- Double -> Int -> Either MathError LossScaler
initialLossScaler (LossScaler -> Double
lossScale LossScaler
state) (LossScaler -> Int
scaleGrowthInterval LossScaler
state)
    if any (\Double
x -> Double -> Bool
forall a. RealFloat a => a -> Bool
isNaN Double
x Bool -> Bool -> Bool
|| Double -> Bool
forall a. RealFloat a => a -> Bool
isInfinite Double
x) gradients
        then Right (False, [])
        else do
            let gs = (Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ LossScaler -> Double
lossScale LossScaler
state) [Double]
gradients
            finite gs
            Right (True, gs)

-- | Halve scale after overflow; double after growthInterval finite updates, with a bounded maximum.
updateLossScaler :: Bool -> LossScaler -> Either MathError LossScaler
updateLossScaler :: Bool -> LossScaler -> Either MathError LossScaler
updateLossScaler Bool
successful LossScaler
state = do
    _ <- Double -> Int -> Either MathError LossScaler
initialLossScaler (LossScaler -> Double
lossScale LossScaler
state) (LossScaler -> Int
scaleGrowthInterval LossScaler
state)
    if finiteSteps state < 0 then Left (InvalidArgument "negative finite scaler steps") else Right ()
    let steps = if Bool
successful then LossScaler -> Int
finiteSteps LossScaler
state Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 else Int
0
        grow = Int
steps Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= LossScaler -> Int
scaleGrowthInterval LossScaler
state
        scale =
            if Bool -> Bool
not Bool
successful
                then Double -> Double -> Double
forall a. Ord a => a -> a -> a
max Double
1 (LossScaler -> Double
lossScale LossScaler
state Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2)
                else if Bool
grow then Double -> Double -> Double
forall a. Ord a => a -> a -> a
min (Double
2 Double -> Double -> Double
forall a. Floating a => a -> a -> a
** Double
120) (LossScaler -> Double
lossScale LossScaler
state Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
2) else LossScaler -> Double
lossScale LossScaler
state
    Right (LossScaler scale (scaleGrowthInterval state) (if grow then 0 else steps))

-- | Compute a checked scaled Euclidean norm; unrepresentable final norms fail explicitly.
gradientNorm :: [Double] -> Either MathError Double
gradientNorm :: [Double] -> Either MathError Double
gradientNorm [Double]
xs = do
    [Double] -> Either MathError ()
finite [Double]
xs
    let m :: Double
m = [Double] -> Double
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum (Double
0 Double -> [Double] -> [Double]
forall a. a -> [a] -> [a]
: (Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Double -> Double
forall a. Num a => a -> a
abs [Double]
xs)
        y :: Double
y = if Double
m Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 then Double
0 else Double
m Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
sqrt ([Double] -> Double
forall {b}. Num b => [b] -> b
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [(Double
x Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
m) Double -> Int -> Double
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int) | Double
x <- [Double]
xs])
    [Double] -> Either MathError ()
finite [Double
y]
    Double -> Either MathError Double
forall a b. b -> Either a b
Right Double
y

-- | Clip norm through the shared overflow-safe automatic-differentiation gradient utility.
clipGradients :: Double -> [Double] -> Either MathError [Double]
clipGradients :: Double -> [Double] -> Either MathError [Double]
clipGradients = Double -> [Double] -> Either MathError [Double]
AD.clipGradient