viva_tensor

-module(viva_tensor@quant@layout).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_tensor/quant/layout.gleam").
-export([nvfp4_block_scaled/1, int2_progressive/2, int3_progressive/2, element_count/1, memory_bits/1, memory_bytes/1, compression_ratio_against/2, format_name/1, is_rubin_native_candidate/1]).
-export_type([quant_format/0, scale_granularity/0, accumulator_format/0, quant_layout/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 quant_format() :: quant_int8 |
    quant_int4 |
    quant_nf4 |
    quant_nvfp4 |
    quant_int3 |
    quant_int2 |
    quant_fp2.

-type scale_granularity() :: per_tensor_scale |
    {per_block_scale, integer()} |
    {per_micro_block_scale, integer()} |
    {per_channel_scale, integer()}.

-type accumulator_format() :: accumulate_float32 |
    accumulate_float16 |
    accumulate_int32.

-type quant_layout() :: {quant_layout,
        list(integer()),
        quant_format(),
        scale_granularity(),
        accumulator_format(),
        integer(),
        boolean(),
        boolean(),
        integer(),
        boolean()}.

-file("src/viva_tensor/quant/layout.gleam", 48).
?DOC(false).
-spec nvfp4_block_scaled(list(integer())) -> quant_layout().
nvfp4_block_scaled(Shape) ->
    {quant_layout,
        Shape,
        quant_nvfp4,
        {per_micro_block_scale, 16},
        accumulate_float32,
        4,
        false,
        false,
        16,
        true}.

-file("src/viva_tensor/quant/layout.gleam", 62).
?DOC(false).
-spec int2_progressive(list(integer()), integer()) -> {ok, quant_layout()} |
    {error, viva_tensor@core@error:tensor_error()}.
int2_progressive(Shape, Block_size) ->
    case Block_size > 0 of
        false ->
            {error,
                {invalid_shape,
                    <<"INT2 progressive layout requires block_size > 0"/utf8>>}};

        true ->
            {ok,
                {quant_layout,
                    Shape,
                    quant_int2,
                    {per_block_scale, Block_size},
                    accumulate_float32,
                    2,
                    true,
                    true,
                    16,
                    true}}
    end.

-file("src/viva_tensor/quant/layout.gleam", 84).
?DOC(false).
-spec int3_progressive(list(integer()), integer()) -> {ok, quant_layout()} |
    {error, viva_tensor@core@error:tensor_error()}.
int3_progressive(Shape, Block_size) ->
    case Block_size > 0 of
        false ->
            {error,
                {invalid_shape,
                    <<"INT3 progressive layout requires block_size > 0"/utf8>>}};

        true ->
            {ok,
                {quant_layout,
                    Shape,
                    quant_int3,
                    {per_block_scale, Block_size},
                    accumulate_float32,
                    3,
                    true,
                    true,
                    16,
                    true}}
    end.

-file("src/viva_tensor/quant/layout.gleam", 106).
?DOC(false).
-spec element_count(list(integer())) -> integer().
element_count(Shape) ->
    gleam@list:fold(Shape, 1, fun(Total, Dim) -> Total * Dim end).

-file("src/viva_tensor/quant/layout.gleam", 110).
?DOC(false).
-spec memory_bits(quant_layout()) -> integer().
memory_bits(Layout) ->
    element_count(erlang:element(2, Layout)) * erlang:element(6, Layout).

-file("src/viva_tensor/quant/layout.gleam", 114).
?DOC(false).
-spec memory_bytes(quant_layout()) -> integer().
memory_bytes(Layout) ->
    Bits = memory_bits(Layout),
    case Bits rem 8 of
        0 ->
            Bits div 8;

        _ ->
            (Bits div 8) + 1
    end.

-file("src/viva_tensor/quant/layout.gleam", 122).
?DOC(false).
-spec compression_ratio_against(quant_layout(), integer()) -> float().
compression_ratio_against(Layout, Baseline_bits_per_value) ->
    case erlang:element(6, Layout) =< 0 of
        true ->
            +0.0;

        false ->
            case erlang:float(erlang:element(6, Layout)) of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator -> erlang:float(Baseline_bits_per_value) / Gleam@denominator
            end
    end.

-file("src/viva_tensor/quant/layout.gleam", 134).
?DOC(false).
-spec format_name(quant_format()) -> binary().
format_name(Format) ->
    case Format of
        quant_int8 ->
            <<"int8"/utf8>>;

        quant_int4 ->
            <<"int4"/utf8>>;

        quant_nf4 ->
            <<"nf4"/utf8>>;

        quant_nvfp4 ->
            <<"nvfp4"/utf8>>;

        quant_int3 ->
            <<"int3"/utf8>>;

        quant_int2 ->
            <<"int2"/utf8>>;

        quant_fp2 ->
            <<"fp2"/utf8>>
    end.

-file("src/viva_tensor/quant/layout.gleam", 146).
?DOC(false).
-spec is_rubin_native_candidate(quant_layout()) -> boolean().
is_rubin_native_candidate(Layout) ->
    case {erlang:element(3, Layout), erlang:element(9, Layout)} of
        {quant_nvfp4, 16} ->
            true;

        {quant_fp2, 16} ->
            true;

        {quant_int2, 16} ->
            erlang:element(8, Layout);

        {_, _} ->
            false
    end.