viva_tensor

-module(viva_tensor@core@einsum).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_tensor/core/einsum.gleam").
-export([run/2, parse/1, validate/3, general/4]).
-export_type([operand/0, parsed/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 operand() :: {operand, list(float()), list(integer())}.

-type parsed() :: {parsed, list(list(binary())), list(binary())}.

-file("src/viva_tensor/core/einsum.gleam", 416).
?DOC(false).
-spec nth_loop(list(float()), integer()) -> {ok, float()} | {error, nil}.
nth_loop(Items, Index) ->
    case {Items, Index} of
        {[], _} ->
            {error, nil};

        {[X | _], 0} ->
            {ok, X};

        {[_ | Rest], N} ->
            nth_loop(Rest, N - 1)
    end.

-file("src/viva_tensor/core/einsum.gleam", 409).
?DOC(false).
-spec nth(list(float()), integer()) -> {ok, float()} | {error, nil}.
nth(Items, Index) ->
    case Index < 0 of
        true ->
            {error, nil};

        false ->
            nth_loop(Items, Index)
    end.

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

-file("src/viva_tensor/core/einsum.gleam", 392).
?DOC(false).
-spec flat_index(list(integer()), list(integer())) -> integer().
flat_index(Indices, Shape) ->
    Strides = compute_strides(Shape),
    _pipe = gleam@list:zip(Indices, Strides),
    gleam@list:fold(
        _pipe,
        0,
        fun(Acc, Pair) ->
            Acc + (erlang:element(1, Pair) * erlang:element(2, Pair))
        end
    ).

-file("src/viva_tensor/core/einsum.gleam", 378).
?DOC(false).
-spec lookup_binding(list({binary(), integer()}), binary()) -> {ok, integer()} |
    {error, nil}.
lookup_binding(Bindings, Key) ->
    case Bindings of
        [] ->
            {error, nil};

        [{K, V} | Rest] ->
            case K =:= Key of
                true ->
                    {ok, V};

                false ->
                    lookup_binding(Rest, Key)
            end
    end.

-file("src/viva_tensor/core/einsum.gleam", 366).
?DOC(false).
-spec label_indices(list(binary()), list({binary(), integer()})) -> list(integer()).
label_indices(Labels, Bindings) ->
    gleam@list:map(Labels, fun(Label) -> case lookup_binding(Bindings, Label) of
                {ok, V} ->
                    V;

                {error, _} ->
                    0
            end end).

-file("src/viva_tensor/core/einsum.gleam", 428).
?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/einsum.gleam", 424).
?DOC(false).
-spec range_int(integer(), integer()) -> list(integer()).
range_int(From, To) ->
    range_loop(From, To, []).

-file("src/viva_tensor/core/einsum.gleam", 356).
?DOC(false).
-spec cartesian(list(list(integer()))) -> list(list(integer())).
cartesian(Ranges) ->
    case Ranges of
        [] ->
            [[]];

        [First | Rest] ->
            Tails = cartesian(Rest),
            gleam@list:flat_map(
                First,
                fun(V) -> gleam@list:map(Tails, fun(T) -> [V | T] end) end
            )
    end.

-file("src/viva_tensor/core/einsum.gleam", 231).
?DOC(false).
-spec lookup(list({binary(), integer()}), binary()) -> {ok, integer()} |
    {error, nil}.
lookup(Pairs, Key) ->
    case Pairs of
        [] ->
            {error, nil};

        [{K, V} | Rest] ->
            case K =:= Key of
                true ->
                    {ok, V};

                false ->
                    lookup(Rest, Key)
            end
    end.

-file("src/viva_tensor/core/einsum.gleam", 345).
?DOC(false).
-spec dedupe(list(binary()), list(binary())) -> list(binary()).
dedupe(Items, Seen) ->
    case Items of
        [] ->
            lists:reverse(Seen);

        [X | Rest] ->
            case gleam@list:contains(Seen, X) of
                true ->
                    dedupe(Rest, Seen);

                false ->
                    dedupe(Rest, [X | Seen])
            end
    end.

-file("src/viva_tensor/core/einsum.gleam", 285).
?DOC(false).
-spec general_kernel(
    list(operand()),
    list(list(binary())),
    list(binary()),
    list({binary(), integer()})
) -> {ok, operand()} | {error, viva_tensor@core@error:tensor_error()}.
general_kernel(Operands, Lhs, Rhs, Dim_map) ->
    All_labels = dedupe(lists:append(Lhs), []),
    Summed = gleam@list:filter(
        All_labels,
        fun(Label) -> not gleam@list:contains(Rhs, Label) end
    ),
    Out_shape = gleam@list:map(
        Rhs,
        fun(Label@1) -> case lookup(Dim_map, Label@1) of
                {ok, D} ->
                    D;

                {error, _} ->
                    0
            end end
    ),
    Out_combos = cartesian(
        gleam@list:map(Out_shape, fun(D@1) -> range_int(0, D@1 - 1) end)
    ),
    Sum_dims = gleam@list:map(
        Summed,
        fun(Label@2) -> case lookup(Dim_map, Label@2) of
                {ok, D@2} ->
                    D@2;

                {error, _} ->
                    0
            end end
    ),
    Sum_combos = cartesian(
        gleam@list:map(Sum_dims, fun(D@3) -> range_int(0, D@3 - 1) end)
    ),
    Lhs_data = gleam@list:zip(Lhs, Operands),
    Data = gleam@list:map(
        Out_combos,
        fun(Out_idx) ->
            Out_bindings = gleam@list:zip(Rhs, Out_idx),
            gleam@list:fold(
                Sum_combos,
                +0.0,
                fun(Acc, Sum_idx) ->
                    Sum_bindings = gleam@list:zip(Summed, Sum_idx),
                    Bindings = lists:append(Out_bindings, Sum_bindings),
                    Product = gleam@list:fold(
                        Lhs_data,
                        1.0,
                        fun(P, Pair) ->
                            {Labels, Op} = Pair,
                            Idx = label_indices(Labels, Bindings),
                            Flat = flat_index(Idx, erlang:element(3, Op)),
                            case nth(erlang:element(2, Op), Flat) of
                                {ok, V} ->
                                    P * V;

                                {error, _} ->
                                    P
                            end
                        end
                    ),
                    Acc + Product
                end
            )
        end
    ),
    case Out_shape of
        [] ->
            case Data of
                [V@1] ->
                    {ok, {operand, [V@1], []}};

                _ ->
                    {ok, {operand, [+0.0], []}}
            end;

        _ ->
            {ok, {operand, Data, Out_shape}}
    end.

-file("src/viva_tensor/core/einsum.gleam", 269).
?DOC(false).
-spec has_unique(list(binary()), list(binary())) -> boolean().
has_unique(Items, Seen) ->
    case Items of
        [] ->
            true;

        [X | Rest] ->
            case gleam@list:contains(Seen, X) of
                true ->
                    false;

                false ->
                    has_unique(Rest, [X | Seen])
            end
    end.

-file("src/viva_tensor/core/einsum.gleam", 242).
?DOC(false).
-spec check_rhs_labels(list(binary()), list(list(binary()))) -> {ok, nil} |
    {error, viva_tensor@core@error:tensor_error()}.
check_rhs_labels(Rhs, Lhs) ->
    All_lhs = lists:append(Lhs),
    gleam@result:'try'(
        gleam@list:try_each(
            Rhs,
            fun(Label) -> case gleam@list:contains(All_lhs, Label) of
                    true ->
                        {ok, nil};

                    false ->
                        {error,
                            {invalid_shape,
                                <<<<"einsum: output label '"/utf8,
                                        Label/binary>>/binary,
                                    "' does not appear in any input"/utf8>>}}
                end end
        ),
        fun(_) -> case has_unique(Rhs, []) of
                true ->
                    {ok, nil};

                false ->
                    {error,
                        {invalid_shape,
                            <<"einsum: output labels must be unique (repeated label on RHS)"/utf8>>}}
            end end
    ).

-file("src/viva_tensor/core/einsum.gleam", 214).
?DOC(false).
-spec merge_dims(list({binary(), integer()}), list({binary(), integer()})) -> {ok,
        list({binary(), integer()})} |
    {error, viva_tensor@core@error:tensor_error()}.
merge_dims(Acc, Pairs) ->
    gleam@list:try_fold(
        Pairs,
        Acc,
        fun(Acc@1, Pair) ->
            {Label, Dim} = Pair,
            case lookup(Acc@1, Label) of
                {ok, Existing} ->
                    case Existing =:= Dim of
                        true ->
                            {ok, Acc@1};

                        false ->
                            {error, {shape_mismatch, [Existing], [Dim]}}
                    end;

                {error, _} ->
                    {ok, [{Label, Dim} | Acc@1]}
            end
        end
    ).

-file("src/viva_tensor/core/einsum.gleam", 203).
?DOC(false).
-spec infer_dims(list(list(binary())), list(operand())) -> {ok,
        list({binary(), integer()})} |
    {error, viva_tensor@core@error:tensor_error()}.
infer_dims(Lhs, Operands) ->
    _pipe = gleam@list:zip(Lhs, Operands),
    gleam@list:try_fold(
        _pipe,
        [],
        fun(Acc, Pair) ->
            {Labels, Op} = Pair,
            merge_dims(Acc, gleam@list:zip(Labels, erlang:element(3, Op)))
        end
    ).

-file("src/viva_tensor/core/einsum.gleam", 182).
?DOC(false).
-spec check_rank_match(list(list(binary())), list(operand())) -> {ok, nil} |
    {error, viva_tensor@core@error:tensor_error()}.
check_rank_match(Lhs, Operands) ->
    _pipe = gleam@list:zip(Lhs, Operands),
    gleam@list:try_each(
        _pipe,
        fun(Pair) ->
            {Labels, Op} = Pair,
            R = erlang:length(erlang:element(3, Op)),
            case erlang:length(Labels) =:= R of
                true ->
                    {ok, nil};

                false ->
                    {error,
                        {dimension_error,
                            <<<<<<"einsum: operand rank "/utf8,
                                        (erlang:integer_to_binary(R))/binary>>/binary,
                                    " does not match label count "/utf8>>/binary,
                                (erlang:integer_to_binary(erlang:length(Labels)))/binary>>}}
            end
        end
    ).

-file("src/viva_tensor/core/einsum.gleam", 164).
?DOC(false).
-spec check_operand_arity(list(list(binary())), list(operand())) -> {ok, nil} |
    {error, viva_tensor@core@error:tensor_error()}.
check_operand_arity(Lhs, Operands) ->
    Expected = erlang:length(Lhs),
    Got = erlang:length(Operands),
    case Expected =:= Got of
        true ->
            {ok, nil};

        false ->
            {error,
                {dimension_error,
                    <<<<<<"einsum: equation expects "/utf8,
                                (erlang:integer_to_binary(Expected))/binary>>/binary,
                            " operand(s) but got "/utf8>>/binary,
                        (erlang:integer_to_binary(Got))/binary>>}}
    end.

-file("src/viva_tensor/core/einsum.gleam", 150).
?DOC(false).
-spec is_label_char(binary()) -> boolean().
is_label_char(G) ->
    case G of
        <<"a"/utf8>> ->
            true;

        <<"b"/utf8>> ->
            true;

        <<"c"/utf8>> ->
            true;

        <<"d"/utf8>> ->
            true;

        <<"e"/utf8>> ->
            true;

        <<"f"/utf8>> ->
            true;

        <<"g"/utf8>> ->
            true;

        <<"h"/utf8>> ->
            true;

        <<"i"/utf8>> ->
            true;

        <<"j"/utf8>> ->
            true;

        <<"k"/utf8>> ->
            true;

        <<"l"/utf8>> ->
            true;

        <<"m"/utf8>> ->
            true;

        <<"n"/utf8>> ->
            true;

        <<"o"/utf8>> ->
            true;

        <<"p"/utf8>> ->
            true;

        <<"q"/utf8>> ->
            true;

        <<"r"/utf8>> ->
            true;

        <<"s"/utf8>> ->
            true;

        <<"t"/utf8>> ->
            true;

        <<"u"/utf8>> ->
            true;

        <<"v"/utf8>> ->
            true;

        <<"w"/utf8>> ->
            true;

        <<"x"/utf8>> ->
            true;

        <<"y"/utf8>> ->
            true;

        <<"z"/utf8>> ->
            true;

        <<"A"/utf8>> ->
            true;

        <<"B"/utf8>> ->
            true;

        <<"C"/utf8>> ->
            true;

        <<"D"/utf8>> ->
            true;

        <<"E"/utf8>> ->
            true;

        <<"F"/utf8>> ->
            true;

        <<"G"/utf8>> ->
            true;

        <<"H"/utf8>> ->
            true;

        <<"I"/utf8>> ->
            true;

        <<"J"/utf8>> ->
            true;

        <<"K"/utf8>> ->
            true;

        <<"L"/utf8>> ->
            true;

        <<"M"/utf8>> ->
            true;

        <<"N"/utf8>> ->
            true;

        <<"O"/utf8>> ->
            true;

        <<"P"/utf8>> ->
            true;

        <<"Q"/utf8>> ->
            true;

        <<"R"/utf8>> ->
            true;

        <<"S"/utf8>> ->
            true;

        <<"T"/utf8>> ->
            true;

        <<"U"/utf8>> ->
            true;

        <<"V"/utf8>> ->
            true;

        <<"W"/utf8>> ->
            true;

        <<"X"/utf8>> ->
            true;

        <<"Y"/utf8>> ->
            true;

        <<"Z"/utf8>> ->
            true;

        _ ->
            false
    end.

-file("src/viva_tensor/core/einsum.gleam", 137).
?DOC(false).
-spec check_labels(list(binary()), binary()) -> {ok, nil} |
    {error, viva_tensor@core@error:tensor_error()}.
check_labels(Labels, Context) ->
    case gleam@list:all(Labels, fun is_label_char/1) of
        true ->
            {ok, nil};

        false ->
            {error,
                {invalid_shape,
                    <<<<"einsum: invalid label in "/utf8, Context/binary>>/binary,
                        " (expected ASCII letters)"/utf8>>}}
    end.

-file("src/viva_tensor/core/einsum.gleam", 130).
?DOC(false).
-spec remove_whitespace(binary()) -> binary().
remove_whitespace(S) ->
    _pipe = S,
    _pipe@1 = gleam@string:to_graphemes(_pipe),
    _pipe@2 = gleam@list:filter(
        _pipe@1,
        fun(G) ->
            ((G /= <<" "/utf8>>) andalso (G /= <<"\t"/utf8>>)) andalso (G /= <<"\n"/utf8>>)
        end
    ),
    erlang:list_to_binary(_pipe@2).

-file("src/viva_tensor/core/einsum.gleam", 99).
?DOC(false).
-spec parse_equation(binary()) -> {ok, parsed()} |
    {error, viva_tensor@core@error:tensor_error()}.
parse_equation(Equation) ->
    Normalized = remove_whitespace(Equation),
    case gleam_stdlib:contains_string(Normalized, <<"..."/utf8>>) of
        true ->
            {error,
                {dimension_error,
                    <<"einsum: feature not supported in v1: ellipsis"/utf8>>}};

        false ->
            case gleam@string:split(Normalized, <<"->"/utf8>>) of
                [Lhs, Rhs] ->
                    Lhs_parts = gleam@string:split(Lhs, <<","/utf8>>),
                    Rhs_labels = gleam@string:to_graphemes(Rhs),
                    gleam@result:'try'(
                        check_labels(Rhs_labels, <<"right-hand side"/utf8>>),
                        fun(_) ->
                            Lhs_labels = gleam@list:map(
                                Lhs_parts,
                                fun gleam@string:to_graphemes/1
                            ),
                            gleam@result:'try'(
                                gleam@list:try_each(
                                    Lhs_labels,
                                    fun(Labels) ->
                                        check_labels(
                                            Labels,
                                            <<"left-hand side"/utf8>>
                                        )
                                    end
                                ),
                                fun(_) ->
                                    {ok, {parsed, Lhs_labels, Rhs_labels}}
                                end
                            )
                        end
                    );

                [_] ->
                    {error,
                        {dimension_error,
                            <<"einsum: feature not supported in v1: implicit output mode (missing '->')"/utf8>>}};

                _ ->
                    {error,
                        {invalid_shape,
                            <<"einsum: malformed equation (multiple '->' separators)"/utf8>>}}
            end
    end.

-file("src/viva_tensor/core/einsum.gleam", 39).
?DOC(false).
-spec run(binary(), list(operand())) -> {ok, operand()} |
    {error, viva_tensor@core@error:tensor_error()}.
run(Equation, Operands) ->
    gleam@result:'try'(
        parse_equation(Equation),
        fun(Parsed) ->
            {parsed, Lhs_labels, Rhs_labels} = Parsed,
            gleam@result:'try'(
                check_operand_arity(Lhs_labels, Operands),
                fun(_) ->
                    gleam@result:'try'(
                        check_rank_match(Lhs_labels, Operands),
                        fun(_) ->
                            gleam@result:'try'(
                                infer_dims(Lhs_labels, Operands),
                                fun(Dim_map) ->
                                    gleam@result:'try'(
                                        check_rhs_labels(Rhs_labels, Lhs_labels),
                                        fun(_) ->
                                            case erlang:length(Operands) > 2 of
                                                true ->
                                                    {error,
                                                        {dimension_error,
                                                            <<"einsum: feature not supported in v1: more than 2 input operands"/utf8>>}};

                                                false ->
                                                    general_kernel(
                                                        Operands,
                                                        Lhs_labels,
                                                        Rhs_labels,
                                                        Dim_map
                                                    )
                                            end
                                        end
                                    )
                                end
                            )
                        end
                    )
                end
            )
        end
    ).

-file("src/viva_tensor/core/einsum.gleam", 60).
?DOC(false).
-spec parse(binary()) -> {ok, {list(list(binary())), list(binary())}} |
    {error, viva_tensor@core@error:tensor_error()}.
parse(Equation) ->
    gleam@result:'try'(
        parse_equation(Equation),
        fun(Parsed) ->
            {parsed, Lhs, Rhs} = Parsed,
            {ok, {Lhs, Rhs}}
        end
    ).

-file("src/viva_tensor/core/einsum.gleam", 70).
?DOC(false).
-spec validate(list(list(binary())), list(binary()), list(operand())) -> {ok,
        list({binary(), integer()})} |
    {error, viva_tensor@core@error:tensor_error()}.
validate(Lhs, Rhs, Operands) ->
    gleam@result:'try'(
        check_operand_arity(Lhs, Operands),
        fun(_) ->
            gleam@result:'try'(
                check_rank_match(Lhs, Operands),
                fun(_) ->
                    gleam@result:'try'(
                        infer_dims(Lhs, Operands),
                        fun(Dim_map) ->
                            gleam@result:'try'(
                                check_rhs_labels(Rhs, Lhs),
                                fun(_) -> {ok, Dim_map} end
                            )
                        end
                    )
                end
            )
        end
    ).

-file("src/viva_tensor/core/einsum.gleam", 84).
?DOC(false).
-spec general(
    list(operand()),
    list(list(binary())),
    list(binary()),
    list({binary(), integer()})
) -> {ok, operand()} | {error, viva_tensor@core@error:tensor_error()}.
general(Operands, Lhs, Rhs, Dim_map) ->
    general_kernel(Operands, Lhs, Rhs, Dim_map).