Current section

Files

Jump to
viva_tensor src viva_tensor@nn@backward.erl
Raw

src/viva_tensor@nn@backward.erl

-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).