-module(viva_tensor@nn@backward).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_tensor/nn/backward.gleam").
-export([relu_backward/2, sigmoid_backward/2, tanh_backward/2, gelu_backward/2, leaky_relu_backward/3, elu_backward/3, mse_loss_backward/4, l1_loss_backward/4, bce_loss_backward/4, cross_entropy_loss_backward/4, matmul_backward/3, linear_backward/3, layer_norm_backward/6, rms_norm_backward/5, softmax_backward/3]).

-if(?OTP_RELEASE >= 27).
-define(MODULEDOC(Str), -moduledoc(Str)).
-define(DOC(Str), -doc(Str)).
-else.
-define(MODULEDOC(Str), -compile([])).
-define(DOC(Str), -compile([])).
-endif.

?MODULEDOC(false).

-file("src/viva_tensor/nn/backward.gleam", 754).
?DOC(false).
-spec elementwise2(
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor(),
    fun((float(), float()) -> float())
) -> {ok, viva_tensor@tensor:tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
elementwise2(A, B, F) ->
    gleam@result:'try'(
        viva_tensor@tensor:try_to_list(A),
        fun(A_data) ->
            gleam@result:'try'(
                viva_tensor@tensor:try_to_list(B),
                fun(B_data) ->
                    Out = begin
                        _pipe = gleam@list:zip(A_data, B_data),
                        gleam@list:map(
                            _pipe,
                            fun(P) ->
                                {X, Y} = P,
                                F(X, Y)
                            end
                        )
                    end,
                    {ok, {tensor, Out, viva_tensor@tensor:shape(A)}}
                end
            )
        end
    ).

-file("src/viva_tensor/nn/backward.gleam", 770).
?DOC(false).
-spec ensure_same_shape(
    binary(),
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor()
) -> {ok, nil} | {error, viva_tensor@core@error:tensor_error()}.
ensure_same_shape(_, A, B) ->
    case viva_tensor@tensor:shape(A) =:= viva_tensor@tensor:shape(B) of
        true ->
            {ok, nil};

        false ->
            {error,
                {shape_mismatch,
                    viva_tensor@tensor:shape(A),
                    viva_tensor@tensor:shape(B)}}
    end.

-file("src/viva_tensor/nn/backward.gleam", 54).
?DOC(false).
-spec relu_backward(viva_tensor@tensor:tensor(), viva_tensor@tensor:tensor()) -> {ok,
        viva_tensor@tensor:tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
relu_backward(Grad_out, Input) ->
    _pipe = ensure_same_shape(<<"relu_backward"/utf8>>, Grad_out, Input),
    gleam@result:'try'(
        _pipe,
        fun(_) -> elementwise2(Grad_out, Input, fun(G, X) -> case X > +0.0 of
                        true ->
                            G;

                        false ->
                            +0.0
                    end end) end
    ).

-file("src/viva_tensor/nn/backward.gleam", 74).
?DOC(false).
-spec sigmoid_backward(viva_tensor@tensor:tensor(), viva_tensor@tensor:tensor()) -> {ok,
        viva_tensor@tensor:tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
sigmoid_backward(Grad_out, Output) ->
    _pipe = ensure_same_shape(<<"sigmoid_backward"/utf8>>, Grad_out, Output),
    gleam@result:'try'(
        _pipe,
        fun(_) ->
            elementwise2(Grad_out, Output, fun(G, Y) -> (G * Y) * (1.0 - Y) end)
        end
    ).

-file("src/viva_tensor/nn/backward.gleam", 88).
?DOC(false).
-spec tanh_backward(viva_tensor@tensor:tensor(), viva_tensor@tensor:tensor()) -> {ok,
        viva_tensor@tensor:tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
tanh_backward(Grad_out, Output) ->
    _pipe = ensure_same_shape(<<"tanh_backward"/utf8>>, Grad_out, Output),
    gleam@result:'try'(
        _pipe,
        fun(_) ->
            elementwise2(Grad_out, Output, fun(G, Y) -> G * (1.0 - (Y * Y)) end)
        end
    ).

-file("src/viva_tensor/nn/backward.gleam", 106).
?DOC(false).
-spec gelu_backward(viva_tensor@tensor:tensor(), viva_tensor@tensor:tensor()) -> {ok,
        viva_tensor@tensor:tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
gelu_backward(Grad_out, Input) ->
    _pipe = ensure_same_shape(<<"gelu_backward"/utf8>>, Grad_out, Input),
    gleam@result:'try'(
        _pipe,
        fun(_) ->
            elementwise2(
                Grad_out,
                Input,
                fun(G, X) ->
                    Phi_part = 1.0 + math:erf(X * 0.7071067811865475),
                    Pdf_part = (X * 0.7978845608028654) * math:exp(
                        (-0.5 * X) * X
                    ),
                    (G * 0.5) * (Phi_part + Pdf_part)
                end
            )
        end
    ).

-file("src/viva_tensor/nn/backward.gleam", 125).
?DOC(false).
-spec leaky_relu_backward(
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor(),
    float()
) -> {ok, viva_tensor@tensor:tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
leaky_relu_backward(Grad_out, Input, Negative_slope) ->
    _pipe = ensure_same_shape(<<"leaky_relu_backward"/utf8>>, Grad_out, Input),
    gleam@result:'try'(
        _pipe,
        fun(_) -> elementwise2(Grad_out, Input, fun(G, X) -> case X > +0.0 of
                        true ->
                            G;

                        false ->
                            G * Negative_slope
                    end end) end
    ).

-file("src/viva_tensor/nn/backward.gleam", 146).
?DOC(false).
-spec elu_backward(
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor(),
    float()
) -> {ok, viva_tensor@tensor:tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
elu_backward(Grad_out, Input, Alpha) ->
    _pipe = ensure_same_shape(<<"elu_backward"/utf8>>, Grad_out, Input),
    gleam@result:'try'(
        _pipe,
        fun(_) -> elementwise2(Grad_out, Input, fun(G, X) -> case X > +0.0 of
                        true ->
                            G;

                        false ->
                            (G * Alpha) * math:exp(X)
                    end end) end
    ).

-file("src/viva_tensor/nn/backward.gleam", 813).
?DOC(false).
-spec reduction_inv_n(
    viva_tensor@tensor:tensor(),
    viva_tensor@nn@losses:reduction()
) -> float().
reduction_inv_n(Prediction, Reduction) ->
    case Reduction of
        reduction_mean ->
            N = viva_tensor@tensor:size(Prediction),
            case N =< 0 of
                true ->
                    1.0;

                false ->
                    case erlang:float(N) of
                        +0.0 -> +0.0;
                        -0.0 -> -0.0;
                        Gleam@denominator -> 1.0 / Gleam@denominator
                    end
            end;

        _ ->
            1.0
    end.

-file("src/viva_tensor/nn/backward.gleam", 802).
?DOC(false).
-spec grad_scalar_value(viva_tensor@tensor:tensor()) -> {ok, float()} |
    {error, viva_tensor@core@error:tensor_error()}.
grad_scalar_value(Grad_out) ->
    case viva_tensor@tensor:try_to_list(Grad_out) of
        {ok, [V]} ->
            {ok, V};

        {ok, _} ->
            {error,
                {invalid_shape,
                    <<"expected scalar (shape [1]) grad_out for reduced loss backward"/utf8>>}};

        {error, E} ->
            {error, E}
    end.

-file("src/viva_tensor/nn/backward.gleam", 782).
?DOC(false).
-spec loss_scale(
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor(),
    viva_tensor@nn@losses:reduction()
) -> {ok, float()} | {error, viva_tensor@core@error:tensor_error()}.
loss_scale(Grad_out, Prediction, Reduction) ->
    case Reduction of
        reduction_none ->
            case viva_tensor@tensor:shape(Grad_out) =:= viva_tensor@tensor:shape(
                Prediction
            ) of
                true ->
                    {ok, 1.0};

                false ->
                    {error,
                        {shape_mismatch,
                            viva_tensor@tensor:shape(Prediction),
                            viva_tensor@tensor:shape(Grad_out)}}
            end;

        _ ->
            grad_scalar_value(Grad_out)
    end.

-file("src/viva_tensor/nn/backward.gleam", 176).
?DOC(false).
-spec mse_loss_backward(
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor(),
    viva_tensor@nn@losses:reduction()
) -> {ok, viva_tensor@tensor:tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
mse_loss_backward(Grad_out, Prediction, Target, Reduction) ->
    gleam@result:'try'(
        ensure_same_shape(<<"mse_loss_backward"/utf8>>, Prediction, Target),
        fun(_) ->
            gleam@result:'try'(
                viva_tensor@tensor:try_to_list(Prediction),
                fun(Pred_data) ->
                    gleam@result:'try'(
                        viva_tensor@tensor:try_to_list(Target),
                        fun(Target_data) ->
                            gleam@result:'try'(
                                loss_scale(Grad_out, Prediction, Reduction),
                                fun(Scale) ->
                                    N_inv = reduction_inv_n(
                                        Prediction,
                                        Reduction
                                    ),
                                    Pred_shape = viva_tensor@tensor:shape(
                                        Prediction
                                    ),
                                    case Reduction of
                                        reduction_none ->
                                            gleam@result:'try'(
                                                viva_tensor@tensor:try_to_list(
                                                    Grad_out
                                                ),
                                                fun(Grad_data) ->
                                                    Out = gleam@list:map(
                                                        gleam@list:zip(
                                                            Grad_data,
                                                            gleam@list:zip(
                                                                Pred_data,
                                                                Target_data
                                                            )
                                                        ),
                                                        fun(T) ->
                                                            {G, Rest} = T,
                                                            {P, Y} = Rest,
                                                            (G * 2.0) * (P - Y)
                                                        end
                                                    ),
                                                    {ok,
                                                        {tensor,
                                                            Out,
                                                            Pred_shape}}
                                                end
                                            );

                                        _ ->
                                            Out@1 = gleam@list:map(
                                                gleam@list:zip(
                                                    Pred_data,
                                                    Target_data
                                                ),
                                                fun(Pair) ->
                                                    {P@1, Y@1} = Pair,
                                                    ((Scale * 2.0) * (P@1 - Y@1))
                                                    * N_inv
                                                end
                                            ),
                                            {ok, {tensor, Out@1, Pred_shape}}
                                    end
                                end
                            )
                        end
                    )
                end
            )
        end
    ).

-file("src/viva_tensor/nn/backward.gleam", 826).
?DOC(false).
-spec sign(float()) -> float().
sign(X) ->
    case {X > +0.0, X < +0.0} of
        {true, _} ->
            1.0;

        {_, true} ->
            -1.0;

        {_, _} ->
            +0.0
    end.

-file("src/viva_tensor/nn/backward.gleam", 217).
?DOC(false).
-spec l1_loss_backward(
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor(),
    viva_tensor@nn@losses:reduction()
) -> {ok, viva_tensor@tensor:tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
l1_loss_backward(Grad_out, Prediction, Target, Reduction) ->
    gleam@result:'try'(
        ensure_same_shape(<<"l1_loss_backward"/utf8>>, Prediction, Target),
        fun(_) ->
            gleam@result:'try'(
                viva_tensor@tensor:try_to_list(Prediction),
                fun(Pred_data) ->
                    gleam@result:'try'(
                        viva_tensor@tensor:try_to_list(Target),
                        fun(Target_data) ->
                            Pred_shape = viva_tensor@tensor:shape(Prediction),
                            case Reduction of
                                reduction_none ->
                                    gleam@result:'try'(
                                        viva_tensor@tensor:try_to_list(Grad_out),
                                        fun(Grad_data) ->
                                            Out = gleam@list:map(
                                                gleam@list:zip(
                                                    Grad_data,
                                                    gleam@list:zip(
                                                        Pred_data,
                                                        Target_data
                                                    )
                                                ),
                                                fun(T) ->
                                                    {G, Rest} = T,
                                                    {P, Y} = Rest,
                                                    G * sign(P - Y)
                                                end
                                            ),
                                            {ok, {tensor, Out, Pred_shape}}
                                        end
                                    );

                                _ ->
                                    gleam@result:'try'(
                                        loss_scale(
                                            Grad_out,
                                            Prediction,
                                            Reduction
                                        ),
                                        fun(Scale) ->
                                            N_inv = reduction_inv_n(
                                                Prediction,
                                                Reduction
                                            ),
                                            Out@1 = gleam@list:map(
                                                gleam@list:zip(
                                                    Pred_data,
                                                    Target_data
                                                ),
                                                fun(Pair) ->
                                                    {P@1, Y@1} = Pair,
                                                    (Scale * sign(P@1 - Y@1)) * N_inv
                                                end
                                            ),
                                            {ok, {tensor, Out@1, Pred_shape}}
                                        end
                                    )
                            end
                        end
                    )
                end
            )
        end
    ).

-file("src/viva_tensor/nn/backward.gleam", 259).
?DOC(false).
-spec bce_loss_backward(
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor(),
    viva_tensor@nn@losses:reduction()
) -> {ok, viva_tensor@tensor:tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
bce_loss_backward(Grad_out, Prediction, Target, Reduction) ->
    gleam@result:'try'(
        ensure_same_shape(<<"bce_loss_backward"/utf8>>, Prediction, Target),
        fun(_) ->
            gleam@result:'try'(
                viva_tensor@tensor:try_to_list(Prediction),
                fun(Pred_data) ->
                    gleam@result:'try'(
                        viva_tensor@tensor:try_to_list(Target),
                        fun(Target_data) ->
                            Eps = 1.0e-7,
                            One_minus_eps = 1.0 - Eps,
                            Pred_shape = viva_tensor@tensor:shape(Prediction),
                            case Reduction of
                                reduction_none ->
                                    gleam@result:'try'(
                                        viva_tensor@tensor:try_to_list(Grad_out),
                                        fun(Grad_data) ->
                                            Out = gleam@list:map(
                                                gleam@list:zip(
                                                    Grad_data,
                                                    gleam@list:zip(
                                                        Pred_data,
                                                        Target_data
                                                    )
                                                ),
                                                fun(T) ->
                                                    {G, Rest} = T,
                                                    {P, Y} = Rest,
                                                    P_c = gleam@float:clamp(
                                                        P,
                                                        Eps,
                                                        One_minus_eps
                                                    ),
                                                    case (P_c * (1.0 - P_c)) of
                                                        +0.0 -> +0.0;
                                                        -0.0 -> -0.0;
                                                        Gleam@denominator -> G * (P_c
                                                        - Y)
                                                        / Gleam@denominator
                                                    end
                                                end
                                            ),
                                            {ok, {tensor, Out, Pred_shape}}
                                        end
                                    );

                                _ ->
                                    gleam@result:'try'(
                                        loss_scale(
                                            Grad_out,
                                            Prediction,
                                            Reduction
                                        ),
                                        fun(Scale) ->
                                            N_inv = reduction_inv_n(
                                                Prediction,
                                                Reduction
                                            ),
                                            Out@1 = gleam@list:map(
                                                gleam@list:zip(
                                                    Pred_data,
                                                    Target_data
                                                ),
                                                fun(Pair) ->
                                                    {P@1, Y@1} = Pair,
                                                    P_c@1 = gleam@float:clamp(
                                                        P@1,
                                                        Eps,
                                                        One_minus_eps
                                                    ),
                                                    (case (P_c@1 * (1.0 - P_c@1)) of
                                                        +0.0 -> +0.0;
                                                        -0.0 -> -0.0;
                                                        Gleam@denominator@1 -> Scale
                                                        * (P_c@1 - Y@1)
                                                        / Gleam@denominator@1
                                                    end)
                                                    * N_inv
                                                end
                                            ),
                                            {ok, {tensor, Out@1, Pred_shape}}
                                        end
                                    )
                            end
                        end
                    )
                end
            )
        end
    ).

-file("src/viva_tensor/nn/backward.gleam", 885).
?DOC(false).
-spec do_chunk_every(list(float()), integer(), list(list(float()))) -> list(list(float())).
do_chunk_every(Items, N, Acc) ->
    case Items of
        [] ->
            lists:reverse(Acc);

        _ ->
            Chunk = gleam@list:take(Items, N),
            Rest = gleam@list:drop(Items, N),
            do_chunk_every(Rest, N, [Chunk | Acc])
    end.

-file("src/viva_tensor/nn/backward.gleam", 878).
?DOC(false).
-spec chunk_every(list(float()), integer()) -> list(list(float())).
chunk_every(Items, N) ->
    case N =< 0 of
        true ->
            [];

        false ->
            do_chunk_every(Items, N, [])
    end.

-file("src/viva_tensor/nn/backward.gleam", 308).
?DOC(false).
-spec cross_entropy_loss_backward(
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor(),
    viva_tensor@nn@losses:reduction()
) -> {ok, viva_tensor@tensor:tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
cross_entropy_loss_backward(Grad_out, Logits, Targets, Reduction) ->
    Logits_shape = viva_tensor@tensor:shape(Logits),
    gleam@result:'try'(case Logits_shape of
            [B, C] ->
                {ok, {B, C}};

            _ ->
                {error,
                    {rank_mismatch,
                        <<"cross_entropy_loss_backward"/utf8>>,
                        2,
                        Logits_shape}}
        end, fun(_use0) ->
            {Batch, Num_classes} = _use0,
            Targets_shape = viva_tensor@tensor:shape(Targets),
            gleam@result:'try'(case Targets_shape of
                    [T_batch] when T_batch =:= Batch ->
                        {ok, nil};

                    _ ->
                        {error,
                            {operand_shape_mismatch,
                                <<"cross_entropy_loss_backward"/utf8>>,
                                <<"targets"/utf8>>,
                                <<"[batch]"/utf8>>,
                                Targets_shape}}
                end, fun(_) ->
                    gleam@result:'try'(
                        viva_tensor@tensor:softmax_axis(Logits, 1),
                        fun(Softmaxed) ->
                            gleam@result:'try'(
                                viva_tensor@tensor:try_to_list(Softmaxed),
                                fun(Sm_data) ->
                                    gleam@result:'try'(
                                        viva_tensor@tensor:try_to_list(Targets),
                                        fun(Target_data) ->
                                            Inv_batch = case Reduction of
                                                reduction_mean ->
                                                    case erlang:float(Batch) of
                                                        +0.0 -> +0.0;
                                                        -0.0 -> -0.0;
                                                        Gleam@denominator -> 1.0
                                                        / Gleam@denominator
                                                    end;

                                                _ ->
                                                    1.0
                                            end,
                                            case Reduction of
                                                reduction_none ->
                                                    gleam@result:'try'(
                                                        viva_tensor@tensor:try_to_list(
                                                            Grad_out
                                                        ),
                                                        fun(Grad_data) ->
                                                            Rows = chunk_every(
                                                                Sm_data,
                                                                Num_classes
                                                            ),
                                                            Zipped = gleam@list:zip(
                                                                Rows,
                                                                gleam@list:zip(
                                                                    Target_data,
                                                                    Grad_data
                                                                )
                                                            ),
                                                            Grad_rows = gleam@list:map(
                                                                Zipped,
                                                                fun(T) ->
                                                                    {Row, Rest} = T,
                                                                    {Target_f,
                                                                        G} = Rest,
                                                                    Class_idx = erlang:round(
                                                                        Target_f
                                                                    ),
                                                                    gleam@list:index_map(
                                                                        Row,
                                                                        fun(
                                                                            P,
                                                                            I
                                                                        ) ->
                                                                            One_hot = case I
                                                                            =:= Class_idx of
                                                                                true ->
                                                                                    1.0;

                                                                                false ->
                                                                                    +0.0
                                                                            end,
                                                                            G * (P
                                                                            - One_hot)
                                                                        end
                                                                    )
                                                                end
                                                            ),
                                                            {ok,
                                                                {tensor,
                                                                    lists:append(
                                                                        Grad_rows
                                                                    ),
                                                                    [Batch,
                                                                        Num_classes]}}
                                                        end
                                                    );

                                                _ ->
                                                    gleam@result:'try'(
                                                        grad_scalar_value(
                                                            Grad_out
                                                        ),
                                                        fun(Scale) ->
                                                            Rows@1 = chunk_every(
                                                                Sm_data,
                                                                Num_classes
                                                            ),
                                                            Zipped@1 = gleam@list:zip(
                                                                Rows@1,
                                                                Target_data
                                                            ),
                                                            Grad_rows@1 = gleam@list:map(
                                                                Zipped@1,
                                                                fun(Pair) ->
                                                                    {Row@1,
                                                                        Target_f@1} = Pair,
                                                                    Class_idx@1 = erlang:round(
                                                                        Target_f@1
                                                                    ),
                                                                    gleam@list:index_map(
                                                                        Row@1,
                                                                        fun(
                                                                            P@1,
                                                                            I@1
                                                                        ) ->
                                                                            One_hot@1 = case I@1
                                                                            =:= Class_idx@1 of
                                                                                true ->
                                                                                    1.0;

                                                                                false ->
                                                                                    +0.0
                                                                            end,
                                                                            (Scale
                                                                            * (P@1
                                                                            - One_hot@1))
                                                                            * Inv_batch
                                                                        end
                                                                    )
                                                                end
                                                            ),
                                                            {ok,
                                                                {tensor,
                                                                    lists:append(
                                                                        Grad_rows@1
                                                                    ),
                                                                    [Batch,
                                                                        Num_classes]}}
                                                        end
                                                    )
                                            end
                                        end
                                    )
                                end
                            )
                        end
                    )
                end)
        end).

-file("src/viva_tensor/nn/backward.gleam", 422).
?DOC(false).
-spec matmul_backward(
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor()
) -> {ok, {viva_tensor@tensor:tensor(), viva_tensor@tensor:tensor()}} |
    {error, viva_tensor@core@error:tensor_error()}.
matmul_backward(Grad_out, A, B) ->
    gleam@result:'try'(
        viva_tensor@tensor:transpose(B),
        fun(B_t) ->
            gleam@result:'try'(
                viva_tensor@tensor:transpose(A),
                fun(A_t) ->
                    gleam@result:'try'(
                        viva_tensor@tensor:matmul(Grad_out, B_t),
                        fun(Grad_a) ->
                            gleam@result:'try'(
                                viva_tensor@tensor:matmul(A_t, Grad_out),
                                fun(Grad_b) -> {ok, {Grad_a, Grad_b}} end
                            )
                        end
                    )
                end
            )
        end
    ).

-file("src/viva_tensor/nn/backward.gleam", 406).
?DOC(false).
-spec linear_backward(
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor()
) -> {ok, {viva_tensor@tensor:tensor(), viva_tensor@tensor:tensor()}} |
    {error, viva_tensor@core@error:tensor_error()}.
linear_backward(Grad_out, Input, Weight) ->
    matmul_backward(Grad_out, Input, Weight).

-file("src/viva_tensor/nn/backward.gleam", 834).
?DOC(false).
-spec safe_sqrt(float()) -> float().
safe_sqrt(X) ->
    case gleam@float:square_root(X) of
        {ok, V} ->
            V;

        {error, _} ->
            +0.0
    end.

-file("src/viva_tensor/nn/backward.gleam", 860).
?DOC(false).
-spec check_stat_size(binary(), viva_tensor@tensor:tensor(), integer()) -> {ok,
        nil} |
    {error, viva_tensor@core@error:tensor_error()}.
check_stat_size(_, T, Expected) ->
    N = viva_tensor@tensor:size(T),
    case N =:= Expected of
        true ->
            {ok, nil};

        false ->
            {error,
                {invalid_shape,
                    <<<<<<"expected "/utf8,
                                (erlang:integer_to_binary(Expected))/binary>>/binary,
                            " entries, got "/utf8>>/binary,
                        (erlang:integer_to_binary(N))/binary>>}}
    end.

-file("src/viva_tensor/nn/backward.gleam", 856).
?DOC(false).
-spec product(list(integer())) -> integer().
product(Shape) ->
    gleam@list:fold(Shape, 1, fun(A, B) -> A * B end).

-file("src/viva_tensor/nn/backward.gleam", 841).
?DOC(false).
-spec last_dim_of(list(integer())) -> {ok, integer()} |
    {error, viva_tensor@core@error:tensor_error()}.
last_dim_of(Shape) ->
    case gleam@list:last(Shape) of
        {ok, D} ->
            {ok, D};

        {error, _} ->
            {error,
                {invalid_shape, <<"expected non-empty shape, got []"/utf8>>}}
    end.

-file("src/viva_tensor/nn/backward.gleam", 459).
?DOC(false).
-spec layer_norm_backward(
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor(),
    float()
) -> {ok,
        {viva_tensor@tensor:tensor(),
            viva_tensor@tensor:tensor(),
            viva_tensor@tensor:tensor()}} |
    {error, viva_tensor@core@error:tensor_error()}.
layer_norm_backward(Grad_out, Input, Scale, Mean, Variance, Eps) ->
    Input_shape = viva_tensor@tensor:shape(Input),
    gleam@result:'try'(
        case viva_tensor@tensor:shape(Grad_out) =:= Input_shape of
            true ->
                {ok, nil};

            false ->
                {error,
                    {shape_mismatch,
                        Input_shape,
                        viva_tensor@tensor:shape(Grad_out)}}
        end,
        fun(_) ->
            gleam@result:'try'(
                last_dim_of(Input_shape),
                fun(D) ->
                    Scale_shape = viva_tensor@tensor:shape(Scale),
                    gleam@result:'try'(case Scale_shape =:= [D] of
                            true ->
                                {ok, nil};

                            false ->
                                {error, {shape_mismatch, [D], Scale_shape}}
                        end, fun(_) ->
                            Outer = case gleam@int:max(D, 1) of
                                0 -> 0;
                                Gleam@denominator -> product(Input_shape) div Gleam@denominator
                            end,
                            gleam@result:'try'(
                                check_stat_size(
                                    <<"layer_norm_backward.mean"/utf8>>,
                                    Mean,
                                    Outer
                                ),
                                fun(_) ->
                                    gleam@result:'try'(
                                        check_stat_size(
                                            <<"layer_norm_backward.variance"/utf8>>,
                                            Variance,
                                            Outer
                                        ),
                                        fun(_) ->
                                            gleam@result:'try'(
                                                viva_tensor@tensor:try_to_list(
                                                    Input
                                                ),
                                                fun(Data) ->
                                                    gleam@result:'try'(
                                                        viva_tensor@tensor:try_to_list(
                                                            Grad_out
                                                        ),
                                                        fun(Grad_data) ->
                                                            gleam@result:'try'(
                                                                viva_tensor@tensor:try_to_list(
                                                                    Scale
                                                                ),
                                                                fun(Scale_data) ->
                                                                    gleam@result:'try'(
                                                                        viva_tensor@tensor:try_to_list(
                                                                            Mean
                                                                        ),
                                                                        fun(
                                                                            Mean_data
                                                                        ) ->
                                                                            gleam@result:'try'(
                                                                                viva_tensor@tensor:try_to_list(
                                                                                    Variance
                                                                                ),
                                                                                fun(
                                                                                    Var_data
                                                                                ) ->
                                                                                    Input_rows = chunk_every(
                                                                                        Data,
                                                                                        D
                                                                                    ),
                                                                                    Grad_rows = chunk_every(
                                                                                        Grad_data,
                                                                                        D
                                                                                    ),
                                                                                    Combined = gleam@list:zip(
                                                                                        Input_rows,
                                                                                        gleam@list:zip(
                                                                                            Grad_rows,
                                                                                            gleam@list:zip(
                                                                                                Mean_data,
                                                                                                Var_data
                                                                                            )
                                                                                        )
                                                                                    ),
                                                                                    Init_acc = {[],
                                                                                        gleam@list:repeat(
                                                                                            +0.0,
                                                                                            D
                                                                                        ),
                                                                                        gleam@list:repeat(
                                                                                            +0.0,
                                                                                            D
                                                                                        )},
                                                                                    {Rev_grad_x,
                                                                                        Grad_scale_data,
                                                                                        Grad_bias_data} = gleam@list:fold(
                                                                                        Combined,
                                                                                        Init_acc,
                                                                                        fun(
                                                                                            Acc,
                                                                                            Row
                                                                                        ) ->
                                                                                            {Rev_gx,
                                                                                                Gs,
                                                                                                Gb} = Acc,
                                                                                            {X_row,
                                                                                                Rest1} = Row,
                                                                                            {G_row,
                                                                                                Stats} = Rest1,
                                                                                            {Mu,
                                                                                                Var} = Stats,
                                                                                            Std = safe_sqrt(
                                                                                                Var
                                                                                                + Eps
                                                                                            ),
                                                                                            Inv_std = case Std of
                                                                                                +0.0 -> +0.0;
                                                                                                -0.0 -> -0.0;
                                                                                                Gleam@denominator@1 -> 1.0
                                                                                                / Gleam@denominator@1
                                                                                            end,
                                                                                            X_hat = gleam@list:map(
                                                                                                X_row,
                                                                                                fun(
                                                                                                    X
                                                                                                ) ->
                                                                                                    (X
                                                                                                    - Mu)
                                                                                                    * Inv_std
                                                                                                end
                                                                                            ),
                                                                                            G_scaled = gleam@list:map(
                                                                                                gleam@list:zip(
                                                                                                    G_row,
                                                                                                    Scale_data
                                                                                                ),
                                                                                                fun(
                                                                                                    P
                                                                                                ) ->
                                                                                                    {G,
                                                                                                        S} = P,
                                                                                                    G
                                                                                                    * S
                                                                                                end
                                                                                            ),
                                                                                            D_f = erlang:float(
                                                                                                D
                                                                                            ),
                                                                                            M1 = case D_f of
                                                                                                +0.0 -> +0.0;
                                                                                                -0.0 -> -0.0;
                                                                                                Gleam@denominator@2 -> gleam@list:fold(
                                                                                                    G_scaled,
                                                                                                    +0.0,
                                                                                                    fun(
                                                                                                        A,
                                                                                                        V
                                                                                                    ) ->
                                                                                                        A
                                                                                                        + V
                                                                                                    end
                                                                                                )
                                                                                                / Gleam@denominator@2
                                                                                            end,
                                                                                            M2 = case D_f of
                                                                                                +0.0 -> +0.0;
                                                                                                -0.0 -> -0.0;
                                                                                                Gleam@denominator@3 -> gleam@list:fold(
                                                                                                    gleam@list:zip(
                                                                                                        G_scaled,
                                                                                                        X_hat
                                                                                                    ),
                                                                                                    +0.0,
                                                                                                    fun(
                                                                                                        A@1,
                                                                                                        P@1
                                                                                                    ) ->
                                                                                                        {Gs2,
                                                                                                            Xh} = P@1,
                                                                                                        A@1
                                                                                                        + (Gs2
                                                                                                        * Xh)
                                                                                                    end
                                                                                                )
                                                                                                / Gleam@denominator@3
                                                                                            end,
                                                                                            Grad_x_slice = gleam@list:map(
                                                                                                gleam@list:zip(
                                                                                                    G_scaled,
                                                                                                    X_hat
                                                                                                ),
                                                                                                fun(
                                                                                                    P@2
                                                                                                ) ->
                                                                                                    {Gs2@1,
                                                                                                        Xh@1} = P@2,
                                                                                                    ((Gs2@1
                                                                                                    - M1)
                                                                                                    - (Xh@1
                                                                                                    * M2))
                                                                                                    * Inv_std
                                                                                                end
                                                                                            ),
                                                                                            New_gs = gleam@list:map(
                                                                                                gleam@list:zip(
                                                                                                    Gs,
                                                                                                    gleam@list:zip(
                                                                                                        G_row,
                                                                                                        X_hat
                                                                                                    )
                                                                                                ),
                                                                                                fun(
                                                                                                    T
                                                                                                ) ->
                                                                                                    {Acc_v,
                                                                                                        Rest2} = T,
                                                                                                    {G@1,
                                                                                                        Xh@2} = Rest2,
                                                                                                    Acc_v
                                                                                                    + (G@1
                                                                                                    * Xh@2)
                                                                                                end
                                                                                            ),
                                                                                            New_gb = gleam@list:map(
                                                                                                gleam@list:zip(
                                                                                                    Gb,
                                                                                                    G_row
                                                                                                ),
                                                                                                fun(
                                                                                                    P@3
                                                                                                ) ->
                                                                                                    {Acc_v@1,
                                                                                                        G@2} = P@3,
                                                                                                    Acc_v@1
                                                                                                    + G@2
                                                                                                end
                                                                                            ),
                                                                                            {[Grad_x_slice |
                                                                                                    Rev_gx],
                                                                                                New_gs,
                                                                                                New_gb}
                                                                                        end
                                                                                    ),
                                                                                    Grad_x_data = begin
                                                                                        _pipe = Rev_grad_x,
                                                                                        _pipe@1 = lists:reverse(
                                                                                            _pipe
                                                                                        ),
                                                                                        lists:append(
                                                                                            _pipe@1
                                                                                        )
                                                                                    end,
                                                                                    {ok,
                                                                                        {{tensor,
                                                                                                Grad_x_data,
                                                                                                Input_shape},
                                                                                            {tensor,
                                                                                                Grad_scale_data,
                                                                                                [D]},
                                                                                            {tensor,
                                                                                                Grad_bias_data,
                                                                                                [D]}}}
                                                                                end
                                                                            )
                                                                        end
                                                                    )
                                                                end
                                                            )
                                                        end
                                                    )
                                                end
                                            )
                                        end
                                    )
                                end
                            )
                        end)
                end
            )
        end
    ).

-file("src/viva_tensor/nn/backward.gleam", 577).
?DOC(false).
-spec rms_norm_backward(
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor(),
    float()
) -> {ok, {viva_tensor@tensor:tensor(), viva_tensor@tensor:tensor()}} |
    {error, viva_tensor@core@error:tensor_error()}.
rms_norm_backward(Grad_out, Input, Scale, Rms, _) ->
    Input_shape = viva_tensor@tensor:shape(Input),
    gleam@result:'try'(
        case viva_tensor@tensor:shape(Grad_out) =:= Input_shape of
            true ->
                {ok, nil};

            false ->
                {error,
                    {shape_mismatch,
                        Input_shape,
                        viva_tensor@tensor:shape(Grad_out)}}
        end,
        fun(_) ->
            gleam@result:'try'(
                last_dim_of(Input_shape),
                fun(D) ->
                    Scale_shape = viva_tensor@tensor:shape(Scale),
                    gleam@result:'try'(case Scale_shape =:= [D] of
                            true ->
                                {ok, nil};

                            false ->
                                {error, {shape_mismatch, [D], Scale_shape}}
                        end, fun(_) ->
                            Outer = case gleam@int:max(D, 1) of
                                0 -> 0;
                                Gleam@denominator -> product(Input_shape) div Gleam@denominator
                            end,
                            gleam@result:'try'(
                                check_stat_size(
                                    <<"rms_norm_backward.rms"/utf8>>,
                                    Rms,
                                    Outer
                                ),
                                fun(_) ->
                                    gleam@result:'try'(
                                        viva_tensor@tensor:try_to_list(Input),
                                        fun(Data) ->
                                            gleam@result:'try'(
                                                viva_tensor@tensor:try_to_list(
                                                    Grad_out
                                                ),
                                                fun(Grad_data) ->
                                                    gleam@result:'try'(
                                                        viva_tensor@tensor:try_to_list(
                                                            Scale
                                                        ),
                                                        fun(Scale_data) ->
                                                            gleam@result:'try'(
                                                                viva_tensor@tensor:try_to_list(
                                                                    Rms
                                                                ),
                                                                fun(Rms_data) ->
                                                                    Input_rows = chunk_every(
                                                                        Data,
                                                                        D
                                                                    ),
                                                                    Grad_rows = chunk_every(
                                                                        Grad_data,
                                                                        D
                                                                    ),
                                                                    Combined = gleam@list:zip(
                                                                        Input_rows,
                                                                        gleam@list:zip(
                                                                            Grad_rows,
                                                                            Rms_data
                                                                        )
                                                                    ),
                                                                    Init_acc = {[],
                                                                        gleam@list:repeat(
                                                                            +0.0,
                                                                            D
                                                                        )},
                                                                    {Rev_grad_x,
                                                                        Grad_scale_data} = gleam@list:fold(
                                                                        Combined,
                                                                        Init_acc,
                                                                        fun(
                                                                            Acc,
                                                                            Row
                                                                        ) ->
                                                                            {Rev_gx,
                                                                                Gs} = Acc,
                                                                            {X_row,
                                                                                Rest} = Row,
                                                                            {G_row,
                                                                                R} = Rest,
                                                                            Inv_r = case R of
                                                                                +0.0 -> +0.0;
                                                                                -0.0 -> -0.0;
                                                                                Gleam@denominator@1 -> 1.0
                                                                                / Gleam@denominator@1
                                                                            end,
                                                                            Inv_r2 = Inv_r
                                                                            * Inv_r,
                                                                            G_scaled = gleam@list:map(
                                                                                gleam@list:zip(
                                                                                    G_row,
                                                                                    Scale_data
                                                                                ),
                                                                                fun(
                                                                                    P
                                                                                ) ->
                                                                                    {G,
                                                                                        S} = P,
                                                                                    G
                                                                                    * S
                                                                                end
                                                                            ),
                                                                            D_f = erlang:float(
                                                                                D
                                                                            ),
                                                                            Dot = case D_f of
                                                                                +0.0 -> +0.0;
                                                                                -0.0 -> -0.0;
                                                                                Gleam@denominator@2 -> gleam@list:fold(
                                                                                    gleam@list:zip(
                                                                                        G_scaled,
                                                                                        X_row
                                                                                    ),
                                                                                    +0.0,
                                                                                    fun(
                                                                                        A,
                                                                                        P@1
                                                                                    ) ->
                                                                                        {Gs2,
                                                                                            X} = P@1,
                                                                                        A
                                                                                        + (Gs2
                                                                                        * X)
                                                                                    end
                                                                                )
                                                                                / Gleam@denominator@2
                                                                            end,
                                                                            Grad_x_slice = gleam@list:map(
                                                                                gleam@list:zip(
                                                                                    G_scaled,
                                                                                    X_row
                                                                                ),
                                                                                fun(
                                                                                    P@2
                                                                                ) ->
                                                                                    {Gs2@1,
                                                                                        X@1} = P@2,
                                                                                    (Gs2@1
                                                                                    - ((X@1
                                                                                    * Dot)
                                                                                    * Inv_r2))
                                                                                    * Inv_r
                                                                                end
                                                                            ),
                                                                            New_gs = gleam@list:map(
                                                                                gleam@list:zip(
                                                                                    Gs,
                                                                                    gleam@list:zip(
                                                                                        G_row,
                                                                                        X_row
                                                                                    )
                                                                                ),
                                                                                fun(
                                                                                    T
                                                                                ) ->
                                                                                    {Acc_v,
                                                                                        Rest2} = T,
                                                                                    {G@1,
                                                                                        X@2} = Rest2,
                                                                                    Acc_v
                                                                                    + ((G@1
                                                                                    * X@2)
                                                                                    * Inv_r)
                                                                                end
                                                                            ),
                                                                            {[Grad_x_slice |
                                                                                    Rev_gx],
                                                                                New_gs}
                                                                        end
                                                                    ),
                                                                    Grad_x_data = begin
                                                                        _pipe = Rev_grad_x,
                                                                        _pipe@1 = lists:reverse(
                                                                            _pipe
                                                                        ),
                                                                        lists:append(
                                                                            _pipe@1
                                                                        )
                                                                    end,
                                                                    {ok,
                                                                        {{tensor,
                                                                                Grad_x_data,
                                                                                Input_shape},
                                                                            {tensor,
                                                                                Grad_scale_data,
                                                                                [D]}}}
                                                                end
                                                            )
                                                        end
                                                    )
                                                end
                                            )
                                        end
                                    )
                                end
                            )
                        end)
                end
            )
        end
    ).

-file("src/viva_tensor/nn/backward.gleam", 904).
?DOC(false).
-spec range_loop(integer(), integer(), list(integer())) -> list(integer()).
range_loop(From, To, Acc) ->
    case From > To of
        true ->
            lists:reverse(Acc);

        false ->
            range_loop(From + 1, To, [From | Acc])
    end.

-file("src/viva_tensor/nn/backward.gleam", 900).
?DOC(false).
-spec range_int(integer(), integer()) -> list(integer()).
range_int(From, To) ->
    range_loop(From, To, []).

-file("src/viva_tensor/nn/backward.gleam", 712).
?DOC(false).
-spec softmax_backward_data(
    list(float()),
    list(float()),
    integer(),
    integer(),
    integer()
) -> list(float()).
softmax_backward_data(Grad, Out, Outer, Axis_size, Inner) ->
    Grad_arr = viva_tensor@core@ffi:list_to_array(Grad),
    Out_arr = viva_tensor@core@ffi:list_to_array(Out),
    _pipe = range_int(0, Outer - 1),
    gleam@list:flat_map(
        _pipe,
        fun(O) ->
            Outer_offset = (O * Axis_size) * Inner,
            Sums = begin
                _pipe@1 = range_int(0, Inner - 1),
                gleam@list:map(
                    _pipe@1,
                    fun(Inner_idx) -> _pipe@2 = range_int(0, Axis_size - 1),
                        gleam@list:fold(
                            _pipe@2,
                            +0.0,
                            fun(Acc, K) ->
                                Idx = (Outer_offset + (K * Inner)) + Inner_idx,
                                Acc + (viva_tensor@core@ffi:array_get(
                                    Grad_arr,
                                    Idx
                                )
                                * viva_tensor@core@ffi:array_get(Out_arr, Idx))
                            end
                        ) end
                )
            end,
            Sums_arr = viva_tensor@core@ffi:list_to_array(Sums),
            _pipe@3 = range_int(0, Axis_size - 1),
            gleam@list:flat_map(
                _pipe@3,
                fun(K@1) -> _pipe@4 = range_int(0, Inner - 1),
                    gleam@list:map(
                        _pipe@4,
                        fun(Inner_idx@1) ->
                            Idx@1 = (Outer_offset + (K@1 * Inner)) + Inner_idx@1,
                            G = viva_tensor@core@ffi:array_get(Grad_arr, Idx@1),
                            Y = viva_tensor@core@ffi:array_get(Out_arr, Idx@1),
                            S = viva_tensor@core@ffi:array_get(
                                Sums_arr,
                                Inner_idx@1
                            ),
                            Y * (G - S)
                        end
                    ) end
            )
        end
    ).

-file("src/viva_tensor/nn/backward.gleam", 848).
?DOC(false).
-spec nth_dim(list(integer()), integer()) -> integer().
nth_dim(Shape, Idx) ->
    case {Shape, Idx} of
        {[], _} ->
            0;

        {[D | _], 0} ->
            D;

        {[_ | Rest], I} ->
            nth_dim(Rest, I - 1)
    end.

-file("src/viva_tensor/nn/backward.gleam", 669).
?DOC(false).
-spec softmax_backward(
    viva_tensor@tensor:tensor(),
    viva_tensor@tensor:tensor(),
    integer()
) -> {ok, viva_tensor@tensor:tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
softmax_backward(Grad_out, Output, Axis) ->
    Shp = viva_tensor@tensor:shape(Output),
    gleam@result:'try'(case viva_tensor@tensor:shape(Grad_out) =:= Shp of
            true ->
                {ok, nil};

            false ->
                {error,
                    {shape_mismatch, Shp, viva_tensor@tensor:shape(Grad_out)}}
        end, fun(_) ->
            Rnk = erlang:length(Shp),
            case (Axis >= 0) andalso (Axis < Rnk) of
                false ->
                    {error,
                        {dimension_error,
                            <<"Invalid axis for softmax_backward"/utf8>>}};

                true ->
                    Axis_size = nth_dim(Shp, Axis),
                    Inner = begin
                        _pipe = Shp,
                        _pipe@1 = gleam@list:drop(_pipe, Axis + 1),
                        product(_pipe@1)
                    end,
                    Outer = begin
                        _pipe@2 = Shp,
                        _pipe@3 = gleam@list:take(_pipe@2, Axis),
                        product(_pipe@3)
                    end,
                    case Axis_size =< 0 of
                        true ->
                            {ok, {tensor, [], Shp}};

                        false ->
                            gleam@result:'try'(
                                viva_tensor@tensor:try_to_list(Grad_out),
                                fun(Grad_data) ->
                                    gleam@result:'try'(
                                        viva_tensor@tensor:try_to_list(Output),
                                        fun(Out_data) ->
                                            Buffer = softmax_backward_data(
                                                Grad_data,
                                                Out_data,
                                                Outer,
                                                Axis_size,
                                                Inner
                                            ),
                                            {ok, {tensor, Buffer, Shp}}
                                        end
                                    )
                                end
                            )
                    end
            end
        end).