viva_tensor

-module(viva_tensor@nn@layers).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_tensor/nn/layers.gleam").
-export([linear/3, linear_forward/3, relu/2, mse_loss/3]).
-export_type([linear/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(
    " Neural Network Layers - The Building Blocks of Deep Learning\n"
    "\n"
    " \"A neural network is just a differentiable program.\"\n"
    "   — Yann LeCun (paraphrased, but he'd probably agree)\n"
    "\n"
    " References:\n"
    " - Rumelhart, Hinton & Williams (1986). \"Learning representations by\n"
    "   back-propagating errors.\" Nature. THE paper that started it all.\n"
    " - Glorot & Bengio (2010). \"Understanding the difficulty of training\n"
    "   deep feedforward neural networks.\" Xavier initialization lives here.\n"
    " - He et al. (2015). \"Delving Deep into Rectifiers.\" Kaiming init for ReLU.\n"
    "\n"
    " Design philosophy:\n"
    " - Layers are values, not classes. No hidden state, no surprises.\n"
    " - Forward pass returns Traced(Variable) - computation graph included.\n"
    " - Everything flows through the tape. Backprop just works.\n"
).

-type linear() :: {linear,
        viva_tensor@nn@autograd:variable(),
        viva_tensor@nn@autograd:variable()}.

-file("src/viva_tensor/nn/layers.gleam", 51).
?DOC(
    " Creates a new Linear layer with Xavier/Glorot initialization.\n"
    "\n"
    " Xavier init: W ~ Uniform(-sqrt(6/(fan_in+fan_out)), sqrt(6/(fan_in+fan_out)))\n"
    " This keeps variance stable across layers for tanh/sigmoid.\n"
    " For ReLU, you'd want He init (scale by sqrt(2/fan_in)) instead.\n"
    "\n"
    " The bias starts at zero. Some argue for small positive values\n"
    " to ensure ReLU neurons fire initially, but zero works fine.\n"
).
-spec linear(viva_tensor@nn@autograd:tape(), integer(), integer()) -> viva_tensor@nn@autograd:traced(linear()).
linear(Tape, In_features, Out_features) ->
    W_data = viva_tensor@core@tensor:xavier_init(In_features, Out_features),
    B_data = viva_tensor@core@tensor:zeros([Out_features]),
    {traced, W, Tape1} = viva_tensor@nn@autograd:new_variable(Tape, W_data),
    {traced, B, Tape2} = viva_tensor@nn@autograd:new_variable(Tape1, B_data),
    {traced, {linear, W, B}, Tape2}.

-file("src/viva_tensor/nn/layers.gleam", 71).
?DOC(
    " Forward pass: y = xW^T + b\n"
    " Instrumented: records forward pass latency.\n"
    "\n"
    " The order of operations matters for gradient computation:\n"
    " 1. Transpose W: [out, in] -> [in, out]\n"
    " 2. Matmul: [batch, in] @ [in, out] -> [batch, out]\n"
    " 3. Add bias: [batch, out] + [out] (broadcast over batch dim)\n"
    "\n"
    " Each step is traced, so backward() will compute dL/dW, dL/db, dL/dx.\n"
).
-spec linear_forward(
    viva_tensor@nn@autograd:tape(),
    linear(),
    viva_tensor@nn@autograd:variable()
) -> {ok, viva_tensor@nn@autograd:traced(viva_tensor@nn@autograd:variable())} |
    {error, viva_tensor@core@error:tensor_error()}.
linear_forward(Tape, Layer, X) ->
    T0 = viva_tensor@core@ffi:now_microseconds(),
    gleam@result:'try'(
        viva_tensor@nn@autograd:transpose(Tape, erlang:element(2, Layer)),
        fun(_use0) ->
            {traced, Wt, Tape1} = _use0,
            gleam@result:'try'(
                viva_tensor@nn@autograd:matmul(Tape1, X, Wt),
                fun(_use0@1) ->
                    {traced, Xw, Tape2} = _use0@1,
                    Result = viva_tensor@nn@autograd:add(
                        Tape2,
                        Xw,
                        erlang:element(3, Layer)
                    ),
                    viva_tensor@telemetry:record_op(
                        <<"linear_forward"/utf8>>,
                        viva_tensor@core@ffi:now_microseconds() - T0
                    ),
                    Result
                end
            )
        end
    ).

-file("src/viva_tensor/nn/layers.gleam", 117).
?DOC(
    " ReLU activation: f(x) = max(0, x)\n"
    "\n"
    " The most popular activation function. Simple, effective, sometimes dead.\n"
    "\n"
    " Why ReLU wins:\n"
    " - Sparse activations (many zeros = efficient)\n"
    " - No vanishing gradient for positive values (gradient = 1)\n"
    " - Computationally trivial (just a comparison)\n"
    "\n"
    " Why ReLU loses:\n"
    " - \"Dying ReLU\": neurons that output 0 have 0 gradient forever\n"
    " - Unbounded output can cause numerical issues\n"
    "\n"
    " Alternatives: LeakyReLU, GELU, SiLU/Swish. Each has tradeoffs.\n"
).
-spec relu(viva_tensor@nn@autograd:tape(), viva_tensor@nn@autograd:variable()) -> viva_tensor@nn@autograd:traced(viva_tensor@nn@autograd:variable()).
relu(Tape, X) ->
    viva_tensor@nn@autograd:relu(Tape, X).

-file("src/viva_tensor/nn/layers.gleam", 150).
?DOC(
    " Mean Squared Error loss: L = mean((pred - target)^2)\n"
    "\n"
    " MSE: the L2 norm's favorite child.\n"
    "\n"
    " Properties:\n"
    " - Gradient: dL/dpred = 2(pred - target) / n\n"
    " - Strongly convex (unique minimum)\n"
    " - Penalizes large errors quadratically (sensitive to outliers)\n"
    " - The MLE under Gaussian noise assumption\n"
    "\n"
    " When to use:\n"
    " - Regression problems with Gaussian-distributed errors\n"
    " - When you want smooth, well-behaved gradients\n"
    "\n"
    " When NOT to use:\n"
    " - Outlier-heavy data (use Huber or MAE instead)\n"
    " - Classification (use cross-entropy)\n"
).
-spec mse_loss(
    viva_tensor@nn@autograd:tape(),
    viva_tensor@nn@autograd:variable(),
    viva_tensor@nn@autograd:variable()
) -> {ok, viva_tensor@nn@autograd:traced(viva_tensor@nn@autograd:variable())} |
    {error, viva_tensor@core@error:tensor_error()}.
mse_loss(Tape, Pred, Target) ->
    gleam@result:'try'(
        viva_tensor@nn@autograd:sub(Tape, Pred, Target),
        fun(_use0) ->
            {traced, Diff, Tape1} = _use0,
            gleam@result:'try'(
                viva_tensor@nn@autograd:mul(Tape1, Diff, Diff),
                fun(_use0@1) ->
                    {traced, Square, Tape2} = _use0@1,
                    {ok, viva_tensor@nn@autograd:mean(Tape2, Square)}
                end
            )
        end
    ).