viva_tensor

-module(viva_tensor@core@layout_math).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_tensor/core/layout_math.gleam").
-export([size/1, indices/1, compute_strides/1, flat_to_multi/2, multi_to_flat/2, at/2, dim_at/2, value_at/2, replace_at/3, broadcast_strides/3]).

-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).

-file("src/viva_tensor/core/layout_math.gleam", 10).
?DOC(false).
-spec size(list(integer())) -> integer().
size(Shape) ->
    gleam@list:fold(Shape, 1, fun(Acc, Dim) -> Acc * Dim end).

-file("src/viva_tensor/core/layout_math.gleam", 120).
?DOC(false).
-spec range_loop(integer(), integer(), list(integer())) -> list(integer()).
range_loop(From, To, Acc) ->
    case From > To of
        true ->
            lists:reverse(Acc);

        false ->
            range_loop(From + 1, To, [From | Acc])
    end.

-file("src/viva_tensor/core/layout_math.gleam", 116).
?DOC(false).
-spec range_int(integer(), integer()) -> list(integer()).
range_int(From, To) ->
    range_loop(From, To, []).

-file("src/viva_tensor/core/layout_math.gleam", 15).
?DOC(false).
-spec indices(integer()) -> list(integer()).
indices(Size) ->
    case Size =< 0 of
        true ->
            [];

        false ->
            range_int(0, Size - 1)
    end.

-file("src/viva_tensor/core/layout_math.gleam", 23).
?DOC(false).
-spec compute_strides(list(integer())) -> list(integer()).
compute_strides(Shape) ->
    Reversed = lists:reverse(Shape),
    {Strides, _} = gleam@list:fold(
        Reversed,
        {[], 1},
        fun(Acc, Dim) ->
            {S, Running} = Acc,
            {[Running | S], Running * Dim}
        end
    ),
    Strides.

-file("src/viva_tensor/core/layout_math.gleam", 34).
?DOC(false).
-spec flat_to_multi(integer(), list(integer())) -> list(integer()).
flat_to_multi(Flat, Shape) ->
    Reversed = lists:reverse(Shape),
    {Coordinates, _} = gleam@list:fold(
        Reversed,
        {[], Flat},
        fun(Acc, Dim) ->
            {Coords, Remaining} = Acc,
            Coord = case Dim of
                0 -> 0;
                Gleam@denominator -> Remaining rem Gleam@denominator
            end,
            Next = case Dim of
                0 -> 0;
                Gleam@denominator@1 -> Remaining div Gleam@denominator@1
            end,
            {[Coord | Coords], Next}
        end
    ),
    Coordinates.

-file("src/viva_tensor/core/layout_math.gleam", 47).
?DOC(false).
-spec multi_to_flat(list(integer()), list(integer())) -> integer().
multi_to_flat(Coordinates, Shape) ->
    Strides = compute_strides(Shape),
    _pipe = gleam@list:zip(Coordinates, Strides),
    gleam@list:fold(
        _pipe,
        0,
        fun(Acc, Pair) ->
            {Coordinate, Stride} = Pair,
            Acc + (Coordinate * Stride)
        end
    ).

-file("src/viva_tensor/core/layout_math.gleam", 57).
?DOC(false).
-spec at(list(KZU), integer()) -> {ok, KZU} | {error, nil}.
at(Values, Index) ->
    case Index < 0 of
        true ->
            {error, nil};

        false ->
            _pipe = Values,
            _pipe@1 = gleam@list:drop(_pipe, Index),
            gleam@list:first(_pipe@1)
    end.

-file("src/viva_tensor/core/layout_math.gleam", 68).
?DOC(false).
-spec dim_at(list(integer()), integer()) -> integer().
dim_at(Values, Index) ->
    case at(Values, Index) of
        {ok, Value} ->
            Value;

        {error, _} ->
            0
    end.

-file("src/viva_tensor/core/layout_math.gleam", 76).
?DOC(false).
-spec value_at(list(float()), integer()) -> float().
value_at(Values, Index) ->
    case at(Values, Index) of
        {ok, Value} ->
            Value;

        {error, _} ->
            +0.0
    end.

-file("src/viva_tensor/core/layout_math.gleam", 84).
?DOC(false).
-spec replace_at(list(integer()), integer(), integer()) -> list(integer()).
replace_at(Values, Index, Value) ->
    _pipe = Values,
    gleam@list:index_map(_pipe, fun(Item, I) -> case I =:= Index of
                true ->
                    Value;

                false ->
                    Item
            end end).

-file("src/viva_tensor/core/layout_math.gleam", 97).
?DOC(false).
-spec broadcast_strides(list(integer()), list(integer()), list(integer())) -> list(integer()).
broadcast_strides(Src_shape, Src_strides, Target_shape) ->
    Diff = erlang:length(Target_shape) - erlang:length(Src_shape),
    Padded_shape = lists:append(gleam@list:repeat(1, Diff), Src_shape),
    Padded_strides = lists:append(gleam@list:repeat(0, Diff), Src_strides),
    _pipe = gleam@list:zip(
        gleam@list:zip(Padded_shape, Target_shape),
        Padded_strides
    ),
    gleam@list:map(
        _pipe,
        fun(Item) ->
            {{Src_dim, Target_dim}, Stride} = Item,
            case Src_dim =:= Target_dim of
                true ->
                    Stride;

                false ->
                    0
            end
        end
    ).