viva_tensor

-module(viva_tensor@core@tensor).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_tensor/core/tensor.gleam").
-export([new/2, zeros/1, ones/1, fill/2, from_list/1, from_list2d/1, vector/1, matrix/3, eye/1, arange/3, linspace/3, random_uniform/1, random_normal/3, xavier_init/2, he_init/2, shape/1, to_list/1, size/1, rank/1, dim/2, rows/1, cols/1, get/2, get2d/3, get_row/2, get_col/2, to_strided/1, to_dense/1, to_contiguous/1, is_contiguous/1, transpose_strided/1, is_native/1, native_ref/1, from_native_ref/2, native_zeros/1, native_ones/1, native_fill/2, native_from_list/2]).
-export_type([tensor/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).

-opaque tensor() :: {dense, list(float()), list(integer())} |
    {strided,
        viva_tensor@core@ffi:erlang_array(),
        list(integer()),
        list(integer()),
        integer()} |
    {native, viva_tensor@core@ffi:native_tensor_ref(), list(integer())}.

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

-file("src/viva_tensor/core/tensor.gleam", 54).
?DOC(false).
-spec new(list(float()), list(integer())) -> {ok, tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
new(Data, Shape) ->
    Expected_size = compute_size(Shape),
    Actual_size = erlang:length(Data),
    case Expected_size =:= Actual_size of
        true ->
            {ok, {dense, Data, Shape}};

        false ->
            {error,
                {invalid_shape,
                    <<<<<<<<<<<<"Data size "/utf8,
                                            (erlang:integer_to_binary(
                                                Actual_size
                                            ))/binary>>/binary,
                                        " doesn't match shape "/utf8>>/binary,
                                    (viva_tensor@core@error:shape_to_string(
                                        Shape
                                    ))/binary>>/binary,
                                " (expected "/utf8>>/binary,
                            (erlang:integer_to_binary(Expected_size))/binary>>/binary,
                        ")"/utf8>>}}
    end.

-file("src/viva_tensor/core/tensor.gleam", 74).
?DOC(false).
-spec zeros(list(integer())) -> tensor().
zeros(Shape) ->
    Size = compute_size(Shape),
    {dense, gleam@list:repeat(+0.0, Size), Shape}.

-file("src/viva_tensor/core/tensor.gleam", 80).
?DOC(false).
-spec ones(list(integer())) -> tensor().
ones(Shape) ->
    Size = compute_size(Shape),
    {dense, gleam@list:repeat(1.0, Size), Shape}.

-file("src/viva_tensor/core/tensor.gleam", 86).
?DOC(false).
-spec fill(list(integer()), float()) -> tensor().
fill(Shape, Value) ->
    Size = compute_size(Shape),
    {dense, gleam@list:repeat(Value, Size), Shape}.

-file("src/viva_tensor/core/tensor.gleam", 92).
?DOC(false).
-spec from_list(list(float())) -> tensor().
from_list(Data) ->
    {dense, Data, [erlang:length(Data)]}.

-file("src/viva_tensor/core/tensor.gleam", 97).
?DOC(false).
-spec from_list2d(list(list(float()))) -> {ok, tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
from_list2d(Rows) ->
    case Rows of
        [] ->
            {ok, {dense, [], [0, 0]}};

        [First | Rest] ->
            Cols = erlang:length(First),
            Valid = gleam@list:all(
                Rest,
                fun(Row) -> erlang:length(Row) =:= Cols end
            ),
            case Valid of
                false ->
                    {error,
                        {invalid_shape, <<"Rows have different lengths"/utf8>>}};

                true ->
                    Data = lists:append(Rows),
                    Num_rows = erlang:length(Rows),
                    {ok, {dense, Data, [Num_rows, Cols]}}
            end
    end.

-file("src/viva_tensor/core/tensor.gleam", 117).
?DOC(false).
-spec vector(list(float())) -> tensor().
vector(Data) ->
    from_list(Data).

-file("src/viva_tensor/core/tensor.gleam", 122).
?DOC(false).
-spec matrix(integer(), integer(), list(float())) -> {ok, tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
matrix(Rows, Cols, Data) ->
    new(Data, [Rows, Cols]).

-file("src/viva_tensor/core/tensor.gleam", 132).
?DOC(false).
-spec eye(integer()) -> tensor().
eye(N) ->
    Data = begin
        _pipe = gleam@list:range(0, N - 1),
        gleam@list:flat_map(
            _pipe,
            fun(I) -> _pipe@1 = gleam@list:range(0, N - 1),
                gleam@list:map(_pipe@1, fun(J) -> case I =:= J of
                            true ->
                                1.0;

                            false ->
                                +0.0
                        end end) end
        )
    end,
    {dense, Data, [N, N]}.

-file("src/viva_tensor/core/tensor.gleam", 154).
?DOC(false).
-spec arange_loop(float(), float(), float(), list(float())) -> list(float()).
arange_loop(Current, End, Step, Acc) ->
    case Current >= End of
        true ->
            Acc;

        false ->
            arange_loop(Current + Step, End, Step, [Current | Acc])
    end.

-file("src/viva_tensor/core/tensor.gleam", 149).
?DOC(false).
-spec arange(float(), float(), float()) -> tensor().
arange(Start, End, Step) ->
    Data = arange_loop(Start, End, Step, []),
    from_list(lists:reverse(Data)).

-file("src/viva_tensor/core/tensor.gleam", 167).
?DOC(false).
-spec linspace(float(), float(), integer()) -> tensor().
linspace(Start, End, Num) ->
    case Num =< 1 of
        true ->
            from_list([Start]);

        false ->
            Step = case erlang:float(Num - 1) of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator -> (End - Start) / Gleam@denominator
            end,
            Data = begin
                _pipe = gleam@list:range(0, Num - 1),
                gleam@list:map(
                    _pipe,
                    fun(I) -> Start + (erlang:float(I) * Step) end
                )
            end,
            from_list(Data)
    end.

-file("src/viva_tensor/core/tensor.gleam", 186).
?DOC(false).
-spec random_uniform(list(integer())) -> tensor().
random_uniform(Shape) ->
    Size = compute_size(Shape),
    Data = begin
        _pipe = gleam@list:range(1, Size),
        gleam@list:map(
            _pipe,
            fun(_) -> viva_tensor@core@ffi:random_uniform() end
        )
    end,
    {dense, Data, Shape}.

-file("src/viva_tensor/core/tensor.gleam", 197).
?DOC(false).
-spec random_normal(list(integer()), float(), float()) -> tensor().
random_normal(Shape, Mean, Std) ->
    Size = compute_size(Shape),
    Data = begin
        _pipe = gleam@list:range(1, Size),
        gleam@list:map(
            _pipe,
            fun(_) ->
                U1 = gleam@float:max(
                    viva_tensor@core@ffi:random_uniform(),
                    0.0001
                ),
                U2 = viva_tensor@core@ffi:random_uniform(),
                Z = viva_tensor@core@ffi:sqrt(
                    -2.0 * viva_tensor@core@ffi:log(U1)
                )
                * viva_tensor@core@ffi:cos((2.0 * 3.14159265358979323846) * U2),
                Mean + (Z * Std)
            end
        )
    end,
    {dense, Data, Shape}.

-file("src/viva_tensor/core/tensor.gleam", 217).
?DOC(false).
-spec xavier_init(integer(), integer()) -> tensor().
xavier_init(Fan_in, Fan_out) ->
    Limit = viva_tensor@core@ffi:sqrt(case erlang:float(Fan_in + Fan_out) of
            +0.0 -> +0.0;
            -0.0 -> -0.0;
            Gleam@denominator -> 6.0 / Gleam@denominator
        end),
    Data = begin
        _pipe = gleam@list:range(1, Fan_in * Fan_out),
        gleam@list:map(
            _pipe,
            fun(_) ->
                R = viva_tensor@core@ffi:random_uniform(),
                ((R * 2.0) * Limit) - Limit
            end
        )
    end,
    {dense, Data, [Fan_out, Fan_in]}.

-file("src/viva_tensor/core/tensor.gleam", 232).
?DOC(false).
-spec he_init(integer(), integer()) -> tensor().
he_init(Fan_in, Fan_out) ->
    Std = viva_tensor@core@ffi:sqrt(case erlang:float(Fan_in) of
            +0.0 -> +0.0;
            -0.0 -> -0.0;
            Gleam@denominator -> 2.0 / Gleam@denominator
        end),
    random_normal([Fan_out, Fan_in], +0.0, Std).

-file("src/viva_tensor/core/tensor.gleam", 240).
?DOC(false).
-spec shape(tensor()) -> list(integer()).
shape(T) ->
    case T of
        {dense, _, S} ->
            S;

        {strided, _, S@1, _, _} ->
            S@1;

        {native, _, S@2} ->
            S@2
    end.

-file("src/viva_tensor/core/tensor.gleam", 496).
?DOC(false).
-spec flat_to_multi(integer(), list(integer())) -> list(integer()).
flat_to_multi(Flat, Shape) ->
    Reversed = lists:reverse(Shape),
    {Indices, _} = gleam@list:fold(
        Reversed,
        {[], Flat},
        fun(Acc, Dim) ->
            {Idxs, Remaining} = Acc,
            Idx = 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,
            {[Idx | Idxs], Next}
        end
    ),
    Indices.

-file("src/viva_tensor/core/tensor.gleam", 249).
?DOC(false).
-spec to_list(tensor()) -> list(float()).
to_list(T) ->
    case T of
        {dense, Data, _} ->
            Data;

        {native, Ref, _} ->
            case viva_tensor@core@ffi:nt_to_list(Ref) of
                {ok, Data@1} ->
                    Data@1;

                {error, _} ->
                    []
            end;

        {strided, Storage, Shp, Strides, Offset} ->
            Total_size = compute_size(Shp),
            _pipe = gleam@list:range(0, Total_size - 1),
            gleam@list:map(
                _pipe,
                fun(Flat_idx) ->
                    Indices = flat_to_multi(Flat_idx, Shp),
                    Idx = begin
                        _pipe@1 = gleam@list:zip(Indices, Strides),
                        gleam@list:fold(
                            _pipe@1,
                            Offset,
                            fun(Acc, Pair) ->
                                {I, S} = Pair,
                                Acc + (I * S)
                            end
                        )
                    end,
                    viva_tensor@core@ffi:array_get(Storage, Idx)
                end
            )
    end.

-file("src/viva_tensor/core/tensor.gleam", 275).
?DOC(false).
-spec size(tensor()) -> integer().
size(T) ->
    compute_size(shape(T)).

-file("src/viva_tensor/core/tensor.gleam", 280).
?DOC(false).
-spec rank(tensor()) -> integer().
rank(T) ->
    erlang:length(shape(T)).

-file("src/viva_tensor/core/tensor.gleam", 508).
?DOC(false).
-spec list_at(list(JQZ), integer()) -> {ok, JQZ} | {error, nil}.
list_at(Lst, Index) ->
    case Index < 0 of
        true ->
            {error, nil};

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

-file("src/viva_tensor/core/tensor.gleam", 285).
?DOC(false).
-spec dim(tensor(), integer()) -> {ok, integer()} |
    {error, viva_tensor@core@error:tensor_error()}.
dim(T, Axis) ->
    _pipe = list_at(shape(T), Axis),
    gleam@result:map_error(
        _pipe,
        fun(_) ->
            {dimension_error,
                <<<<"Axis "/utf8, (erlang:integer_to_binary(Axis))/binary>>/binary,
                    " out of bounds"/utf8>>}
        end
    ).

-file("src/viva_tensor/core/tensor.gleam", 293).
?DOC(false).
-spec rows(tensor()) -> integer().
rows(T) ->
    case shape(T) of
        [R | _] ->
            R;

        [] ->
            0
    end.

-file("src/viva_tensor/core/tensor.gleam", 301).
?DOC(false).
-spec cols(tensor()) -> integer().
cols(T) ->
    case shape(T) of
        [_, C | _] ->
            C;

        [N] ->
            N;

        [] ->
            0
    end.

-file("src/viva_tensor/core/tensor.gleam", 518).
?DOC(false).
-spec list_at_float(list(float()), integer()) -> {ok, float()} | {error, nil}.
list_at_float(Lst, Index) ->
    list_at(Lst, Index).

-file("src/viva_tensor/core/tensor.gleam", 312).
?DOC(false).
-spec get(tensor(), integer()) -> {ok, float()} |
    {error, viva_tensor@core@error:tensor_error()}.
get(T, Index) ->
    case T of
        {dense, Data, _} ->
            _pipe = list_at_float(Data, Index),
            gleam@result:map_error(
                _pipe,
                fun(_) -> {index_out_of_bounds, Index, size(T)} end
            );

        {native, Ref, _} ->
            case viva_tensor@core@ffi:nt_to_list(Ref) of
                {ok, Data@1} ->
                    _pipe@1 = list_at_float(Data@1, Index),
                    gleam@result:map_error(
                        _pipe@1,
                        fun(_) -> {index_out_of_bounds, Index, size(T)} end
                    );

                {error, _} ->
                    {error, {index_out_of_bounds, Index, 0}}
            end;

        {strided, Storage, Shp, Strides, Offset} ->
            Indices = flat_to_multi(Index, Shp),
            Flat_idx = begin
                _pipe@2 = gleam@list:zip(Indices, Strides),
                gleam@list:fold(
                    _pipe@2,
                    Offset,
                    fun(Acc, Pair) ->
                        {I, S} = Pair,
                        Acc + (I * S)
                    end
                )
            end,
            {ok, viva_tensor@core@ffi:array_get(Storage, Flat_idx)}
    end.

-file("src/viva_tensor/core/tensor.gleam", 340).
?DOC(false).
-spec get2d(tensor(), integer(), integer()) -> {ok, float()} |
    {error, viva_tensor@core@error:tensor_error()}.
get2d(T, Row, Col) ->
    case shape(T) of
        [_, Num_cols] ->
            get(T, (Row * Num_cols) + Col);

        _ ->
            {error, {dimension_error, <<"Tensor is not 2D"/utf8>>}}
    end.

-file("src/viva_tensor/core/tensor.gleam", 348).
?DOC(false).
-spec get_row(tensor(), integer()) -> {ok, tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
get_row(T, Row_idx) ->
    case shape(T) of
        [Num_rows, Num_cols] ->
            case (Row_idx >= 0) andalso (Row_idx < Num_rows) of
                true ->
                    Data = to_list(T),
                    Start = Row_idx * Num_cols,
                    Row_data = begin
                        _pipe = Data,
                        _pipe@1 = gleam@list:drop(_pipe, Start),
                        gleam@list:take(_pipe@1, Num_cols)
                    end,
                    {ok, from_list(Row_data)};

                false ->
                    {error, {index_out_of_bounds, Row_idx, Num_rows}}
            end;

        _ ->
            {error, {dimension_error, <<"Tensor is not 2D"/utf8>>}}
    end.

-file("src/viva_tensor/core/tensor.gleam", 369).
?DOC(false).
-spec get_col(tensor(), integer()) -> {ok, tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
get_col(T, Col_idx) ->
    case shape(T) of
        [Num_rows, Num_cols] ->
            case (Col_idx >= 0) andalso (Col_idx < Num_cols) of
                true ->
                    Col_data = begin
                        _pipe = gleam@list:range(0, Num_rows - 1),
                        gleam@list:filter_map(
                            _pipe,
                            fun(Row) -> get2d(T, Row, Col_idx) end
                        )
                    end,
                    {ok, from_list(Col_data)};

                false ->
                    {error, {index_out_of_bounds, Col_idx, Num_cols}}
            end;

        _ ->
            {error, {dimension_error, <<"Tensor is not 2D"/utf8>>}}
    end.

-file("src/viva_tensor/core/tensor.gleam", 486).
?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/tensor.gleam", 395).
?DOC(false).
-spec to_strided(tensor()) -> tensor().
to_strided(T) ->
    case T of
        {strided, _, _, _, _} ->
            T;

        {native, _, _} ->
            T;

        {dense, Data, Shp} ->
            Storage = viva_tensor@core@ffi:list_to_array(Data),
            Strides = compute_strides(Shp),
            {strided, Storage, Shp, Strides, 0}
    end.

-file("src/viva_tensor/core/tensor.gleam", 408).
?DOC(false).
-spec to_dense(tensor()) -> tensor().
to_dense(T) ->
    case T of
        {dense, _, _} ->
            T;

        {native, _, _} ->
            Data = to_list(T),
            {dense, Data, shape(T)};

        {strided, _, _, _, _} ->
            Data@1 = to_list(T),
            {dense, Data@1, shape(T)}
    end.

-file("src/viva_tensor/core/tensor.gleam", 423).
?DOC(false).
-spec to_contiguous(tensor()) -> tensor().
to_contiguous(T) ->
    to_dense(T).

-file("src/viva_tensor/core/tensor.gleam", 428).
?DOC(false).
-spec is_contiguous(tensor()) -> boolean().
is_contiguous(T) ->
    case T of
        {dense, _, _} ->
            true;

        {native, _, _} ->
            true;

        {strided, _, Shp, Strides, _} ->
            Expected_strides = compute_strides(Shp),
            Strides =:= Expected_strides
    end.

-file("src/viva_tensor/core/tensor.gleam", 440).
?DOC(false).
-spec transpose_strided(tensor()) -> {ok, tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
transpose_strided(T) ->
    case T of
        {native, Ref, Shp} ->
            case Shp of
                [_, _] ->
                    case viva_tensor@core@ffi:nt_transpose(Ref) of
                        {ok, Ref_t} ->
                            {ok, {native, Ref_t, lists:reverse(Shp)}};

                        {error, _} ->
                            Dense = to_dense(T),
                            transpose_strided(Dense)
                    end;

                _ ->
                    {error,
                        {dimension_error,
                            <<"Transpose requires 2D tensor"/utf8>>}}
            end;

        {dense, _, Shp@1} ->
            case Shp@1 of
                [_, _] ->
                    Strided = to_strided(T),
                    transpose_strided(Strided);

                _ ->
                    {error,
                        {dimension_error,
                            <<"Transpose requires 2D tensor"/utf8>>}}
            end;

        {strided, Storage, Shp@2, Strides, Offset} ->
            case {Shp@2, Strides} of
                {[M, N], [S0, S1]} ->
                    {ok, {strided, Storage, [N, M], [S1, S0], Offset}};

                {_, _} ->
                    {error,
                        {dimension_error,
                            <<"Transpose requires 2D tensor"/utf8>>}}
            end
    end.

-file("src/viva_tensor/core/tensor.gleam", 527).
?DOC(false).
-spec is_native(tensor()) -> boolean().
is_native(T) ->
    case T of
        {native, _, _} ->
            true;

        _ ->
            false
    end.

-file("src/viva_tensor/core/tensor.gleam", 535).
?DOC(false).
-spec native_ref(tensor()) -> {ok, viva_tensor@core@ffi:native_tensor_ref()} |
    {error, nil}.
native_ref(T) ->
    case T of
        {native, Ref, _} ->
            {ok, Ref};

        _ ->
            {error, nil}
    end.

-file("src/viva_tensor/core/tensor.gleam", 543).
?DOC(false).
-spec from_native_ref(viva_tensor@core@ffi:native_tensor_ref(), list(integer())) -> tensor().
from_native_ref(Ref, Shape) ->
    {native, Ref, Shape}.

-file("src/viva_tensor/core/tensor.gleam", 548).
?DOC(false).
-spec native_zeros(list(integer())) -> {ok, tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
native_zeros(Shape) ->
    case viva_tensor@core@ffi:nt_zeros(Shape) of
        {ok, Ref} ->
            {ok, {native, Ref, Shape}};

        {error, _} ->
            {error, {invalid_shape, <<"NIF resource allocation failed"/utf8>>}}
    end.

-file("src/viva_tensor/core/tensor.gleam", 556).
?DOC(false).
-spec native_ones(list(integer())) -> {ok, tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
native_ones(Shape) ->
    case viva_tensor@core@ffi:nt_ones(Shape) of
        {ok, Ref} ->
            {ok, {native, Ref, Shape}};

        {error, _} ->
            {error, {invalid_shape, <<"NIF resource allocation failed"/utf8>>}}
    end.

-file("src/viva_tensor/core/tensor.gleam", 564).
?DOC(false).
-spec native_fill(list(integer()), float()) -> {ok, tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
native_fill(Shape, Value) ->
    case viva_tensor@core@ffi:nt_fill(Shape, Value) of
        {ok, Ref} ->
            {ok, {native, Ref, Shape}};

        {error, _} ->
            {error, {invalid_shape, <<"NIF resource allocation failed"/utf8>>}}
    end.

-file("src/viva_tensor/core/tensor.gleam", 575).
?DOC(false).
-spec native_from_list(list(float()), list(integer())) -> {ok, tensor()} |
    {error, viva_tensor@core@error:tensor_error()}.
native_from_list(Data, Shape) ->
    case viva_tensor@core@ffi:nt_from_list(Data, Shape) of
        {ok, Ref} ->
            {ok, {native, Ref, Shape}};

        {error, _} ->
            Expected = compute_size(Shape),
            Actual = erlang:length(Data),
            case Expected =:= Actual of
                true ->
                    {error,
                        {invalid_shape,
                            <<"NIF resource allocation failed"/utf8>>}};

                false ->
                    {error,
                        {invalid_shape,
                            <<<<<<"Data size "/utf8,
                                        (erlang:integer_to_binary(Actual))/binary>>/binary,
                                    " doesn't match shape "/utf8>>/binary,
                                (viva_tensor@core@error:shape_to_string(Shape))/binary>>}}
            end
    end.