viva_tensor

-module(viva_tensor@nn@scheduler).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_tensor/nn/scheduler.gleam").
-export([step_lr/3, cosine_annealing_lr/3, linear_warmup/2, one_cycle_lr/4, exponential_lr/2, scheduler_lr/1, scheduler_step/1, apply_to_optimizer/2]).
-export_type([scheduler_kind/0, scheduler/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 scheduler_kind() :: step_lr |
    cosine_annealing_lr |
    linear_warmup |
    one_cycle_lr |
    exponential_lr.

-type scheduler() :: {scheduler,
        scheduler_kind(),
        float(),
        integer(),
        integer(),
        float(),
        integer(),
        float(),
        integer(),
        float(),
        integer(),
        float()}.

-file("src/viva_tensor/nn/scheduler.gleam", 68).
?DOC(false).
-spec step_lr(float(), integer(), float()) -> scheduler().
step_lr(Base_lr, Step_size, Gamma) ->
    {scheduler,
        step_lr,
        Base_lr,
        0,
        Step_size,
        Gamma,
        0,
        +0.0,
        0,
        +0.0,
        0,
        +0.0}.

-file("src/viva_tensor/nn/scheduler.gleam", 90).
?DOC(false).
-spec cosine_annealing_lr(float(), integer(), float()) -> scheduler().
cosine_annealing_lr(Base_lr, T_max, Eta_min) ->
    {scheduler,
        cosine_annealing_lr,
        Base_lr,
        0,
        0,
        +0.0,
        T_max,
        Eta_min,
        0,
        +0.0,
        0,
        +0.0}.

-file("src/viva_tensor/nn/scheduler.gleam", 116).
?DOC(false).
-spec linear_warmup(float(), integer()) -> scheduler().
linear_warmup(Base_lr, Warmup_steps) ->
    {scheduler,
        linear_warmup,
        Base_lr,
        0,
        0,
        +0.0,
        0,
        +0.0,
        Warmup_steps,
        +0.0,
        0,
        +0.0}.

-file("src/viva_tensor/nn/scheduler.gleam", 143).
?DOC(false).
-spec one_cycle_lr(float(), float(), integer(), float()) -> scheduler().
one_cycle_lr(Base_lr, Max_lr, Total_steps, Pct_start) ->
    {scheduler,
        one_cycle_lr,
        Base_lr,
        0,
        0,
        +0.0,
        0,
        +0.0,
        0,
        Max_lr,
        Total_steps,
        Pct_start}.

-file("src/viva_tensor/nn/scheduler.gleam", 170).
?DOC(false).
-spec exponential_lr(float(), float()) -> scheduler().
exponential_lr(Base_lr, Gamma) ->
    {scheduler, exponential_lr, Base_lr, 0, 0, Gamma, 0, +0.0, 0, +0.0, 0, +0.0}.

-file("src/viva_tensor/nn/scheduler.gleam", 274).
?DOC(false).
-spec pow(float(), float()) -> float().
pow(Base, Exp) ->
    case gleam@float:power(Base, Exp) of
        {ok, V} ->
            V;

        {error, _} ->
            +0.0
    end.

-file("src/viva_tensor/nn/scheduler.gleam", 270).
?DOC(false).
-spec exponential_value(scheduler()) -> float().
exponential_value(S) ->
    erlang:element(3, S) * pow(
        erlang:element(6, S),
        erlang:float(erlang:element(4, S))
    ).

-file("src/viva_tensor/nn/scheduler.gleam", 246).
?DOC(false).
-spec one_cycle_value(scheduler()) -> float().
one_cycle_value(S) ->
    Total_f = erlang:float(erlang:element(11, S)),
    Warmup_end_f = erlang:element(12, S) * Total_f,
    Step_f = erlang:float(erlang:element(4, S)),
    case Step_f =< Warmup_end_f of
        true ->
            Ramp = case Warmup_end_f > +0.0 of
                true ->
                    case Warmup_end_f of
                        +0.0 -> +0.0;
                        -0.0 -> -0.0;
                        Gleam@denominator -> Step_f / Gleam@denominator
                    end;

                false ->
                    1.0
            end,
            erlang:element(3, S) + ((erlang:element(10, S) - erlang:element(
                3,
                S
            ))
            * Ramp);

        false ->
            Denom = Total_f - Warmup_end_f,
            Progress = case Denom > +0.0 of
                true ->
                    case Denom of
                        +0.0 -> +0.0;
                        -0.0 -> -0.0;
                        Gleam@denominator@1 -> (Step_f - Warmup_end_f) / Gleam@denominator@1
                    end;

                false ->
                    1.0
            end,
            Cos_term = gleam_community@maths:cos(
                gleam_community@maths:pi() * Progress
            ),
            erlang:element(3, S) + ((0.5 * (erlang:element(10, S) - erlang:element(
                3,
                S
            )))
            * (1.0 + Cos_term))
    end.

-file("src/viva_tensor/nn/scheduler.gleam", 235).
?DOC(false).
-spec warmup_value(scheduler()) -> float().
warmup_value(S) ->
    case erlang:element(9, S) of
        0 ->
            erlang:element(3, S);

        W ->
            Frac = case erlang:float(W) of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator -> erlang:float(erlang:element(4, S)) / Gleam@denominator
            end,
            Clamped = gleam@float:min(Frac, 1.0),
            erlang:element(3, S) * Clamped
    end.

-file("src/viva_tensor/nn/scheduler.gleam", 226).
?DOC(false).
-spec cosine_value(scheduler()) -> float().
cosine_value(S) ->
    Progress = case erlang:element(7, S) of
        0 ->
            +0.0;

        T ->
            case erlang:float(T) of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator -> erlang:float(erlang:element(4, S)) / Gleam@denominator
            end
    end,
    Cos_term = gleam_community@maths:cos(gleam_community@maths:pi() * Progress),
    erlang:element(8, S) + ((0.5 * (erlang:element(3, S) - erlang:element(8, S)))
    * (1.0 + Cos_term)).

-file("src/viva_tensor/nn/scheduler.gleam", 218).
?DOC(false).
-spec step_lr_value(scheduler()) -> float().
step_lr_value(S) ->
    Drops = case erlang:element(5, S) of
        0 ->
            0;

        Size ->
            case Size of
                0 -> 0;
                Gleam@denominator -> erlang:element(4, S) div Gleam@denominator
            end
    end,
    erlang:element(3, S) * pow(erlang:element(6, S), erlang:float(Drops)).

-file("src/viva_tensor/nn/scheduler.gleam", 190).
?DOC(false).
-spec scheduler_lr(scheduler()) -> float().
scheduler_lr(S) ->
    case erlang:element(2, S) of
        step_lr ->
            step_lr_value(S);

        cosine_annealing_lr ->
            cosine_value(S);

        linear_warmup ->
            warmup_value(S);

        one_cycle_lr ->
            one_cycle_value(S);

        exponential_lr ->
            exponential_value(S)
    end.

-file("src/viva_tensor/nn/scheduler.gleam", 201).
?DOC(false).
-spec scheduler_step(scheduler()) -> {scheduler(), float()}.
scheduler_step(S) ->
    Next = {scheduler,
        erlang:element(2, S),
        erlang:element(3, S),
        erlang:element(4, S) + 1,
        erlang:element(5, S),
        erlang:element(6, S),
        erlang:element(7, S),
        erlang:element(8, S),
        erlang:element(9, S),
        erlang:element(10, S),
        erlang:element(11, S),
        erlang:element(12, S)},
    {Next, scheduler_lr(Next)}.

-file("src/viva_tensor/nn/scheduler.gleam", 208).
?DOC(false).
-spec apply_to_optimizer(scheduler(), viva_tensor@nn@optim:optimizer()) -> {scheduler(),
    viva_tensor@nn@optim:optimizer()}.
apply_to_optimizer(S, Opt) ->
    {Next, Lr} = scheduler_step(S),
    {Next,
        {optimizer,
            erlang:element(2, Opt),
            Lr,
            erlang:element(4, Opt),
            erlang:element(5, Opt),
            erlang:element(6, Opt),
            erlang:element(7, Opt),
            erlang:element(8, Opt),
            erlang:element(9, Opt)}}.