mlx

-module(mlx).

%% Type definitions
-type array() :: reference().
-type shape() :: [integer()].
-type dtype() :: atom().

%% Core array creation functions
-export([array/1, array/2, zeros/1, zeros/2, ones/1, ones/2, full/2, full/3, eye/1, eye/2, eye/3,
         arange/1, arange/2, arange/3, linspace/2, linspace/3, linspace/4]).

%% Basic arithmetic operations
-export([add/2, subtract/2, multiply/2, divide/2, floor_divide/2, remainder/2,
         divmod/2, power/2, abs/1, negative/1, reciprocal/1, sign/1]).

%% Trigonometric operations
-export([sin/1, cos/1, tan/1, arcsin/1, arccos/1, arctan/1, arctan2/2,
         sinh/1, cosh/1, tanh/1, arcsinh/1, arccosh/1, arctanh/1]).

%% Exponential and logarithmic operations
-export([exp/1, expm1/1, log/1, log2/1, log10/1, log1p/1, logaddexp/2,
         sqrt/1, square/1, rsqrt/1]).

%% Array manipulation operations
-export([reshape/2, transpose/1, transpose/2, squeeze/1, squeeze/2,
         expand_dims/2, flatten/1, flatten/2, flatten/3, swapaxes/3,
         moveaxis/3, roll/2, roll/3]).

%% Reduction operations
-export([sum/1, sum/2, sum/3, mean/1, mean/2, mean/3, max/1, max/2, max/3,
         min/1, min/2, min/3, prod/1, prod/2, prod/3, std/1, std/2, std/3,
         var/1, var/2, var/3]).

%% Cumulative operations
-export([cumsum/1, cumsum/2, cumprod/1, cumprod/2, cummax/1, cummax/2,
         cummin/1, cummin/2]).

%% Linear algebra operations
-export([matmul/2, dot/2, inner/2, outer/2, tensordot/2, tensordot/3]).

%% Comparison operations
-export([equal/2, not_equal/2, less/2, less_equal/2, greater/2, greater_equal/2,
         allclose/2, allclose/3, array_equal/2, isfinite/1, isinf/1, isnan/1]).

%% Logical operations
-export([logical_and/2, logical_or/2, logical_not/1, all/1, all/2, all/3,
         any/1, any/2, any/3]).

%% Bitwise operations
-export([bitwise_and/2, bitwise_or/2, bitwise_xor/2, bitwise_invert/1,
         left_shift/2, right_shift/2]).

%% Selection and indexing operations
-export([where/3, take/2, take/3, argmax/1, argmax/2, argmin/1, argmin/2,
         sort/1, sort/2, argsort/1, argsort/2, topk/2, topk/3]).

%% Array concatenation and stacking
-export([concatenate/1, concatenate/2, stack/1, stack/2, split/2, split/3]).

%% Broadcasting operations
-export([broadcast_to/2, broadcast_arrays/1]).

%% Clipping and rounding operations
-export([clip/3, ceil/1, floor/1, round/1, round/2]).

%% Array properties and conversion
-export([shape/1, size/1, ndim/1, dtype/1, eval/1, to_list/1]).

%% Neural network operations
-export([softmax/1, softmax/2, sigmoid/1, conv1d/2, conv1d/3, conv2d/2, conv2d/3]).

%% Device management
-export([use_cpu/0, use_gpu/0]).

%% Testing and utilities
-export([test_basic/0, version/0]).


%% Array creation functions
-spec array(term()) -> {ok, array()} | {error, term()}.
array(Data) ->
    array(Data, float32).

-spec array(term(), dtype()) -> {ok, array()} | {error, term()}.
array(Data, Dtype) ->
    mlx_nif:array(Data, Dtype).

-spec zeros(shape()) -> {ok, array()} | {error, term()}.
zeros(Shape) ->
    zeros(Shape, float32).

-spec zeros(shape(), dtype()) -> {ok, array()} | {error, term()}.
zeros(Shape, Dtype) ->
    mlx_nif:zeros(Shape, Dtype).

-spec ones(shape()) -> {ok, array()} | {error, term()}.
ones(Shape) ->
    ones(Shape, float32).

-spec ones(shape(), dtype()) -> {ok, array()} | {error, term()}.
ones(Shape, Dtype) ->
    mlx_nif:ones(Shape, Dtype).

-spec full(shape(), number()) -> {ok, array()} | {error, term()}.
full(Shape, Value) ->
    full(Shape, Value, float32).

-spec full(shape(), number(), dtype()) -> {ok, array()} | {error, term()}.
full(Shape, Value, Dtype) ->
    mlx_nif:full(Shape, Value, Dtype).

-spec eye(integer()) -> {ok, array()} | {error, term()}.
eye(N) ->
    eye(N, N, float32).

-spec eye(integer(), integer()) -> {ok, array()} | {error, term()}.
eye(N, M) ->
    eye(N, M, float32).

-spec eye(integer(), integer(), dtype()) -> {ok, array()} | {error, term()}.
eye(N, M, Dtype) ->
    case N of
        M -> mlx_nif:eye(N, Dtype);
        _ -> {error, non_square_matrix_not_supported}
    end.

-spec arange(number()) -> {ok, array()} | {error, term()}.
arange(Stop) ->
    arange(0, Stop, 1).

-spec arange(number(), number()) -> {ok, array()} | {error, term()}.
arange(Start, Stop) ->
    arange(Start, Stop, 1).

-spec arange(number(), number(), number()) -> {ok, array()} | {error, term()}.
arange(Start, Stop, Step) ->
    mlx_nif:arange(Start, Stop, Step).

-spec linspace(number(), number()) -> {ok, array()} | {error, term()}.
linspace(Start, Stop) ->
    linspace(Start, Stop, 50).

-spec linspace(number(), number(), integer()) -> {ok, array()} | {error, term()}.
linspace(Start, Stop, Num) ->
    linspace(Start, Stop, Num, float32).

-spec linspace(number(), number(), integer(), dtype()) -> {ok, array()} | {error, term()}.
linspace(Start, Stop, Num, Dtype) ->
    case Dtype of
        float32 -> mlx_nif:linspace(Start, Stop, Num);
        _ -> {error, unsupported_dtype_use_float32}
    end.

%% Basic arithmetic operations
-spec add(array(), array()) -> {ok, array()} | {error, term()}.
add(A, B) ->
    mlx_nif:add(A, B).

-spec subtract(array(), array()) -> {ok, array()} | {error, term()}.
subtract(A, B) ->
    mlx_nif:subtract(A, B).

-spec multiply(array(), array()) -> {ok, array()} | {error, term()}.
multiply(A, B) ->
    mlx_nif:multiply(A, B).

-spec divide(array(), array()) -> {ok, array()} | {error, term()}.
divide(A, B) ->
    mlx_nif:divide(A, B).

-spec floor_divide(array(), array()) -> {ok, array()} | {error, term()}.
floor_divide(A, B) ->
    mlx_nif:floor_divide(A, B).

-spec remainder(array(), array()) -> {ok, array()} | {error, term()}.
remainder(A, B) ->
    mlx_nif:remainder(A, B).

-spec divmod(array(), array()) -> {ok, {array(), array()}} | {error, term()}.
divmod(A, B) ->
    mlx_nif:divmod(A, B).

-spec power(array(), array()) -> {ok, array()} | {error, term()}.
power(A, B) ->
    mlx_nif:power(A, B).

-spec abs(array()) -> {ok, array()} | {error, term()}.
abs(A) ->
    mlx_nif:abs(A).

-spec negative(array()) -> {ok, array()} | {error, term()}.
negative(A) ->
    mlx_nif:negative(A).

-spec reciprocal(array()) -> {ok, array()} | {error, term()}.
reciprocal(A) ->
    mlx_nif:reciprocal(A).

-spec sign(array()) -> {ok, array()} | {error, term()}.
sign(A) ->
    mlx_nif:sign(A).

%% Trigonometric operations
-spec sin(array()) -> {ok, array()} | {error, term()}.
sin(A) ->
    mlx_nif:sin(A).

-spec cos(array()) -> {ok, array()} | {error, term()}.
cos(A) ->
    mlx_nif:cos(A).

-spec tan(array()) -> {ok, array()} | {error, term()}.
tan(A) ->
    mlx_nif:tan(A).

-spec arcsin(array()) -> {ok, array()} | {error, term()}.
arcsin(A) ->
    mlx_nif:arcsin(A).

-spec arccos(array()) -> {ok, array()} | {error, term()}.
arccos(A) ->
    mlx_nif:arccos(A).

-spec arctan(array()) -> {ok, array()} | {error, term()}.
arctan(A) ->
    mlx_nif:arctan(A).

-spec arctan2(array(), array()) -> {ok, array()} | {error, term()}.
arctan2(A, B) ->
    mlx_nif:arctan2(A, B).

-spec sinh(array()) -> {ok, array()} | {error, term()}.
sinh(A) ->
    mlx_nif:sinh(A).

-spec cosh(array()) -> {ok, array()} | {error, term()}.
cosh(A) ->
    mlx_nif:cosh(A).

-spec tanh(array()) -> {ok, array()} | {error, term()}.
tanh(A) ->
    mlx_nif:tanh(A).

-spec arcsinh(array()) -> {ok, array()} | {error, term()}.
arcsinh(A) ->
    mlx_nif:arcsinh(A).

-spec arccosh(array()) -> {ok, array()} | {error, term()}.
arccosh(A) ->
    mlx_nif:arccosh(A).

-spec arctanh(array()) -> {ok, array()} | {error, term()}.
arctanh(A) ->
    mlx_nif:arctanh(A).

%% Exponential and logarithmic operations
-spec exp(array()) -> {ok, array()} | {error, term()}.
exp(A) ->
    mlx_nif:exp(A).

-spec expm1(array()) -> {ok, array()} | {error, term()}.
expm1(A) ->
    mlx_nif:expm1(A).

-spec log(array()) -> {ok, array()} | {error, term()}.
log(A) ->
    mlx_nif:log(A).

-spec log2(array()) -> {ok, array()} | {error, term()}.
log2(A) ->
    mlx_nif:log2(A).

-spec log10(array()) -> {ok, array()} | {error, term()}.
log10(A) ->
    mlx_nif:log10(A).

-spec log1p(array()) -> {ok, array()} | {error, term()}.
log1p(A) ->
    mlx_nif:log1p(A).

-spec logaddexp(array(), array()) -> {ok, array()} | {error, term()}.
logaddexp(A, B) ->
    mlx_nif:logaddexp(A, B).

-spec sqrt(array()) -> {ok, array()} | {error, term()}.
sqrt(A) ->
    mlx_nif:sqrt(A).

-spec square(array()) -> {ok, array()} | {error, term()}.
square(A) ->
    mlx_nif:square(A).

-spec rsqrt(array()) -> {ok, array()} | {error, term()}.
rsqrt(A) ->
    mlx_nif:rsqrt(A).

%% Array manipulation operations
-spec reshape(array(), shape()) -> {ok, array()} | {error, term()}.
reshape(A, Shape) ->
    mlx_nif:reshape(A, Shape).

-spec transpose(array()) -> {ok, array()} | {error, term()}.
transpose(A) ->
    transpose(A, []).

-spec transpose(array(), list()) -> {ok, array()} | {error, term()}.
transpose(A, Axes) ->
    mlx_nif:transpose(A, Axes).

-spec squeeze(array()) -> {ok, array()} | {error, term()}.
squeeze(A) ->
    squeeze(A, []).

-spec squeeze(array(), list()) -> {ok, array()} | {error, term()}.
squeeze(A, Axis) ->
    mlx_nif:squeeze(A, Axis).

-spec expand_dims(array(), integer()) -> {ok, array()} | {error, term()}.
expand_dims(A, Axis) ->
    mlx_nif:expand_dims(A, Axis).

-spec flatten(array()) -> {ok, array()} | {error, term()}.
flatten(A) ->
    flatten(A, 0, -1).

-spec flatten(array(), integer()) -> {ok, array()} | {error, term()}.
flatten(A, StartAxis) ->
    flatten(A, StartAxis, -1).

-spec flatten(array(), integer(), integer()) -> {ok, array()} | {error, term()}.
flatten(A, StartAxis, EndAxis) ->
    mlx_nif:flatten(A, StartAxis, EndAxis).

-spec swapaxes(array(), integer(), integer()) -> {ok, array()} | {error, term()}.
swapaxes(A, Axis1, Axis2) ->
    mlx_nif:swapaxes(A, Axis1, Axis2).

-spec moveaxis(array(), integer(), integer()) -> {ok, array()} | {error, term()}.
moveaxis(A, Source, Destination) ->
    mlx_nif:moveaxis(A, Source, Destination).

-spec roll(array(), integer()) -> {ok, array()} | {error, term()}.
roll(A, Shift) ->
    roll(A, Shift, []).

-spec roll(array(), integer(), list()) -> {ok, array()} | {error, term()}.
roll(A, Shift, Axis) ->
    mlx_nif:roll(A, Shift, Axis).

%% Reduction operations
-spec sum(array()) -> {ok, array()} | {error, term()}.
sum(A) ->
    mlx_nif:sum(A).

-spec sum(array(), list()) -> {ok, array()} | {error, term()}.
sum(A, Axis) ->
    mlx_nif:sum(A, Axis).

-spec sum(array(), list(), boolean()) -> {ok, array()} | {error, term()}.
sum(A, Axis, Keepdims) ->
    mlx_nif:sum(A, Axis, Keepdims).

-spec mean(array()) -> {ok, array()} | {error, term()}.
mean(A) ->
    mlx_nif:mean(A).

-spec mean(array(), list()) -> {ok, array()} | {error, term()}.
mean(A, Axis) ->
    mlx_nif:mean(A, Axis).

-spec mean(array(), list(), boolean()) -> {ok, array()} | {error, term()}.
mean(A, Axis, Keepdims) ->
    mlx_nif:mean(A, Axis, Keepdims).

-spec max(array()) -> {ok, array()} | {error, term()}.
max(A) ->
    mlx_nif:max(A).

-spec max(array(), list()) -> {ok, array()} | {error, term()}.
max(A, Axis) ->
    mlx_nif:max(A, Axis).

-spec max(array(), list(), boolean()) -> {ok, array()} | {error, term()}.
max(A, Axis, Keepdims) ->
    mlx_nif:max(A, Axis, Keepdims).

-spec min(array()) -> {ok, array()} | {error, term()}.
min(A) ->
    mlx_nif:min(A).

-spec min(array(), list()) -> {ok, array()} | {error, term()}.
min(A, Axis) ->
    mlx_nif:min(A, Axis).

-spec min(array(), list(), boolean()) -> {ok, array()} | {error, term()}.
min(A, Axis, Keepdims) ->
    mlx_nif:min(A, Axis, Keepdims).

-spec prod(array()) -> {ok, array()} | {error, term()}.
prod(A) ->
    prod(A, [], true).

-spec prod(array(), list()) -> {ok, array()} | {error, term()}.
prod(A, Axis) ->
    prod(A, Axis, false).

-spec prod(array(), list(), boolean()) -> {ok, array()} | {error, term()}.
prod(A, Axis, Keepdims) ->
    mlx_nif:prod(A, Axis, Keepdims).

-spec std(array()) -> {ok, array()} | {error, term()}.
std(A) ->
    std(A, [], true).

-spec std(array(), list()) -> {ok, array()} | {error, term()}.
std(A, Axis) ->
    std(A, Axis, false).

-spec std(array(), list(), boolean()) -> {ok, array()} | {error, term()}.
std(A, Axis, Keepdims) ->
    mlx_nif:std(A, Axis, Keepdims).

-spec var(array()) -> {ok, array()} | {error, term()}.
var(A) ->
    var(A, [], true).

-spec var(array(), list()) -> {ok, array()} | {error, term()}.
var(A, Axis) ->
    var(A, Axis, false).

-spec var(array(), list(), boolean()) -> {ok, array()} | {error, term()}.
var(A, Axis, Keepdims) ->
    mlx_nif:var(A, Axis, Keepdims).

%% Cumulative operations
-spec cumsum(array()) -> {ok, array()} | {error, term()}.
cumsum(A) ->
    cumsum(A, 0).

-spec cumsum(array(), integer()) -> {ok, array()} | {error, term()}.
cumsum(A, Axis) ->
    mlx_nif:cumsum(A, Axis).

-spec cumprod(array()) -> {ok, array()} | {error, term()}.
cumprod(A) ->
    cumprod(A, 0).

-spec cumprod(array(), integer()) -> {ok, array()} | {error, term()}.
cumprod(A, Axis) ->
    mlx_nif:cumprod(A, Axis).

-spec cummax(array()) -> {ok, array()} | {error, term()}.
cummax(A) ->
    cummax(A, 0).

-spec cummax(array(), integer()) -> {ok, array()} | {error, term()}.
cummax(A, Axis) ->
    mlx_nif:cummax(A, Axis).

-spec cummin(array()) -> {ok, array()} | {error, term()}.
cummin(A) ->
    cummin(A, 0).

-spec cummin(array(), integer()) -> {ok, array()} | {error, term()}.
cummin(A, Axis) ->
    mlx_nif:cummin(A, Axis).

%% Linear algebra operations
-spec matmul(array(), array()) -> {ok, array()} | {error, term()}.
matmul(A, B) ->
    mlx_nif:matmul(A, B).

-spec dot(array(), array()) -> {ok, array()} | {error, term()}.
dot(A, B) ->
    mlx_nif:dot(A, B).

-spec inner(array(), array()) -> {ok, array()} | {error, term()}.
inner(A, B) ->
    mlx_nif:inner(A, B).

-spec outer(array(), array()) -> {ok, array()} | {error, term()}.
outer(A, B) ->
    mlx_nif:outer(A, B).

-spec tensordot(array(), array()) -> {ok, array()} | {error, term()}.
tensordot(A, B) ->
    tensordot(A, B, 2).

-spec tensordot(array(), array(), integer()) -> {ok, array()} | {error, term()}.
tensordot(A, B, Axes) ->
    mlx_nif:tensordot(A, B, Axes).

%% Comparison operations
-spec equal(array(), array()) -> {ok, array()} | {error, term()}.
equal(A, B) ->
    mlx_nif:equal(A, B).

-spec not_equal(array(), array()) -> {ok, array()} | {error, term()}.
not_equal(A, B) ->
    mlx_nif:not_equal(A, B).

-spec less(array(), array()) -> {ok, array()} | {error, term()}.
less(A, B) ->
    mlx_nif:less(A, B).

-spec less_equal(array(), array()) -> {ok, array()} | {error, term()}.
less_equal(A, B) ->
    mlx_nif:less_equal(A, B).

-spec greater(array(), array()) -> {ok, array()} | {error, term()}.
greater(A, B) ->
    mlx_nif:greater(A, B).

-spec greater_equal(array(), array()) -> {ok, array()} | {error, term()}.
greater_equal(A, B) ->
    mlx_nif:greater_equal(A, B).

-spec maximum(array(), array()) -> {ok, array()} | {error, term()}.
maximum(A, B) ->
    mlx_nif:maximum(A, B).

-spec minimum(array(), array()) -> {ok, array()} | {error, term()}.
minimum(A, B) ->
    mlx_nif:minimum(A, B).

-spec allclose(array(), array()) -> {ok, boolean()} | {error, term()}.
allclose(A, B) ->
    allclose(A, B, 1.0e-5).

-spec allclose(array(), array(), float()) -> {ok, boolean()} | {error, term()}.
allclose(A, B, Tolerance) ->
    mlx_nif:allclose(A, B, Tolerance).

-spec array_equal(array(), array()) -> {ok, boolean()} | {error, term()}.
array_equal(A, B) ->
    mlx_nif:array_equal(A, B).

-spec isfinite(array()) -> {ok, array()} | {error, term()}.
isfinite(A) ->
    mlx_nif:isfinite(A).

-spec isinf(array()) -> {ok, array()} | {error, term()}.
isinf(A) ->
    mlx_nif:isinf(A).

-spec isnan(array()) -> {ok, array()} | {error, term()}.
isnan(A) ->
    mlx_nif:isnan(A).

%% Logical operations
-spec logical_and(array(), array()) -> {ok, array()} | {error, term()}.
logical_and(A, B) ->
    mlx_nif:logical_and(A, B).

-spec logical_or(array(), array()) -> {ok, array()} | {error, term()}.
logical_or(A, B) ->
    mlx_nif:logical_or(A, B).

-spec logical_not(array()) -> {ok, array()} | {error, term()}.
logical_not(A) ->
    mlx_nif:logical_not(A).

-spec all(array()) -> {ok, array()} | {error, term()}.
all(A) ->
    all(A, [], true).

-spec all(array(), list()) -> {ok, array()} | {error, term()}.
all(A, Axis) ->
    all(A, Axis, false).

-spec all(array(), list(), boolean()) -> {ok, array()} | {error, term()}.
all(A, Axis, Keepdims) ->
    mlx_nif:all(A, Axis, Keepdims).

-spec any(array()) -> {ok, array()} | {error, term()}.
any(A) ->
    any(A, [], true).

-spec any(array(), list()) -> {ok, array()} | {error, term()}.
any(A, Axis) ->
    any(A, Axis, false).

-spec any(array(), list(), boolean()) -> {ok, array()} | {error, term()}.
any(A, Axis, Keepdims) ->
    mlx_nif:any(A, Axis, Keepdims).

%% Bitwise operations
-spec bitwise_and(array(), array()) -> {ok, array()} | {error, term()}.
bitwise_and(A, B) ->
    mlx_nif:bitwise_and(A, B).

-spec bitwise_or(array(), array()) -> {ok, array()} | {error, term()}.
bitwise_or(A, B) ->
    mlx_nif:bitwise_or(A, B).

-spec bitwise_xor(array(), array()) -> {ok, array()} | {error, term()}.
bitwise_xor(A, B) ->
    mlx_nif:bitwise_xor(A, B).

-spec bitwise_invert(array()) -> {ok, array()} | {error, term()}.
bitwise_invert(A) ->
    mlx_nif:bitwise_invert(A).

-spec left_shift(array(), array()) -> {ok, array()} | {error, term()}.
left_shift(A, B) ->
    mlx_nif:left_shift(A, B).

-spec right_shift(array(), array()) -> {ok, array()} | {error, term()}.
right_shift(A, B) ->
    mlx_nif:right_shift(A, B).

%% Selection and indexing operations
-spec where(array(), array(), array()) -> {ok, array()} | {error, term()}.
where(Condition, X, Y) ->
    mlx_nif:where(Condition, X, Y).

-spec take(array(), array()) -> {ok, array()} | {error, term()}.
take(A, Indices) ->
    take(A, Indices, 0).

-spec take(array(), array(), integer()) -> {ok, array()} | {error, term()}.
take(A, Indices, Axis) ->
    mlx_nif:take(A, Indices, Axis).

-spec argmax(array()) -> {ok, array()} | {error, term()}.
argmax(A) ->
    argmax(A, 0).

-spec argmax(array(), integer()) -> {ok, array()} | {error, term()}.
argmax(A, Axis) ->
    mlx_nif:argmax(A, Axis).

-spec argmin(array()) -> {ok, array()} | {error, term()}.
argmin(A) ->
    argmin(A, 0).

-spec argmin(array(), integer()) -> {ok, array()} | {error, term()}.
argmin(A, Axis) ->
    mlx_nif:argmin(A, Axis).

-spec sort(array()) -> {ok, array()} | {error, term()}.
sort(A) ->
    sort(A, -1).

-spec sort(array(), integer()) -> {ok, array()} | {error, term()}.
sort(A, Axis) ->
    mlx_nif:sort(A, Axis).

-spec argsort(array()) -> {ok, array()} | {error, term()}.
argsort(A) ->
    argsort(A, -1).

-spec argsort(array(), integer()) -> {ok, array()} | {error, term()}.
argsort(A, Axis) ->
    mlx_nif:argsort(A, Axis).

-spec topk(array(), integer()) -> {ok, array()} | {error, term()}.
topk(A, K) ->
    topk(A, K, -1).

-spec topk(array(), integer(), integer()) -> {ok, array()} | {error, term()}.
topk(A, K, Axis) ->
    mlx_nif:topk(A, K, Axis).

%% Array concatenation and stacking
-spec concatenate(list()) -> {ok, array()} | {error, term()}.
concatenate(Arrays) ->
    concatenate(Arrays, 0).

-spec concatenate(list(), integer()) -> {ok, array()} | {error, term()}.
concatenate(Arrays, Axis) ->
    mlx_nif:concatenate(Arrays, Axis).

-spec stack(list()) -> {ok, array()} | {error, term()}.
stack(Arrays) ->
    stack(Arrays, 0).

-spec stack(list(), integer()) -> {ok, array()} | {error, term()}.
stack(Arrays, Axis) ->
    mlx_nif:stack(Arrays, Axis).

-spec split(array(), list()) -> {ok, list()} | {error, term()}.
split(A, Indices) ->
    split(A, Indices, 0).

-spec split(array(), list(), integer()) -> {ok, list()} | {error, term()}.
split(A, Indices, Axis) ->
    mlx_nif:split(A, Indices, Axis).

%% Broadcasting operations
-spec broadcast_to(array(), shape()) -> {ok, array()} | {error, term()}.
broadcast_to(A, Shape) ->
    mlx_nif:broadcast_to(A, Shape).

-spec broadcast_arrays(list()) -> {ok, list()} | {error, term()}.
broadcast_arrays(Arrays) ->
    mlx_nif:broadcast_arrays(Arrays).

%% Clipping and rounding operations
-spec clip(array(), number(), number()) -> {ok, array()} | {error, term()}.
clip(A, Min, Max) ->
    mlx_nif:clip(A, Min, Max).

-spec ceil(array()) -> {ok, array()} | {error, term()}.
ceil(A) ->
    mlx_nif:ceil(A).

-spec floor(array()) -> {ok, array()} | {error, term()}.
floor(A) ->
    mlx_nif:floor(A).

-spec round(array()) -> {ok, array()} | {error, term()}.
round(A) ->
    round(A, 0).

-spec round(array(), integer()) -> {ok, array()} | {error, term()}.
round(A, Decimals) ->
    mlx_nif:round(A, Decimals).

%% Neural network operations
-spec softmax(array()) -> {ok, array()} | {error, term()}.
softmax(A) ->
    softmax(A, -1).

-spec softmax(array(), integer()) -> {ok, array()} | {error, term()}.
softmax(A, Axis) ->
    mlx_nif:softmax(A, Axis).

-spec sigmoid(array()) -> {ok, array()} | {error, term()}.
sigmoid(A) ->
    mlx_nif:sigmoid(A).

-spec conv1d(array(), array()) -> {ok, array()} | {error, term()}.
conv1d(Input, Weight) ->
    conv1d(Input, Weight, 1).

-spec conv1d(array(), array(), integer()) -> {ok, array()} | {error, term()}.
conv1d(Input, Weight, Stride) ->
    mlx_nif:conv1d(Input, Weight, Stride).

-spec conv2d(array(), array()) -> {ok, array()} | {error, term()}.
conv2d(Input, Weight) ->
    conv2d(Input, Weight, [1, 1]).

-spec conv2d(array(), array(), list()) -> {ok, array()} | {error, term()}.
conv2d(Input, Weight, Stride) ->
    mlx_nif:conv2d(Input, Weight, Stride).

%% Array properties and conversion
-spec shape(array()) -> {ok, shape()} | {error, term()}.
shape(Array) ->
    mlx_nif:shape(Array).

-spec size(array()) -> {ok, integer()} | {error, term()}.
size(Array) ->
    mlx_nif:size(Array).

-spec ndim(array()) -> {ok, integer()} | {error, term()}.
ndim(Array) ->
    mlx_nif:ndim(Array).

-spec dtype(array()) -> {ok, string()} | {error, term()}.
dtype(Array) ->
    mlx_nif:dtype_str(Array).

-spec eval(array()) -> {ok, array()} | {error, term()}.
eval(Array) ->
    mlx_nif:eval(Array).

-spec to_list(array()) -> {ok, list()} | {error, term()}.
to_list(Array) ->
    mlx_nif:to_list(Array).


%% Device management
-spec use_cpu() -> ok | {error, term()}.
use_cpu() ->
    mlx_nif:set_default_device(cpu).

-spec use_gpu() -> ok | {error, term()}.
use_gpu() ->
    mlx_nif:set_default_device(gpu).

%% Testing and utilities
-spec test_basic() -> ok | {error, term()}.
test_basic() ->
    mlx_nif:test_basic().

-spec version() -> {ok, string()} | {error, term()}.
version() ->
    mlx_nif:version().