viva_tensor

-module(viva_tensor@quant@turboquant).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_tensor/quant/turboquant.gleam").
-export([default_config/0, quantize/2, quantize_tensor/2, dequantize/1, dequantize_tensor/1, inner_product/2, compression_ratio/1]).
-export_type([config/0, quantized_vector/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 config() :: {config, integer(), integer(), boolean()}.

-type quantized_vector() :: {quantized_vector,
        list(integer()),
        list(integer()),
        float(),
        float(),
        integer(),
        integer(),
        integer(),
        integer(),
        boolean(),
        integer()}.

-file("src/viva_tensor/quant/turboquant.gleam", 54).
?DOC(false).
-spec default_config() -> config().
default_config() ->
    {config, 3, 0, true}.

-file("src/viva_tensor/quant/turboquant.gleam", 318).
?DOC(false).
-spec estimate_memory_bytes(integer(), integer(), boolean()) -> integer().
estimate_memory_bytes(Dim, Bits, Use_residual) ->
    Residual_bits = case Use_residual of
        true ->
            Dim;

        false ->
            0
    end,
    Payload_bits = (Dim * Bits) + Residual_bits,
    ((Payload_bits + 7) div 8) + 8.

-file("src/viva_tensor/quant/turboquant.gleam", 272).
?DOC(false).
-spec mean_abs(list(float())) -> float().
mean_abs(Values) ->
    case erlang:length(Values) of
        0 ->
            +0.0;

        N ->
            _pipe = Values,
            _pipe@1 = gleam@list:map(_pipe, fun gleam@float:absolute_value/1),
            _pipe@2 = gleam@list:fold(
                _pipe@1,
                +0.0,
                fun(Acc, Value) -> Acc + Value end
            ),
            (fun(Total) -> case erlang:float(N) of
                    +0.0 -> +0.0;
                    -0.0 -> -0.0;
                    Gleam@denominator -> Total / Gleam@denominator
                end end)(_pipe@2)
    end.

-file("src/viva_tensor/quant/turboquant.gleam", 259).
?DOC(false).
-spec sign_bit(float()) -> integer().
sign_bit(Value) ->
    case Value < +0.0 of
        true ->
            -1;

        false ->
            1
    end.

-file("src/viva_tensor/quant/turboquant.gleam", 307).
?DOC(false).
-spec int_pow2(integer()) -> integer().
int_pow2(Exp) ->
    case Exp =< 0 of
        true ->
            1;

        false ->
            2 * int_pow2(Exp - 1)
    end.

-file("src/viva_tensor/quant/turboquant.gleam", 237).
?DOC(false).
-spec decode_code(integer(), integer(), float()) -> float().
decode_code(Code, Bits, Scale) ->
    case Scale =:= +0.0 of
        true ->
            +0.0;

        false ->
            Levels = int_pow2(Bits),
            (-1.0 * Scale) + (case erlang:float(Levels - 1) of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator -> (2.0 * Scale) * erlang:float(Code) / Gleam@denominator
            end)
    end.

-file("src/viva_tensor/quant/turboquant.gleam", 223).
?DOC(false).
-spec nearest_code(float(), integer(), float()) -> integer().
nearest_code(Value, Bits, Scale) ->
    Levels = int_pow2(Bits),
    _pipe = gleam@list:range(0, Levels - 1),
    _pipe@1 = gleam@list:fold(
        _pipe,
        {0, 1.0e308},
        fun(Best, Code) ->
            Decoded = decode_code(Code, Bits, Scale),
            Distance = gleam@float:absolute_value(Value - Decoded),
            case Distance < erlang:element(2, Best) of
                true ->
                    {Code, Distance};

                false ->
                    Best
            end
        end
    ),
    (fun(Best@1) -> erlang:element(1, Best@1) end)(_pipe@1).

-file("src/viva_tensor/quant/turboquant.gleam", 266).
?DOC(false).
-spec max_abs(list(float())) -> float().
max_abs(Values) ->
    _pipe = Values,
    _pipe@1 = gleam@list:map(_pipe, fun gleam@float:absolute_value/1),
    gleam@list:fold(_pipe@1, +0.0, fun gleam@float:max/2).

-file("src/viva_tensor/quant/turboquant.gleam", 196).
?DOC(false).
-spec hadamard(list(float())) -> list(float()).
hadamard(Values) ->
    case Values of
        [] ->
            [];

        [_] ->
            Values;

        _ ->
            Half = erlang:length(Values) div 2,
            Left = begin
                _pipe = Values,
                _pipe@1 = gleam@list:take(_pipe, Half),
                hadamard(_pipe@1)
            end,
            Right = begin
                _pipe@2 = Values,
                _pipe@3 = gleam@list:drop(_pipe@2, Half),
                hadamard(_pipe@3)
            end,
            lists:append(
                gleam@list:map2(Left, Right, fun(A, B) -> A + B end),
                gleam@list:map2(Left, Right, fun(A@1, B@1) -> A@1 - B@1 end)
            )
    end.

-file("src/viva_tensor/quant/turboquant.gleam", 212).
?DOC(false).
-spec normalize_hadamard(list(float())) -> list(float()).
normalize_hadamard(Values) ->
    N = erlang:length(Values),
    case N of
        0 ->
            [];

        _ ->
            Scale = case math:sqrt(erlang:float(N)) of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator -> 1.0 / Gleam@denominator
            end,
            gleam@list:map(Values, fun(Value) -> Value * Scale end)
    end.

-file("src/viva_tensor/quant/turboquant.gleam", 252).
?DOC(false).
-spec random_sign(integer(), integer()) -> integer().
random_sign(Seed, Index) ->
    case erlang:phash2({Seed, Index}, 2) of
        0 ->
            -1;

        _ ->
            1
    end.

-file("src/viva_tensor/quant/turboquant.gleam", 178).
?DOC(false).
-spec randomized_hadamard(list(float()), integer()) -> list(float()).
randomized_hadamard(Values, Seed) ->
    Signed = begin
        _pipe = Values,
        gleam@list:index_map(
            _pipe,
            fun(Value, Index) ->
                Value * erlang:float(random_sign(Seed, Index))
            end
        )
    end,
    normalize_hadamard(hadamard(Signed)).

-file("src/viva_tensor/quant/turboquant.gleam", 288).
?DOC(false).
-spec pad_to(list(float()), integer()) -> list(float()).
pad_to(Values, Size) ->
    Missing = Size - erlang:length(Values),
    case Missing =< 0 of
        true ->
            Values;

        false ->
            lists:append(Values, gleam@list:repeat(+0.0, Missing))
    end.

-file("src/viva_tensor/quant/turboquant.gleam", 300).
?DOC(false).
-spec next_power_of_two_loop(integer(), integer()) -> integer().
next_power_of_two_loop(Current, N) ->
    case Current >= N of
        true ->
            Current;

        false ->
            next_power_of_two_loop(Current * 2, N)
    end.

-file("src/viva_tensor/quant/turboquant.gleam", 296).
?DOC(false).
-spec next_power_of_two(integer()) -> integer().
next_power_of_two(N) ->
    next_power_of_two_loop(1, N).

-file("src/viva_tensor/quant/turboquant.gleam", 314).
?DOC(false).
-spec valid_bits(integer()) -> boolean().
valid_bits(Bits) ->
    (Bits >= 1) andalso (Bits =< 8).

-file("src/viva_tensor/quant/turboquant.gleam", 59).
?DOC(false).
-spec quantize(list(float()), config()) -> {ok, quantized_vector()} |
    {error, viva_tensor@core@error:tensor_error()}.
quantize(Values, Config) ->
    case {erlang:length(Values), valid_bits(erlang:element(2, Config))} of
        {0, _} ->
            {error,
                {invalid_shape,
                    <<"TurboQuant requires a non-empty vector"/utf8>>}};

        {_, false} ->
            {error,
                {invalid_shape,
                    <<"TurboQuant bits must be between 1 and 8"/utf8>>}};

        {Dim, true} ->
            Padded_dim = next_power_of_two(Dim),
            Padded = pad_to(Values, Padded_dim),
            Rotated = randomized_hadamard(Padded, erlang:element(3, Config)),
            Scale = max_abs(Rotated),
            Codes = gleam@list:map(
                Rotated,
                fun(Value) ->
                    nearest_code(Value, erlang:element(2, Config), Scale)
                end
            ),
            Main = gleam@list:map(
                Codes,
                fun(Code) ->
                    decode_code(Code, erlang:element(2, Config), Scale)
                end
            ),
            Residual = gleam@list:map2(
                Rotated,
                Main,
                fun(Value@1, Approx) -> Value@1 - Approx end
            ),
            Residual_signs = case erlang:element(4, Config) of
                true ->
                    gleam@list:map(Residual, fun sign_bit/1);

                false ->
                    []
            end,
            Residual_scale = case erlang:element(4, Config) of
                true ->
                    mean_abs(Residual);

                false ->
                    +0.0
            end,
            {ok,
                {quantized_vector,
                    Codes,
                    Residual_signs,
                    Scale,
                    Residual_scale,
                    Dim,
                    Padded_dim,
                    erlang:element(2, Config),
                    erlang:element(3, Config),
                    erlang:element(4, Config),
                    estimate_memory_bytes(
                        Padded_dim,
                        erlang:element(2, Config),
                        erlang:element(4, Config)
                    )}}
    end.

-file("src/viva_tensor/quant/turboquant.gleam", 329).
?DOC(false).
-spec shape_to_string(list(integer())) -> binary().
shape_to_string(Shape) ->
    Body = begin
        _pipe = Shape,
        _pipe@1 = gleam@list:map(_pipe, fun erlang:integer_to_binary/1),
        gleam@list:fold(
            _pipe@1,
            <<""/utf8>>,
            fun(Acc, Part) -> case Acc =:= <<""/utf8>> of
                    true ->
                        Part;

                    false ->
                        <<<<Acc/binary, ", "/utf8>>/binary, Part/binary>>
                end end
        )
    end,
    <<<<"["/utf8, Body/binary>>/binary, "]"/utf8>>.

-file("src/viva_tensor/quant/turboquant.gleam", 109).
?DOC(false).
-spec quantize_tensor(viva_tensor@tensor:tensor(), config()) -> {ok,
        quantized_vector()} |
    {error, viva_tensor@core@error:tensor_error()}.
quantize_tensor(Tensor, Config) ->
    case erlang:element(3, Tensor) of
        [_] ->
            quantize(viva_tensor@tensor:to_list(Tensor), Config);

        Shape ->
            {error,
                {invalid_shape,
                    <<"TurboQuant currently expects a vector tensor, got "/utf8,
                        (shape_to_string(Shape))/binary>>}}
    end.

-file("src/viva_tensor/quant/turboquant.gleam", 188).
?DOC(false).
-spec inverse_randomized_hadamard(list(float()), integer()) -> list(float()).
inverse_randomized_hadamard(Values, Seed) ->
    _pipe = hadamard(Values),
    _pipe@1 = normalize_hadamard(_pipe),
    gleam@list:index_map(
        _pipe@1,
        fun(Value, Index) -> Value * erlang:float(random_sign(Seed, Index)) end
    ).

-file("src/viva_tensor/quant/turboquant.gleam", 163).
?DOC(false).
-spec dequantize_rotated(quantized_vector()) -> list(float()).
dequantize_rotated(Vector) ->
    Main = gleam@list:map(
        erlang:element(2, Vector),
        fun(Code) ->
            decode_code(
                Code,
                erlang:element(8, Vector),
                erlang:element(4, Vector)
            )
        end
    ),
    case erlang:element(10, Vector) of
        false ->
            Main;

        true ->
            gleam@list:map2(
                Main,
                erlang:element(3, Vector),
                fun(Value, Sign) ->
                    Value + (erlang:float(Sign) * erlang:element(5, Vector))
                end
            )
    end.

-file("src/viva_tensor/quant/turboquant.gleam", 124).
?DOC(false).
-spec dequantize(quantized_vector()) -> list(float()).
dequantize(Vector) ->
    _pipe = dequantize_rotated(Vector),
    _pipe@1 = inverse_randomized_hadamard(_pipe, erlang:element(9, Vector)),
    gleam@list:take(_pipe@1, erlang:element(6, Vector)).

-file("src/viva_tensor/quant/turboquant.gleam", 131).
?DOC(false).
-spec dequantize_tensor(quantized_vector()) -> viva_tensor@tensor:tensor().
dequantize_tensor(Vector) ->
    viva_tensor@tensor:from_list(dequantize(Vector)).

-file("src/viva_tensor/quant/turboquant.gleam", 283).
?DOC(false).
-spec dot(list(float()), list(float())) -> float().
dot(A, B) ->
    _pipe = gleam@list:map2(A, B, fun(X, Y) -> X * Y end),
    gleam@list:fold(_pipe, +0.0, fun(Acc, Value) -> Acc + Value end).

-file("src/viva_tensor/quant/turboquant.gleam", 137).
?DOC(false).
-spec inner_product(list(float()), quantized_vector()) -> {ok, float()} |
    {error, viva_tensor@core@error:tensor_error()}.
inner_product(Query, Vector) ->
    case erlang:length(Query) =:= erlang:element(6, Vector) of
        false ->
            {error,
                {shape_mismatch,
                    [erlang:element(6, Vector)],
                    [erlang:length(Query)]}};

        true ->
            Rotated_query = begin
                _pipe = Query,
                _pipe@1 = pad_to(_pipe, erlang:element(7, Vector)),
                randomized_hadamard(_pipe@1, erlang:element(9, Vector))
            end,
            {ok, dot(Rotated_query, dequantize_rotated(Vector))}
    end.

-file("src/viva_tensor/quant/turboquant.gleam", 158).
?DOC(false).
-spec compression_ratio(quantized_vector()) -> float().
compression_ratio(Vector) ->
    Original_bytes = erlang:element(6, Vector) * 4,
    case erlang:float(erlang:element(11, Vector)) of
        +0.0 -> +0.0;
        -0.0 -> -0.0;
        Gleam@denominator -> erlang:float(Original_bytes) / Gleam@denominator
    end.