Packages
Tensor library for Gleam/BEAM with a pure Gleam API, zero-copy views, and optional native acceleration
Current section
Files
Jump to
Current section
Files
src/viva_tensor@nn@norm.erl
-module(viva_tensor@nn@norm).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_tensor/nn/norm.gleam").
-export([layer_norm_init_with_eps/2, layer_norm_init/1, layer_norm_forward/2, rms_norm_init_with_eps/2, rms_norm_init/1, rms_norm_forward/2, batch_norm_1d_init/1, batch_norm_1d_forward/3, group_norm_init/2, group_norm_forward/2]).
-export_type([layer_norm/0, rms_norm/0, batch_norm1d/0, group_norm/0]).
-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).
-type layer_norm() :: {layer_norm,
viva_tensor@tensor:tensor(),
viva_tensor@tensor:tensor(),
float()}.
-type rms_norm() :: {rms_norm, viva_tensor@tensor:tensor(), float()}.
-type batch_norm1d() :: {batch_norm1d,
viva_tensor@tensor:tensor(),
viva_tensor@tensor:tensor(),
viva_tensor@tensor:tensor(),
viva_tensor@tensor:tensor(),
float(),
float()}.
-type group_norm() :: {group_norm,
integer(),
viva_tensor@tensor:tensor(),
viva_tensor@tensor:tensor(),
float()}.
-file("src/viva_tensor/nn/norm.gleam", 64).
?DOC(false).
-spec layer_norm_init_with_eps(integer(), float()) -> layer_norm().
layer_norm_init_with_eps(Num_features, Eps) ->
{layer_norm,
viva_tensor@tensor:ones([Num_features]),
viva_tensor@tensor:zeros([Num_features]),
Eps}.
-file("src/viva_tensor/nn/norm.gleam", 52).
?DOC(false).
-spec layer_norm_init(integer()) -> layer_norm().
layer_norm_init(Num_features) ->
layer_norm_init_with_eps(Num_features, 1.0e-5).
-file("src/viva_tensor/nn/norm.gleam", 577).
?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/norm.gleam", 556).
?DOC(false).
-spec list_variance(list(float()), float()) -> float().
list_variance(Xs, Mean) ->
N = erlang:length(Xs),
case N of
0 ->
+0.0;
_ ->
case erlang:float(N) of
+0.0 -> +0.0;
-0.0 -> -0.0;
Gleam@denominator -> gleam@list:fold(
Xs,
+0.0,
fun(Acc, X) ->
D = X - Mean,
Acc + (D * D)
end
)
/ Gleam@denominator
end
end.
-file("src/viva_tensor/nn/norm.gleam", 548).
?DOC(false).
-spec list_mean(list(float())) -> float().
list_mean(Xs) ->
N = erlang:length(Xs),
case N of
0 ->
+0.0;
_ ->
case erlang:float(N) of
+0.0 -> +0.0;
-0.0 -> -0.0;
Gleam@denominator -> gleam@list:fold(
Xs,
+0.0,
fun(Acc, X) -> Acc + X end
)
/ Gleam@denominator
end
end.
-file("src/viva_tensor/nn/norm.gleam", 533).
?DOC(false).
-spec chunk_by_loop(list(float()), integer(), list(list(float()))) -> list(list(float())).
chunk_by_loop(Data, N, Acc) ->
case Data of
[] ->
lists:reverse(Acc);
_ ->
Chunk = gleam@list:take(Data, N),
Rest = gleam@list:drop(Data, N),
chunk_by_loop(Rest, N, [Chunk | Acc])
end.
-file("src/viva_tensor/nn/norm.gleam", 526).
?DOC(false).
-spec chunk_by(list(float()), integer()) -> list(list(float())).
chunk_by(Data, N) ->
case N =< 0 of
true ->
[];
false ->
chunk_by_loop(Data, N, [])
end.
-file("src/viva_tensor/nn/norm.gleam", 519).
?DOC(false).
-spec last_dim(list(integer())) -> {ok, integer()} |
{error, viva_tensor@core@error:tensor_error()}.
last_dim(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/norm.gleam", 87).
?DOC(false).
-spec layer_norm_forward(layer_norm(), viva_tensor@tensor:tensor()) -> {ok,
viva_tensor@tensor:tensor()} |
{error, viva_tensor@core@error:tensor_error()}.
layer_norm_forward(Layer, Input) ->
Input_shape = viva_tensor@tensor:shape(Input),
Scale_shape = viva_tensor@tensor:shape(erlang:element(2, Layer)),
gleam@result:'try'(
last_dim(Scale_shape),
fun(Num_features) ->
gleam@result:'try'(
last_dim(Input_shape),
fun(Input_last) -> case Input_last =:= Num_features of
false ->
{error,
{shape_mismatch, [Num_features], [Input_last]}};
true ->
gleam@result:'try'(
viva_tensor@tensor:try_to_list(Input),
fun(Data) ->
gleam@result:'try'(
viva_tensor@tensor:try_to_list(
erlang:element(2, Layer)
),
fun(Scale_data) ->
gleam@result:'try'(
viva_tensor@tensor:try_to_list(
erlang:element(3, Layer)
),
fun(Bias_data) ->
Normalized = begin
_pipe = chunk_by(
Data,
Num_features
),
_pipe@1 = gleam@list:map(
_pipe,
fun(Chunk) ->
Mean = list_mean(
Chunk
),
Var = list_variance(
Chunk,
Mean
),
Denom = safe_sqrt(
Var + erlang:element(
4,
Layer
)
),
gleam@list:map(
gleam@list:zip(
Chunk,
gleam@list:zip(
Scale_data,
Bias_data
)
),
fun(T) ->
{X, Sb} = T,
{S, B} = Sb,
((case Denom of
+0.0 -> +0.0;
-0.0 -> -0.0;
Gleam@denominator -> (X
- Mean)
/ Gleam@denominator
end)
* S)
+ B
end
)
end
),
lists:append(_pipe@1)
end,
{ok,
{tensor,
Normalized,
Input_shape}}
end
)
end
)
end
)
end end
)
end
).
-file("src/viva_tensor/nn/norm.gleam", 153).
?DOC(false).
-spec rms_norm_init_with_eps(integer(), float()) -> rms_norm().
rms_norm_init_with_eps(Num_features, Eps) ->
{rms_norm, viva_tensor@tensor:ones([Num_features]), Eps}.
-file("src/viva_tensor/nn/norm.gleam", 141).
?DOC(false).
-spec rms_norm_init(integer()) -> rms_norm().
rms_norm_init(Num_features) ->
rms_norm_init_with_eps(Num_features, 1.0e-6).
-file("src/viva_tensor/nn/norm.gleam", 569).
?DOC(false).
-spec list_mean_squares(list(float())) -> float().
list_mean_squares(Xs) ->
N = erlang:length(Xs),
case N of
0 ->
+0.0;
_ ->
case erlang:float(N) of
+0.0 -> +0.0;
-0.0 -> -0.0;
Gleam@denominator -> gleam@list:fold(
Xs,
+0.0,
fun(Acc, X) -> Acc + (X * X) end
)
/ Gleam@denominator
end
end.
-file("src/viva_tensor/nn/norm.gleam", 170).
?DOC(false).
-spec rms_norm_forward(rms_norm(), viva_tensor@tensor:tensor()) -> {ok,
viva_tensor@tensor:tensor()} |
{error, viva_tensor@core@error:tensor_error()}.
rms_norm_forward(Layer, Input) ->
Input_shape = viva_tensor@tensor:shape(Input),
Scale_shape = viva_tensor@tensor:shape(erlang:element(2, Layer)),
gleam@result:'try'(
last_dim(Scale_shape),
fun(Num_features) ->
gleam@result:'try'(
last_dim(Input_shape),
fun(Input_last) -> case Input_last =:= Num_features of
false ->
{error,
{shape_mismatch, [Num_features], [Input_last]}};
true ->
gleam@result:'try'(
viva_tensor@tensor:try_to_list(Input),
fun(Data) ->
gleam@result:'try'(
viva_tensor@tensor:try_to_list(
erlang:element(2, Layer)
),
fun(Scale_data) ->
Normalized = begin
_pipe = chunk_by(
Data,
Num_features
),
_pipe@1 = gleam@list:map(
_pipe,
fun(Chunk) ->
Mean_sq = list_mean_squares(
Chunk
),
Rms = safe_sqrt(
Mean_sq + erlang:element(
3,
Layer
)
),
gleam@list:map(
gleam@list:zip(
Chunk,
Scale_data
),
fun(T) ->
{X, S} = T,
(case Rms of
+0.0 -> +0.0;
-0.0 -> -0.0;
Gleam@denominator -> X
/ Gleam@denominator
end) * S
end
)
end
),
lists:append(_pipe@1)
end,
{ok,
{tensor,
Normalized,
Input_shape}}
end
)
end
)
end end
)
end
).
-file("src/viva_tensor/nn/norm.gleam", 236).
?DOC(false).
-spec batch_norm_1d_init(integer()) -> batch_norm1d().
batch_norm_1d_init(Num_features) ->
{batch_norm1d,
viva_tensor@tensor:ones([Num_features]),
viva_tensor@tensor:zeros([Num_features]),
viva_tensor@tensor:zeros([Num_features]),
viva_tensor@tensor:ones([Num_features]),
0.1,
1.0e-5}.
-file("src/viva_tensor/nn/norm.gleam", 645).
?DOC(false).
-spec join_with(list(binary()), binary()) -> binary().
join_with(Parts, Sep) ->
case Parts of
[] ->
<<""/utf8>>;
[S] ->
S;
[S@1 | Rest] ->
<<<<S@1/binary, Sep/binary>>/binary, (join_with(Rest, Sep))/binary>>
end.
-file("src/viva_tensor/nn/norm.gleam", 641).
?DOC(false).
-spec shape_to_string(list(integer())) -> binary().
shape_to_string(Shape) ->
<<<<"["/utf8,
(join_with(
gleam@list:map(Shape, fun erlang:integer_to_binary/1),
<<", "/utf8>>
))/binary>>/binary,
"]"/utf8>>.
-file("src/viva_tensor/nn/norm.gleam", 621).
?DOC(false).
-spec ema_update(list(float()), list(float()), float()) -> list(float()).
ema_update(Running, Batch, Momentum) ->
One_minus = 1.0 - Momentum,
gleam@list:map(
gleam@list:zip(Running, Batch),
fun(P) ->
{R, B} = P,
(One_minus * R) + (Momentum * B)
end
).
-file("src/viva_tensor/nn/norm.gleam", 601).
?DOC(false).
-spec column_variances(list(list(float())), list(float()), integer(), integer()) -> list(float()).
column_variances(Rows, Means, Features, Batch) ->
Zeros = gleam@list:repeat(+0.0, Features),
Sumsq = gleam@list:fold(
Rows,
Zeros,
fun(Acc, Row) ->
gleam@list:map(
gleam@list:zip(Acc, gleam@list:zip(Row, Means)),
fun(P) ->
{A, Rest} = P,
{X, Mu} = Rest,
D = X - Mu,
A + (D * D)
end
)
end
),
N = erlang:float(Batch),
gleam@list:map(Sumsq, fun(S) -> case N of
+0.0 -> +0.0;
-0.0 -> -0.0;
Gleam@denominator -> S / Gleam@denominator
end end).
-file("src/viva_tensor/nn/norm.gleam", 584).
?DOC(false).
-spec column_means(list(list(float())), integer(), integer()) -> list(float()).
column_means(Rows, Features, Batch) ->
Zeros = gleam@list:repeat(+0.0, Features),
Sums = gleam@list:fold(
Rows,
Zeros,
fun(Acc, Row) ->
gleam@list:map(
gleam@list:zip(Acc, Row),
fun(P) ->
{A, B} = P,
A + B
end
)
end
),
N = erlang:float(Batch),
gleam@list:map(Sums, fun(S) -> case N of
+0.0 -> +0.0;
-0.0 -> -0.0;
Gleam@denominator -> S / Gleam@denominator
end end).
-file("src/viva_tensor/nn/norm.gleam", 307).
?DOC(false).
-spec batch_norm_1d_apply(
batch_norm1d(),
viva_tensor@tensor:tensor(),
integer(),
integer(),
boolean()
) -> {ok, {batch_norm1d(), viva_tensor@tensor:tensor()}} |
{error, viva_tensor@core@error:tensor_error()}.
batch_norm_1d_apply(Layer, Input, Batch, Features, Training) ->
gleam@result:'try'(
viva_tensor@tensor:try_to_list(Input),
fun(Data) ->
gleam@result:'try'(
viva_tensor@tensor:try_to_list(erlang:element(2, Layer)),
fun(Scale_data) ->
gleam@result:'try'(
viva_tensor@tensor:try_to_list(erlang:element(3, Layer)),
fun(Bias_data) ->
gleam@result:'try'(
viva_tensor@tensor:try_to_list(
erlang:element(4, Layer)
),
fun(Running_mean_data) ->
gleam@result:'try'(
viva_tensor@tensor:try_to_list(
erlang:element(5, Layer)
),
fun(Running_var_data) ->
Rows = chunk_by(Data, Features),
Batch_mean = column_means(
Rows,
Features,
Batch
),
Batch_var = column_variances(
Rows,
Batch_mean,
Features,
Batch
),
{Use_mean, Use_var} = case Training of
true ->
{Batch_mean, Batch_var};
false ->
{Running_mean_data,
Running_var_data}
end,
Normalized_rows = gleam@list:map(
Rows,
fun(Row) ->
gleam@list:map(
gleam@list:zip(
Row,
gleam@list:zip(
Use_mean,
gleam@list:zip(
Use_var,
gleam@list:zip(
Scale_data,
Bias_data
)
)
)
),
fun(T) ->
{X, Rest} = T,
{Mu, Rest2} = Rest,
{Var, Rest3} = Rest2,
{S, B} = Rest3,
Denom = safe_sqrt(
Var + erlang:element(
7,
Layer
)
),
((case Denom of
+0.0 -> +0.0;
-0.0 -> -0.0;
Gleam@denominator -> (X
- Mu)
/ Gleam@denominator
end) * S) + B
end
)
end
),
Out_data = lists:append(
Normalized_rows
),
Out = {tensor,
Out_data,
[Batch, Features]},
Updated_layer = case Training of
false ->
Layer;
true ->
New_mean = ema_update(
Running_mean_data,
Batch_mean,
erlang:element(6, Layer)
),
New_var = ema_update(
Running_var_data,
Batch_var,
erlang:element(6, Layer)
),
{batch_norm1d,
erlang:element(2, Layer),
erlang:element(3, Layer),
{tensor,
New_mean,
[Features]},
{tensor,
New_var,
[Features]},
erlang:element(6, Layer),
erlang:element(7, Layer)}
end,
{ok, {Updated_layer, Out}}
end
)
end
)
end
)
end
)
end
).
-file("src/viva_tensor/nn/norm.gleam", 270).
?DOC(false).
-spec batch_norm_1d_forward(
batch_norm1d(),
viva_tensor@tensor:tensor(),
boolean()
) -> {ok, {batch_norm1d(), viva_tensor@tensor:tensor()}} |
{error, viva_tensor@core@error:tensor_error()}.
batch_norm_1d_forward(Layer, Input, Training) ->
Scale_shape = viva_tensor@tensor:shape(erlang:element(2, Layer)),
gleam@result:'try'(
last_dim(Scale_shape),
fun(Num_features) -> case viva_tensor@tensor:shape(Input) of
[Batch, Features] ->
case Features =:= Num_features of
false ->
{error,
{shape_mismatch,
[Batch, Num_features],
[Batch, Features]}};
true ->
case Batch =< 0 of
true ->
{error,
{invalid_shape,
<<"batch_norm_1d: batch dimension must be positive"/utf8>>}};
false ->
batch_norm_1d_apply(
Layer,
Input,
Batch,
Features,
Training
)
end
end;
Other ->
{error,
{invalid_shape,
<<"batch_norm_1d expects 2D input [batch, num_features], got "/utf8,
(shape_to_string(Other))/binary>>}}
end end
).
-file("src/viva_tensor/nn/norm.gleam", 401).
?DOC(false).
-spec group_norm_init(integer(), integer()) -> group_norm().
group_norm_init(Num_groups, Num_channels) ->
{group_norm,
Num_groups,
viva_tensor@tensor:ones([Num_channels]),
viva_tensor@tensor:zeros([Num_channels]),
1.0e-5}.
-file("src/viva_tensor/nn/norm.gleam", 633).
?DOC(false).
-spec list_at(list(float()), integer()) -> float().
list_at(Xs, Idx) ->
case {Xs, Idx} of
{[], _} ->
+0.0;
{[X | _], 0} ->
X;
{[_ | Rest], I} ->
list_at(Rest, I - 1)
end.
-file("src/viva_tensor/nn/norm.gleam", 460).
?DOC(false).
-spec group_norm_apply(
group_norm(),
viva_tensor@tensor:tensor(),
integer(),
integer(),
integer(),
integer()
) -> {ok, viva_tensor@tensor:tensor()} |
{error, viva_tensor@core@error:tensor_error()}.
group_norm_apply(Layer, Input, Batch, Channels, Spatial, Num_channels) ->
case Channels =:= Num_channels of
false ->
{error, {shape_mismatch, [Batch, Num_channels], [Batch, Channels]}};
true ->
gleam@result:'try'(
viva_tensor@tensor:try_to_list(Input),
fun(Data) ->
gleam@result:'try'(
viva_tensor@tensor:try_to_list(erlang:element(3, Layer)),
fun(Scale_data) ->
gleam@result:'try'(
viva_tensor@tensor:try_to_list(
erlang:element(4, Layer)
),
fun(Bias_data) ->
Channels_per_group = case erlang:element(
2,
Layer
) of
0 -> 0;
Gleam@denominator -> Channels div Gleam@denominator
end,
Group_size = Channels_per_group * Spatial,
Row_size = Channels * Spatial,
Rows = chunk_by(Data, Row_size),
Normalized_rows = gleam@list:map(
Rows,
fun(Row) ->
Groups = chunk_by(Row, Group_size),
Normalized_groups = gleam@list:index_map(
Groups,
fun(Group_data, Group_idx) ->
Mean = list_mean(Group_data),
Var = list_variance(
Group_data,
Mean
),
Denom = safe_sqrt(
Var + erlang:element(
5,
Layer
)
),
Channels_in_group = chunk_by(
Group_data,
Spatial
),
_pipe = gleam@list:index_map(
Channels_in_group,
fun(
Channel_data,
C_in_group
) ->
Global_channel = (Group_idx
* Channels_per_group)
+ C_in_group,
S = list_at(
Scale_data,
Global_channel
),
B = list_at(
Bias_data,
Global_channel
),
gleam@list:map(
Channel_data,
fun(X) ->
((case Denom of
+0.0 -> +0.0;
-0.0 -> -0.0;
Gleam@denominator@1 -> (X
- Mean)
/ Gleam@denominator@1
end)
* S)
+ B
end
)
end
),
lists:append(_pipe)
end
),
lists:append(Normalized_groups)
end
),
{ok,
{tensor,
lists:append(Normalized_rows),
viva_tensor@tensor:shape(Input)}}
end
)
end
)
end
)
end.
-file("src/viva_tensor/nn/norm.gleam", 429).
?DOC(false).
-spec group_norm_forward(group_norm(), viva_tensor@tensor:tensor()) -> {ok,
viva_tensor@tensor:tensor()} |
{error, viva_tensor@core@error:tensor_error()}.
group_norm_forward(Layer, Input) ->
Scale_shape = viva_tensor@tensor:shape(erlang:element(3, Layer)),
gleam@result:'try'(
last_dim(Scale_shape),
fun(Num_channels) -> case (case erlang:element(2, Layer) of
0 -> 0;
Gleam@denominator -> Num_channels rem Gleam@denominator
end) =:= 0 of
false ->
{error,
{invalid_shape,
<<<<<<<<"group_norm: channels ("/utf8,
(erlang:integer_to_binary(
Num_channels
))/binary>>/binary,
") not divisible by groups ("/utf8>>/binary,
(erlang:integer_to_binary(
erlang:element(2, Layer)
))/binary>>/binary,
")"/utf8>>}};
true ->
case viva_tensor@tensor:shape(Input) of
[Batch, Channels] ->
group_norm_apply(
Layer,
Input,
Batch,
Channels,
1,
Num_channels
);
[Batch@1, Channels@1, Spatial] ->
group_norm_apply(
Layer,
Input,
Batch@1,
Channels@1,
Spatial,
Num_channels
);
Other ->
{error,
{invalid_shape,
<<"group_norm expects [batch, channels] or [batch, channels, spatial], got "/utf8,
(shape_to_string(Other))/binary>>}}
end
end end
).