diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml index 22e099cb2..456963ddf 100644 --- a/.github/workflows/release.yaml +++ b/.github/workflows/release.yaml @@ -171,6 +171,11 @@ jobs: - '"cufft_dev"' - '"nvrtc"' - '"nvrtc_dev"' + - '"cusolver"' + - '"cusolver_dev"' + - '"cusparse"' + - '"cusparse_dev"' + - '"nvjitlink"' - '"crt"' - '"nvvm"' - '"nvptxcompiler"' diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 430a20e0e..8a0fec6c1 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -231,6 +231,11 @@ jobs: - '"cufft_dev"' - '"nvrtc"' - '"nvrtc_dev"' + - '"cusolver"' + - '"cusolver_dev"' + - '"cusparse"' + - '"cusparse_dev"' + - '"nvjitlink"' - '"crt"' - '"nvvm"' - '"nvptxcompiler"' diff --git a/Dockerfile b/Dockerfile index 1771dc197..c18d0656a 100644 --- a/Dockerfile +++ b/Dockerfile @@ -220,7 +220,6 @@ ENV CGO_LDFLAGS="-L/usr/local/cuda-13/lib64 -L/usr/local/cuda-13/targets/x86_64- WORKDIR /go/src/github.com/ollama/ollama COPY CMakeLists.txt CMakePresets.json . COPY cmake cmake -COPY mlx mlx COPY x/mlxrunner/mlx x/mlxrunner/mlx COPY x/mlxrunner/xgrammar/native x/mlxrunner/xgrammar/native COPY go.mod go.sum . diff --git a/MLX_C_VERSION b/MLX_C_VERSION index b3511376b..3f98d19af 100644 --- a/MLX_C_VERSION +++ b/MLX_C_VERSION @@ -1 +1 @@ -fba4470b89073180056c9ea46c443051375f7399 +c74db5307cc8ce122f48d97ef951b30578674e7f \ No newline at end of file diff --git a/MLX_VERSION b/MLX_VERSION index 12232c618..58b7f7c00 100644 --- a/MLX_VERSION +++ b/MLX_VERSION @@ -1 +1 @@ -c793734eb715dbcfdb1ced58e348ec53c2d7ed85 \ No newline at end of file +37c26e5755da637255d57ea34b4879196a485301 \ No newline at end of file diff --git a/cmake/local.cmake b/cmake/local.cmake index 88be7fff3..7e9f685c5 100644 --- a/cmake/local.cmake +++ b/cmake/local.cmake @@ -169,17 +169,6 @@ if(OLLAMA_MLX_BACKENDS) list(APPEND _mlx_source_targets ollama-mlx-source) endif() - # Temporary MLX-C carry patch: regenerated bindings for force_fused and the - # thread-local compile cache, carried until they merge upstream into - # ml-explore/mlx-c. Then bump MLX_C_VERSION and delete mlx/compat/. - find_package(Git REQUIRED) - set(OLLAMA_MLX_C_COMPAT_PATCH_COMMAND - ${CMAKE_COMMAND} - -DPATCH_DIR=${CMAKE_SOURCE_DIR}/mlx/compat - -DPATCH_LABEL=mlx/compat - -P ${CMAKE_SOURCE_DIR}/cmake/apply-git-patches.cmake - CACHE INTERNAL "MLX-C carry patch") - if(DEFINED "FETCHCONTENT_SOURCE_DIR_MLX-C" AND NOT "${FETCHCONTENT_SOURCE_DIR_MLX-C}" STREQUAL "") get_filename_component(OLLAMA_MLX_C_SOURCE_DIR "${FETCHCONTENT_SOURCE_DIR_MLX-C}" ABSOLUTE BASE_DIR "${CMAKE_SOURCE_DIR}") @@ -199,9 +188,7 @@ if(OLLAMA_MLX_BACKENDS) CONFIGURE_COMMAND "" BUILD_COMMAND "" INSTALL_COMMAND "" - PATCH_COMMAND ${OLLAMA_MLX_C_COMPAT_PATCH_COMMAND} - USES_TERMINAL_DOWNLOAD TRUE - USES_TERMINAL_PATCH TRUE) + USES_TERMINAL_DOWNLOAD TRUE) list(APPEND _mlx_source_targets ollama-mlx-c-source) endif() # XGrammar has no pre-fetch: without an override each variant's build diff --git a/cmake/mlx/CMakeLists.txt b/cmake/mlx/CMakeLists.txt index f510c213d..997577d0f 100644 --- a/cmake/mlx/CMakeLists.txt +++ b/cmake/mlx/CMakeLists.txt @@ -128,8 +128,8 @@ list(APPEND MLX_RUNTIME_DIRS ${CMAKE_BINARY_DIR}/_deps/openblas-src/bin) # for windows so we include the DLL in our dependencies. list(APPEND MLX_RUNTIME_DIRS ${CMAKE_BINARY_DIR}/_deps/mlx-build/mlx/distributed/nccl/nccl_stub-prefix/src/nccl_stub-build/Release) -# Base regexes for runtime dependencies (cross-platform). -set(MLX_INCLUDE_REGEXES cublas cublasLt cudart cufft nvrtc nvrtc-builtins cudnn nccl openblas gfortran) +# Non-link-time deps stay explicit; link-time deps derive from the mlx target. +set(MLX_INCLUDE_REGEXES cublas cublasLt cudart cufft nvrtc nvrtc-builtins cudnn nccl cusolver cusparse nv[Jj]it[Ll]ink openblas gfortran) # On Windows, also include dl.dll (dlfcn-win32 POSIX emulation layer). if(WIN32) list(APPEND MLX_INCLUDE_REGEXES "^dl\\.dll$") @@ -170,6 +170,44 @@ install(FILES "${CMAKE_BINARY_DIR}/PICOJSON_LICENSE" DESTINATION ${OLLAMA_LIB_DIR} COMPONENT MLX) +get_target_property(_mlx_license_source mlx SOURCE_DIR) +if(NOT _mlx_license_source OR _mlx_license_source MATCHES "-NOTFOUND$") + message(FATAL_ERROR "MLX source directory not found for license install") +endif() +install(FILES + "${_mlx_license_source}/LICENSE" + DESTINATION ${OLLAMA_LIB_DIR} + RENAME MLX_LICENSE + COMPONENT MLX) +get_target_property(_mlxc_license_source mlxc SOURCE_DIR) +if(NOT _mlxc_license_source OR _mlxc_license_source MATCHES "-NOTFOUND$") + message(FATAL_ERROR "MLX-C source directory not found for license install") +endif() +install(FILES + "${_mlxc_license_source}/LICENSE" + DESTINATION ${OLLAMA_LIB_DIR} + RENAME MLX_C_LICENSE + COMPONENT MLX) +# Ship LICENSE/NOTICE/COPYING from every fetched MLX dependency. +install(CODE " + file(GLOB _dep_dirs + LIST_DIRECTORIES true + \"${CMAKE_BINARY_DIR}/_deps/*-src\") + foreach(_dep \${_dep_dirs}) + get_filename_component(_dep_name \${_dep} NAME) + string(REGEX REPLACE \"-src$\" \"\" _dep_name \"\${_dep_name}\") + string(TOUPPER \"\${_dep_name}\" _dep_name) + string(REGEX REPLACE \"[^A-Z0-9]\" \"_\" _dep_name \"\${_dep_name}\") + file(GLOB _lics \"\${_dep}/LICENSE*\" \"\${_dep}/COPYING*\" \"\${_dep}/NOTICE*\") + foreach(_lic \${_lics}) + get_filename_component(_lic_name \${_lic} NAME) + file(INSTALL DESTINATION \"$ENV{DESTDIR}\${CMAKE_INSTALL_PREFIX}/${OLLAMA_LIB_DIR}\" + TYPE FILE FILES \"\${_lic}\" + RENAME \"\${_dep_name}_\${_lic_name}\") + endforeach() + endforeach() +" COMPONENT MLX) + install(RUNTIME_DEPENDENCY_SET mlx_runtime_deps DIRECTORIES ${MLX_RUNTIME_DIRS} PRE_INCLUDE_REGEXES ${MLX_INCLUDE_REGEXES} @@ -376,8 +414,7 @@ if(WIN32 AND TARGET dl) COMPONENT MLX) endif() -# Manually install CUDA runtime libraries that MLX loads via dlopen -# (not detected by RUNTIME_DEPENDENCIES since they aren't link-time deps). +# dlopen'd runtime libs, derived from MLX's dynamic.c/delayload.cpp if(CUDAToolkit_FOUND) file(GLOB MLX_CUDA_LIBS "${CUDAToolkit_LIBRARY_DIR}/libcudart.so*" @@ -386,11 +423,20 @@ if(CUDAToolkit_FOUND) "${CUDAToolkit_LIBRARY_DIR}/libnvrtc.so*" "${CUDAToolkit_LIBRARY_DIR}/libnvrtc-builtins.so*" "${CUDAToolkit_LIBRARY_DIR}/libcufft.so*" + "${CUDAToolkit_LIBRARY_DIR}/libcusolver.so*" + "${CUDAToolkit_LIBRARY_DIR}/libcusparse.so*" + "${CUDAToolkit_LIBRARY_DIR}/libnvJitLink.so*" "${CUDAToolkit_LIBRARY_DIR}/libcudnn*.so*") if(WIN32) file(GLOB MLX_CUDA_DLLS "${CUDAToolkit_BIN_DIR}/nvrtc-builtins64_*.dll" - "${CUDAToolkit_BIN_DIR}/x64/nvrtc-builtins64_*.dll") + "${CUDAToolkit_BIN_DIR}/x64/nvrtc-builtins64_*.dll" + "${CUDAToolkit_BIN_DIR}/cusolver64_*.dll" + "${CUDAToolkit_BIN_DIR}/x64/cusolver64_*.dll" + "${CUDAToolkit_BIN_DIR}/cusparse64_*.dll" + "${CUDAToolkit_BIN_DIR}/x64/cusparse64_*.dll" + "${CUDAToolkit_BIN_DIR}/nvJitLink_*.dll" + "${CUDAToolkit_BIN_DIR}/x64/nvJitLink_*.dll") list(APPEND MLX_CUDA_LIBS ${MLX_CUDA_DLLS}) endif() find_library(MLX_CUDNN_LIBRARY NAMES cudnn HINTS "$ENV{CUDNN_LIBRARY_PATH}") diff --git a/mlx/compat/0001-mlx-c-regen-0.32.1.patch b/mlx/compat/0001-mlx-c-regen-0.32.1.patch deleted file mode 100644 index d97813685..000000000 --- a/mlx/compat/0001-mlx-c-regen-0.32.1.patch +++ /dev/null @@ -1,88 +0,0 @@ -diff --git a/mlx/c/compile.cpp b/mlx/c/compile.cpp -index bc7b864..260ee86 100644 ---- a/mlx/c/compile.cpp -+++ b/mlx/c/compile.cpp -@@ -42,22 +42,25 @@ extern "C" int mlx_detail_compile( - } - extern "C" int mlx_detail_compile_clear_cache(void) { - try { -- mlx::core::detail::compile_clear_cache(); -+ mlx::core::detail::compile_clear_cache(mlx::core::detail::compile_cache()); - } catch (std::exception& e) { - mlx_error(e.what()); - return 1; - } - return 0; - } -+ - extern "C" int mlx_detail_compile_erase(uintptr_t fun_id) { - try { -- mlx::core::detail::compile_erase(fun_id); -+ mlx::core::detail::compile_erase( -+ mlx::core::detail::compile_cache(), fun_id); - } catch (std::exception& e) { - mlx_error(e.what()); - return 1; - } - return 0; - } -+ - extern "C" int mlx_disable_compile(void) { - try { - mlx::core::disable_compile(); -diff --git a/mlx/c/compile.h b/mlx/c/compile.h -index 2892a1f..ff3abc4 100644 ---- a/mlx/c/compile.h -+++ b/mlx/c/compile.h -@@ -44,7 +44,9 @@ int mlx_detail_compile( - const uint64_t* constants, - size_t constants_num); - int mlx_detail_compile_clear_cache(void); -+ - int mlx_detail_compile_erase(uintptr_t fun_id); -+ - int mlx_disable_compile(void); - int mlx_enable_compile(void); - int mlx_set_compile_mode(mlx_compile_mode mode); -diff --git a/mlx/c/fast.cpp b/mlx/c/fast.cpp -index ac17d54..eb35d16 100644 ---- a/mlx/c/fast.cpp -+++ b/mlx/c/fast.cpp -@@ -3,8 +3,8 @@ - /* This file is auto-generated. Do not edit manually. */ - /* */ - --#include "mlx/c/fast.h" - #include "mlx/c/error.h" -+#include "mlx/c/fast.h" - #include "mlx/c/private/mlx.h" - #include "mlx/fast.h" - -@@ -610,6 +610,7 @@ extern "C" int mlx_fast_scaled_dot_product_attention( - const char* mask_mode, - const mlx_array mask_arr /* may be null */, - const mlx_array sinks /* may be null */, -+ bool force_fused, - const mlx_stream s) { - try { - mlx_array_set_( -@@ -624,6 +625,7 @@ extern "C" int mlx_fast_scaled_dot_product_attention( - : std::nullopt), - (sinks.ctx ? std::make_optional(mlx_array_get_(sinks)) - : std::nullopt), -+ force_fused, - mlx_stream_get_(s))); - } catch (std::exception& e) { - mlx_error(e.what()); -diff --git a/mlx/c/fast.h b/mlx/c/fast.h -index c825d00..e6c5bdf 100644 ---- a/mlx/c/fast.h -+++ b/mlx/c/fast.h -@@ -195,6 +195,7 @@ int mlx_fast_scaled_dot_product_attention( - const char* mask_mode, - const mlx_array mask_arr /* may be null */, - const mlx_array sinks /* may be null */, -+ bool force_fused, - const mlx_stream s); - - /**@}*/ diff --git a/x/mlxrunner/mlx/generated.c b/x/mlxrunner/mlx/generated.c index a775149a9..500875a94 100644 --- a/x/mlxrunner/mlx/generated.c +++ b/x/mlxrunner/mlx/generated.c @@ -232,6 +232,8 @@ int (*mlx_closure_custom_vmap_apply_)( const mlx_vector_array input_0, const int* input_1, size_t input_1_num) = NULL; +mlx_compile_cache (*mlx_compile_cache_new_)() = NULL; +int (*mlx_compile_cache_free_)(mlx_compile_cache cache) = NULL; int (*mlx_compile_)(mlx_closure* res, const mlx_closure fun, bool shapeless) = NULL; int (*mlx_detail_compile_)( mlx_closure* res, @@ -240,8 +242,9 @@ int (*mlx_detail_compile_)( bool shapeless, const uint64_t* constants, size_t constants_num) = NULL; -int (*mlx_detail_compile_clear_cache_)(void) = NULL; -int (*mlx_detail_compile_erase_)(uintptr_t fun_id) = NULL; +int (*mlx_detail_compile_cache_)(mlx_compile_cache* res) = NULL; +int (*mlx_detail_compile_clear_cache_)(const mlx_compile_cache cache) = NULL; +int (*mlx_detail_compile_erase_)(const mlx_compile_cache cache, uintptr_t fun_id) = NULL; int (*mlx_disable_compile_)(void) = NULL; int (*mlx_enable_compile_)(void) = NULL; int (*mlx_set_compile_mode_)(mlx_compile_mode mode) = NULL; @@ -376,6 +379,11 @@ int (*mlx_imported_function_apply_kwargs_)( const mlx_imported_function xfunc, const mlx_vector_array args, const mlx_map_string_to_array kwargs) = NULL; +int (*mlx_fast_cross_entropy_)( + mlx_array* res, + const mlx_array logits, + const mlx_array targets, + const mlx_stream s) = NULL; mlx_fast_cuda_kernel_config (*mlx_fast_cuda_kernel_config_new_)(void) = NULL; void (*mlx_fast_cuda_kernel_config_free_)(mlx_fast_cuda_kernel_config cls) = NULL; int (*mlx_fast_cuda_kernel_config_add_output_arg_)( @@ -741,6 +749,7 @@ int (*mlx_linalg_cross_)( const mlx_array b, int axis, const mlx_stream s) = NULL; +int (*mlx_linalg_det_)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL; int (*mlx_linalg_eig_)( mlx_array* res_0, mlx_array* res_1, @@ -794,6 +803,11 @@ int (*mlx_linalg_qr_)( mlx_array* res_1, const mlx_array a, const mlx_stream s) = NULL; +int (*mlx_linalg_slogdet_)( + mlx_array* res_0, + mlx_array* res_1, + const mlx_array a, + const mlx_stream s) = NULL; int (*mlx_linalg_solve_)( mlx_array* res, const mlx_array a, @@ -865,7 +879,9 @@ int (*mlx_reset_peak_memory_)(void) = NULL; int (*mlx_set_cache_limit_)(size_t* res, size_t limit) = NULL; int (*mlx_set_memory_limit_)(size_t* res, size_t limit) = NULL; int (*mlx_set_wired_limit_)(size_t* res, size_t limit) = NULL; +int (*mlx_metal_get_metallib_path_)(mlx_string* res) = NULL; int (*mlx_metal_is_available_)(bool* res) = NULL; +int (*mlx_metal_set_metallib_path_)(const char* path) = NULL; int (*mlx_metal_start_capture_)(const char* path) = NULL; int (*mlx_metal_stop_capture_)(void) = NULL; int (*mlx_abs_)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL; @@ -1163,7 +1179,38 @@ int (*mlx_conv_transpose3d_)( int (*mlx_copy_)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL; int (*mlx_cos_)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL; int (*mlx_cosh_)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL; +int (*mlx_count_nonzero_axis_)( + mlx_array* res, + const mlx_array a, + int axis, + bool keepdims, + const mlx_stream s) = NULL; +int (*mlx_count_nonzero_axes_)( + mlx_array* res, + const mlx_array a, + const int* axes, + size_t axes_num, + bool keepdims, + const mlx_stream s) = NULL; +int (*mlx_count_nonzero_)( + mlx_array* res, + const mlx_array a, + bool keepdims, + const mlx_stream s) = NULL; +int (*mlx_cummax_axis_)( + mlx_array* res, + const mlx_array a, + int axis, + bool reverse, + bool inclusive, + const mlx_stream s) = NULL; int (*mlx_cummax_)( + mlx_array* res, + const mlx_array a, + bool reverse, + bool inclusive, + const mlx_stream s) = NULL; +int (*mlx_cummin_axis_)( mlx_array* res, const mlx_array a, int axis, @@ -1171,25 +1218,40 @@ int (*mlx_cummax_)( bool inclusive, const mlx_stream s) = NULL; int (*mlx_cummin_)( + mlx_array* res, + const mlx_array a, + bool reverse, + bool inclusive, + const mlx_stream s) = NULL; +int (*mlx_cumprod_axis_)( mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, + mlx_optional_dtype dtype, const mlx_stream s) = NULL; int (*mlx_cumprod_)( mlx_array* res, const mlx_array a, - int axis, bool reverse, bool inclusive, + mlx_optional_dtype dtype, const mlx_stream s) = NULL; -int (*mlx_cumsum_)( +int (*mlx_cumsum_axis_)( mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, + mlx_optional_dtype dtype, + const mlx_stream s) = NULL; +int (*mlx_cumsum_)( + mlx_array* res, + const mlx_array a, + bool reverse, + bool inclusive, + mlx_optional_dtype dtype, const mlx_stream s) = NULL; int (*mlx_degrees_)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL; int (*mlx_depends_)( @@ -1215,6 +1277,12 @@ int (*mlx_diagonal_)( int axis1, int axis2, const mlx_stream s) = NULL; +int (*mlx_diff_)( + mlx_array* res, + const mlx_array a, + int n, + int axis, + const mlx_stream s) = NULL; int (*mlx_divide_)( mlx_array* res, const mlx_array a, @@ -1263,6 +1331,18 @@ int (*mlx_flatten_)( int start_axis, int end_axis, const mlx_stream s) = NULL; +int (*mlx_flip_axes_)( + mlx_array* res, + const mlx_array a, + const int* axes, + size_t axes_num, + const mlx_stream s) = NULL; +int (*mlx_flip_axis_)( + mlx_array* res, + const mlx_array a, + int axis, + const mlx_stream s) = NULL; +int (*mlx_flip_)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL; int (*mlx_floor_)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL; int (*mlx_floor_divide_)( mlx_array* res, @@ -1326,6 +1406,20 @@ int (*mlx_gather_qmm_)( const char* mode, bool sorted_indices, const mlx_stream s) = NULL; +int (*mlx_gather_qqmm_)( + mlx_array* res, + const mlx_array x, + const mlx_array w, + const mlx_array scales_w /* may be null */, + const mlx_array lhs_indices /* may be null */, + const mlx_array rhs_indices /* may be null */, + mlx_optional_int group_size, + mlx_optional_int bits, + const char* mode, + const mlx_array global_scale_x /* may be null */, + const mlx_array global_scale_w /* may be null */, + bool sorted_indices, + const mlx_stream s) = NULL; int (*mlx_greater_)( mlx_array* res, const mlx_array a, @@ -1383,6 +1477,14 @@ int (*mlx_less_equal_)( const mlx_array a, const mlx_array b, const mlx_stream s) = NULL; +int (*mlx_linspace_endpoint_)( + mlx_array* res, + double start, + double stop, + int num, + bool endpoint, + mlx_dtype dtype, + const mlx_stream s) = NULL; int (*mlx_linspace_)( mlx_array* res, double start, @@ -1399,13 +1501,19 @@ int (*mlx_logaddexp_)( const mlx_array a, const mlx_array b, const mlx_stream s) = NULL; -int (*mlx_logcumsumexp_)( +int (*mlx_logcumsumexp_axis_)( mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, const mlx_stream s) = NULL; +int (*mlx_logcumsumexp_)( + mlx_array* res, + const mlx_array a, + bool reverse, + bool inclusive, + const mlx_stream s) = NULL; int (*mlx_logical_and_)( mlx_array* res, const mlx_array a, @@ -1417,6 +1525,11 @@ int (*mlx_logical_or_)( const mlx_array a, const mlx_array b, const mlx_stream s) = NULL; +int (*mlx_logical_xor_)( + mlx_array* res, + const mlx_array a, + const mlx_array b, + const mlx_stream s) = NULL; int (*mlx_logsumexp_axes_)( mlx_array* res, const mlx_array a, @@ -1487,13 +1600,24 @@ int (*mlx_mean_)( const mlx_array a, bool keepdims, const mlx_stream s) = NULL; -int (*mlx_median_)( +int (*mlx_median_axes_)( mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s) = NULL; +int (*mlx_median_axis_)( + mlx_array* res, + const mlx_array a, + int axis, + bool keepdims, + const mlx_stream s) = NULL; +int (*mlx_median_)( + mlx_array* res, + const mlx_array a, + bool keepdims, + const mlx_stream s) = NULL; int (*mlx_meshgrid_)( mlx_vector_array* res, const mlx_vector_array arrays, @@ -1561,6 +1685,11 @@ int (*mlx_ones_)( size_t shape_num, mlx_dtype dtype, const mlx_stream s) = NULL; +int (*mlx_ones_like_dtype_)( + mlx_array* res, + const mlx_array a, + mlx_dtype dtype, + const mlx_stream s) = NULL; int (*mlx_ones_like_)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL; int (*mlx_outer_)( mlx_array* res, @@ -1597,6 +1726,7 @@ int (*mlx_partition_)( const mlx_array a, int kth, const mlx_stream s) = NULL; +int (*mlx_positive_)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL; int (*mlx_power_)( mlx_array* res, const mlx_array a, @@ -1796,6 +1926,12 @@ int (*mlx_scatter_prod_single_)( const mlx_array updates, int axis, const mlx_stream s) = NULL; +int (*mlx_searchsorted_)( + mlx_array* res, + const mlx_array sorted_sequence, + const mlx_array values, + const char* side, + const mlx_stream s) = NULL; int (*mlx_segmented_mm_)( mlx_array* res, const mlx_array a, @@ -2047,7 +2183,7 @@ int (*mlx_topk_axis_)( int axis, const mlx_stream s) = NULL; int (*mlx_topk_)(mlx_array* res, const mlx_array a, int k, const mlx_stream s) = NULL; -int (*mlx_trace_)( +int (*mlx_trace_axes_)( mlx_array* res, const mlx_array a, int offset, @@ -2055,6 +2191,7 @@ int (*mlx_trace_)( int axis2, mlx_dtype dtype, const mlx_stream s) = NULL; +int (*mlx_trace_)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL; int (*mlx_transpose_axes_)( mlx_array* res, const mlx_array a, @@ -2071,6 +2208,7 @@ int (*mlx_tri_)( const mlx_stream s) = NULL; int (*mlx_tril_)(mlx_array* res, const mlx_array x, int k, const mlx_stream s) = NULL; int (*mlx_triu_)(mlx_array* res, const mlx_array x, int k, const mlx_stream s) = NULL; +int (*mlx_trunc_)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL; int (*mlx_unflatten_)( mlx_array* res, const mlx_array a, @@ -2078,6 +2216,12 @@ int (*mlx_unflatten_)( const int* shape, size_t shape_num, const mlx_stream s) = NULL; +int (*mlx_unstack_axis_)( + mlx_vector_array* res, + const mlx_array a, + int axis, + const mlx_stream s) = NULL; +int (*mlx_unstack_)(mlx_vector_array* res, const mlx_array a, const mlx_stream s) = NULL; int (*mlx_var_axes_)( mlx_array* res, const mlx_array a, @@ -2099,6 +2243,12 @@ int (*mlx_var_)( bool keepdims, int ddof, const mlx_stream s) = NULL; +int (*mlx_vecdot_)( + mlx_array* res, + const mlx_array a, + const mlx_array b, + int axis, + const mlx_stream s) = NULL; int (*mlx_view_)( mlx_array* res, const mlx_array a, @@ -2116,6 +2266,11 @@ int (*mlx_zeros_)( size_t shape_num, mlx_dtype dtype, const mlx_stream s) = NULL; +int (*mlx_zeros_like_dtype_)( + mlx_array* res, + const mlx_array a, + mlx_dtype dtype, + const mlx_stream s) = NULL; int (*mlx_zeros_like_)(mlx_array* res, const mlx_array a, const mlx_stream s) = NULL; int (*mlx_random_bernoulli_)( mlx_array* res, @@ -2247,6 +2402,7 @@ int (*mlx_random_uniform_)( const mlx_stream s) = NULL; mlx_stream (*mlx_stream_new_)(void) = NULL; mlx_stream (*mlx_stream_new_device_)(mlx_device dev) = NULL; +mlx_stream (*mlx_stream_new_thread_unsafe_)(mlx_device dev) = NULL; int (*mlx_stream_set_)(mlx_stream* stream, const mlx_stream src) = NULL; int (*mlx_stream_free_)(mlx_stream stream) = NULL; int (*mlx_stream_tostring_)(mlx_string* str, mlx_stream stream) = NULL; @@ -2491,8 +2647,11 @@ int mlx_dynamic_load_symbols(mlx_dynamic_handle handle) { CHECK_LOAD(handle, mlx_closure_custom_vmap_new_func_payload); CHECK_LOAD(handle, mlx_closure_custom_vmap_set); CHECK_LOAD(handle, mlx_closure_custom_vmap_apply); + CHECK_LOAD(handle, mlx_compile_cache_new); + CHECK_LOAD(handle, mlx_compile_cache_free); CHECK_LOAD(handle, mlx_compile); CHECK_LOAD(handle, mlx_detail_compile); + CHECK_LOAD(handle, mlx_detail_compile_cache); CHECK_LOAD(handle, mlx_detail_compile_clear_cache); CHECK_LOAD(handle, mlx_detail_compile_erase); CHECK_LOAD(handle, mlx_disable_compile); @@ -2546,6 +2705,7 @@ int mlx_dynamic_load_symbols(mlx_dynamic_handle handle) { CHECK_LOAD(handle, mlx_imported_function_free); CHECK_LOAD(handle, mlx_imported_function_apply); CHECK_LOAD(handle, mlx_imported_function_apply_kwargs); + CHECK_LOAD(handle, mlx_fast_cross_entropy); CHECK_LOAD(handle, mlx_fast_cuda_kernel_config_new); CHECK_LOAD(handle, mlx_fast_cuda_kernel_config_free); CHECK_LOAD(handle, mlx_fast_cuda_kernel_config_add_output_arg); @@ -2634,6 +2794,7 @@ int mlx_dynamic_load_symbols(mlx_dynamic_handle handle) { CHECK_LOAD(handle, mlx_linalg_cholesky); CHECK_LOAD(handle, mlx_linalg_cholesky_inv); CHECK_LOAD(handle, mlx_linalg_cross); + CHECK_LOAD(handle, mlx_linalg_det); CHECK_LOAD(handle, mlx_linalg_eig); CHECK_LOAD(handle, mlx_linalg_eigh); CHECK_LOAD(handle, mlx_linalg_eigvals); @@ -2646,6 +2807,7 @@ int mlx_dynamic_load_symbols(mlx_dynamic_handle handle) { CHECK_LOAD(handle, mlx_linalg_norm_l2); CHECK_LOAD(handle, mlx_linalg_pinv); CHECK_LOAD(handle, mlx_linalg_qr); + CHECK_LOAD(handle, mlx_linalg_slogdet); CHECK_LOAD(handle, mlx_linalg_solve); CHECK_LOAD(handle, mlx_linalg_solve_triangular); CHECK_LOAD(handle, mlx_linalg_svd); @@ -2675,7 +2837,9 @@ int mlx_dynamic_load_symbols(mlx_dynamic_handle handle) { CHECK_LOAD(handle, mlx_set_cache_limit); CHECK_LOAD(handle, mlx_set_memory_limit); CHECK_LOAD(handle, mlx_set_wired_limit); + CHECK_LOAD(handle, mlx_metal_get_metallib_path); CHECK_LOAD(handle, mlx_metal_is_available); + CHECK_LOAD(handle, mlx_metal_set_metallib_path); CHECK_LOAD(handle, mlx_metal_start_capture); CHECK_LOAD(handle, mlx_metal_stop_capture); CHECK_LOAD(handle, mlx_abs); @@ -2735,15 +2899,23 @@ int mlx_dynamic_load_symbols(mlx_dynamic_handle handle) { CHECK_LOAD(handle, mlx_copy); CHECK_LOAD(handle, mlx_cos); CHECK_LOAD(handle, mlx_cosh); + CHECK_LOAD(handle, mlx_count_nonzero_axis); + CHECK_LOAD(handle, mlx_count_nonzero_axes); + CHECK_LOAD(handle, mlx_count_nonzero); + CHECK_LOAD(handle, mlx_cummax_axis); CHECK_LOAD(handle, mlx_cummax); + CHECK_LOAD(handle, mlx_cummin_axis); CHECK_LOAD(handle, mlx_cummin); + CHECK_LOAD(handle, mlx_cumprod_axis); CHECK_LOAD(handle, mlx_cumprod); + CHECK_LOAD(handle, mlx_cumsum_axis); CHECK_LOAD(handle, mlx_cumsum); CHECK_LOAD(handle, mlx_degrees); CHECK_LOAD(handle, mlx_depends); CHECK_LOAD(handle, mlx_dequantize); CHECK_LOAD(handle, mlx_diag); CHECK_LOAD(handle, mlx_diagonal); + CHECK_LOAD(handle, mlx_diff); CHECK_LOAD(handle, mlx_divide); CHECK_LOAD(handle, mlx_divmod); CHECK_LOAD(handle, mlx_einsum); @@ -2756,6 +2928,9 @@ int mlx_dynamic_load_symbols(mlx_dynamic_handle handle) { CHECK_LOAD(handle, mlx_expm1); CHECK_LOAD(handle, mlx_eye); CHECK_LOAD(handle, mlx_flatten); + CHECK_LOAD(handle, mlx_flip_axes); + CHECK_LOAD(handle, mlx_flip_axis); + CHECK_LOAD(handle, mlx_flip); CHECK_LOAD(handle, mlx_floor); CHECK_LOAD(handle, mlx_floor_divide); CHECK_LOAD(handle, mlx_from_fp8); @@ -2765,6 +2940,7 @@ int mlx_dynamic_load_symbols(mlx_dynamic_handle handle) { CHECK_LOAD(handle, mlx_gather_single); CHECK_LOAD(handle, mlx_gather_mm); CHECK_LOAD(handle, mlx_gather_qmm); + CHECK_LOAD(handle, mlx_gather_qqmm); CHECK_LOAD(handle, mlx_greater); CHECK_LOAD(handle, mlx_greater_equal); CHECK_LOAD(handle, mlx_hadamard_transform); @@ -2783,16 +2959,19 @@ int mlx_dynamic_load_symbols(mlx_dynamic_handle handle) { CHECK_LOAD(handle, mlx_left_shift); CHECK_LOAD(handle, mlx_less); CHECK_LOAD(handle, mlx_less_equal); + CHECK_LOAD(handle, mlx_linspace_endpoint); CHECK_LOAD(handle, mlx_linspace); CHECK_LOAD(handle, mlx_log); CHECK_LOAD(handle, mlx_log10); CHECK_LOAD(handle, mlx_log1p); CHECK_LOAD(handle, mlx_log2); CHECK_LOAD(handle, mlx_logaddexp); + CHECK_LOAD(handle, mlx_logcumsumexp_axis); CHECK_LOAD(handle, mlx_logcumsumexp); CHECK_LOAD(handle, mlx_logical_and); CHECK_LOAD(handle, mlx_logical_not); CHECK_LOAD(handle, mlx_logical_or); + CHECK_LOAD(handle, mlx_logical_xor); CHECK_LOAD(handle, mlx_logsumexp_axes); CHECK_LOAD(handle, mlx_logsumexp_axis); CHECK_LOAD(handle, mlx_logsumexp); @@ -2805,6 +2984,8 @@ int mlx_dynamic_load_symbols(mlx_dynamic_handle handle) { CHECK_LOAD(handle, mlx_mean_axes); CHECK_LOAD(handle, mlx_mean_axis); CHECK_LOAD(handle, mlx_mean); + CHECK_LOAD(handle, mlx_median_axes); + CHECK_LOAD(handle, mlx_median_axis); CHECK_LOAD(handle, mlx_median); CHECK_LOAD(handle, mlx_meshgrid); CHECK_LOAD(handle, mlx_min_axes); @@ -2818,12 +2999,14 @@ int mlx_dynamic_load_symbols(mlx_dynamic_handle handle) { CHECK_LOAD(handle, mlx_not_equal); CHECK_LOAD(handle, mlx_number_of_elements); CHECK_LOAD(handle, mlx_ones); + CHECK_LOAD(handle, mlx_ones_like_dtype); CHECK_LOAD(handle, mlx_ones_like); CHECK_LOAD(handle, mlx_outer); CHECK_LOAD(handle, mlx_pad); CHECK_LOAD(handle, mlx_pad_symmetric); CHECK_LOAD(handle, mlx_partition_axis); CHECK_LOAD(handle, mlx_partition); + CHECK_LOAD(handle, mlx_positive); CHECK_LOAD(handle, mlx_power); CHECK_LOAD(handle, mlx_prod_axes); CHECK_LOAD(handle, mlx_prod_axis); @@ -2856,6 +3039,7 @@ int mlx_dynamic_load_symbols(mlx_dynamic_handle handle) { CHECK_LOAD(handle, mlx_scatter_min_single); CHECK_LOAD(handle, mlx_scatter_prod); CHECK_LOAD(handle, mlx_scatter_prod_single); + CHECK_LOAD(handle, mlx_searchsorted); CHECK_LOAD(handle, mlx_segmented_mm); CHECK_LOAD(handle, mlx_sigmoid); CHECK_LOAD(handle, mlx_sign); @@ -2903,19 +3087,25 @@ int mlx_dynamic_load_symbols(mlx_dynamic_handle handle) { CHECK_LOAD(handle, mlx_to_fp8); CHECK_LOAD(handle, mlx_topk_axis); CHECK_LOAD(handle, mlx_topk); + CHECK_LOAD(handle, mlx_trace_axes); CHECK_LOAD(handle, mlx_trace); CHECK_LOAD(handle, mlx_transpose_axes); CHECK_LOAD(handle, mlx_transpose); CHECK_LOAD(handle, mlx_tri); CHECK_LOAD(handle, mlx_tril); CHECK_LOAD(handle, mlx_triu); + CHECK_LOAD(handle, mlx_trunc); CHECK_LOAD(handle, mlx_unflatten); + CHECK_LOAD(handle, mlx_unstack_axis); + CHECK_LOAD(handle, mlx_unstack); CHECK_LOAD(handle, mlx_var_axes); CHECK_LOAD(handle, mlx_var_axis); CHECK_LOAD(handle, mlx_var); + CHECK_LOAD(handle, mlx_vecdot); CHECK_LOAD(handle, mlx_view); CHECK_LOAD(handle, mlx_where); CHECK_LOAD(handle, mlx_zeros); + CHECK_LOAD(handle, mlx_zeros_like_dtype); CHECK_LOAD(handle, mlx_zeros_like); CHECK_LOAD(handle, mlx_random_bernoulli); CHECK_LOAD(handle, mlx_random_bits); @@ -2938,6 +3128,7 @@ int mlx_dynamic_load_symbols(mlx_dynamic_handle handle) { CHECK_LOAD(handle, mlx_random_uniform); CHECK_LOAD(handle, mlx_stream_new); CHECK_LOAD(handle, mlx_stream_new_device); + CHECK_LOAD(handle, mlx_stream_new_thread_unsafe); CHECK_LOAD(handle, mlx_stream_set); CHECK_LOAD(handle, mlx_stream_free); CHECK_LOAD(handle, mlx_stream_tostring); diff --git a/x/mlxrunner/mlx/generated.h b/x/mlxrunner/mlx/generated.h index 592aede4e..c60ab7c71 100644 --- a/x/mlxrunner/mlx/generated.h +++ b/x/mlxrunner/mlx/generated.h @@ -107,8 +107,11 @@ #define mlx_closure_custom_vmap_new_func_payload mlx_closure_custom_vmap_new_func_payload_mlx_gen_orig_ #define mlx_closure_custom_vmap_set mlx_closure_custom_vmap_set_mlx_gen_orig_ #define mlx_closure_custom_vmap_apply mlx_closure_custom_vmap_apply_mlx_gen_orig_ +#define mlx_compile_cache_new mlx_compile_cache_new_mlx_gen_orig_ +#define mlx_compile_cache_free mlx_compile_cache_free_mlx_gen_orig_ #define mlx_compile mlx_compile_mlx_gen_orig_ #define mlx_detail_compile mlx_detail_compile_mlx_gen_orig_ +#define mlx_detail_compile_cache mlx_detail_compile_cache_mlx_gen_orig_ #define mlx_detail_compile_clear_cache mlx_detail_compile_clear_cache_mlx_gen_orig_ #define mlx_detail_compile_erase mlx_detail_compile_erase_mlx_gen_orig_ #define mlx_disable_compile mlx_disable_compile_mlx_gen_orig_ @@ -162,6 +165,7 @@ #define mlx_imported_function_free mlx_imported_function_free_mlx_gen_orig_ #define mlx_imported_function_apply mlx_imported_function_apply_mlx_gen_orig_ #define mlx_imported_function_apply_kwargs mlx_imported_function_apply_kwargs_mlx_gen_orig_ +#define mlx_fast_cross_entropy mlx_fast_cross_entropy_mlx_gen_orig_ #define mlx_fast_cuda_kernel_config_new mlx_fast_cuda_kernel_config_new_mlx_gen_orig_ #define mlx_fast_cuda_kernel_config_free mlx_fast_cuda_kernel_config_free_mlx_gen_orig_ #define mlx_fast_cuda_kernel_config_add_output_arg mlx_fast_cuda_kernel_config_add_output_arg_mlx_gen_orig_ @@ -250,6 +254,7 @@ #define mlx_linalg_cholesky mlx_linalg_cholesky_mlx_gen_orig_ #define mlx_linalg_cholesky_inv mlx_linalg_cholesky_inv_mlx_gen_orig_ #define mlx_linalg_cross mlx_linalg_cross_mlx_gen_orig_ +#define mlx_linalg_det mlx_linalg_det_mlx_gen_orig_ #define mlx_linalg_eig mlx_linalg_eig_mlx_gen_orig_ #define mlx_linalg_eigh mlx_linalg_eigh_mlx_gen_orig_ #define mlx_linalg_eigvals mlx_linalg_eigvals_mlx_gen_orig_ @@ -262,6 +267,7 @@ #define mlx_linalg_norm_l2 mlx_linalg_norm_l2_mlx_gen_orig_ #define mlx_linalg_pinv mlx_linalg_pinv_mlx_gen_orig_ #define mlx_linalg_qr mlx_linalg_qr_mlx_gen_orig_ +#define mlx_linalg_slogdet mlx_linalg_slogdet_mlx_gen_orig_ #define mlx_linalg_solve mlx_linalg_solve_mlx_gen_orig_ #define mlx_linalg_solve_triangular mlx_linalg_solve_triangular_mlx_gen_orig_ #define mlx_linalg_svd mlx_linalg_svd_mlx_gen_orig_ @@ -291,7 +297,9 @@ #define mlx_set_cache_limit mlx_set_cache_limit_mlx_gen_orig_ #define mlx_set_memory_limit mlx_set_memory_limit_mlx_gen_orig_ #define mlx_set_wired_limit mlx_set_wired_limit_mlx_gen_orig_ +#define mlx_metal_get_metallib_path mlx_metal_get_metallib_path_mlx_gen_orig_ #define mlx_metal_is_available mlx_metal_is_available_mlx_gen_orig_ +#define mlx_metal_set_metallib_path mlx_metal_set_metallib_path_mlx_gen_orig_ #define mlx_metal_start_capture mlx_metal_start_capture_mlx_gen_orig_ #define mlx_metal_stop_capture mlx_metal_stop_capture_mlx_gen_orig_ #define mlx_abs mlx_abs_mlx_gen_orig_ @@ -351,15 +359,23 @@ #define mlx_copy mlx_copy_mlx_gen_orig_ #define mlx_cos mlx_cos_mlx_gen_orig_ #define mlx_cosh mlx_cosh_mlx_gen_orig_ +#define mlx_count_nonzero_axis mlx_count_nonzero_axis_mlx_gen_orig_ +#define mlx_count_nonzero_axes mlx_count_nonzero_axes_mlx_gen_orig_ +#define mlx_count_nonzero mlx_count_nonzero_mlx_gen_orig_ +#define mlx_cummax_axis mlx_cummax_axis_mlx_gen_orig_ #define mlx_cummax mlx_cummax_mlx_gen_orig_ +#define mlx_cummin_axis mlx_cummin_axis_mlx_gen_orig_ #define mlx_cummin mlx_cummin_mlx_gen_orig_ +#define mlx_cumprod_axis mlx_cumprod_axis_mlx_gen_orig_ #define mlx_cumprod mlx_cumprod_mlx_gen_orig_ +#define mlx_cumsum_axis mlx_cumsum_axis_mlx_gen_orig_ #define mlx_cumsum mlx_cumsum_mlx_gen_orig_ #define mlx_degrees mlx_degrees_mlx_gen_orig_ #define mlx_depends mlx_depends_mlx_gen_orig_ #define mlx_dequantize mlx_dequantize_mlx_gen_orig_ #define mlx_diag mlx_diag_mlx_gen_orig_ #define mlx_diagonal mlx_diagonal_mlx_gen_orig_ +#define mlx_diff mlx_diff_mlx_gen_orig_ #define mlx_divide mlx_divide_mlx_gen_orig_ #define mlx_divmod mlx_divmod_mlx_gen_orig_ #define mlx_einsum mlx_einsum_mlx_gen_orig_ @@ -372,6 +388,9 @@ #define mlx_expm1 mlx_expm1_mlx_gen_orig_ #define mlx_eye mlx_eye_mlx_gen_orig_ #define mlx_flatten mlx_flatten_mlx_gen_orig_ +#define mlx_flip_axes mlx_flip_axes_mlx_gen_orig_ +#define mlx_flip_axis mlx_flip_axis_mlx_gen_orig_ +#define mlx_flip mlx_flip_mlx_gen_orig_ #define mlx_floor mlx_floor_mlx_gen_orig_ #define mlx_floor_divide mlx_floor_divide_mlx_gen_orig_ #define mlx_from_fp8 mlx_from_fp8_mlx_gen_orig_ @@ -381,6 +400,7 @@ #define mlx_gather_single mlx_gather_single_mlx_gen_orig_ #define mlx_gather_mm mlx_gather_mm_mlx_gen_orig_ #define mlx_gather_qmm mlx_gather_qmm_mlx_gen_orig_ +#define mlx_gather_qqmm mlx_gather_qqmm_mlx_gen_orig_ #define mlx_greater mlx_greater_mlx_gen_orig_ #define mlx_greater_equal mlx_greater_equal_mlx_gen_orig_ #define mlx_hadamard_transform mlx_hadamard_transform_mlx_gen_orig_ @@ -399,16 +419,19 @@ #define mlx_left_shift mlx_left_shift_mlx_gen_orig_ #define mlx_less mlx_less_mlx_gen_orig_ #define mlx_less_equal mlx_less_equal_mlx_gen_orig_ +#define mlx_linspace_endpoint mlx_linspace_endpoint_mlx_gen_orig_ #define mlx_linspace mlx_linspace_mlx_gen_orig_ #define mlx_log mlx_log_mlx_gen_orig_ #define mlx_log10 mlx_log10_mlx_gen_orig_ #define mlx_log1p mlx_log1p_mlx_gen_orig_ #define mlx_log2 mlx_log2_mlx_gen_orig_ #define mlx_logaddexp mlx_logaddexp_mlx_gen_orig_ +#define mlx_logcumsumexp_axis mlx_logcumsumexp_axis_mlx_gen_orig_ #define mlx_logcumsumexp mlx_logcumsumexp_mlx_gen_orig_ #define mlx_logical_and mlx_logical_and_mlx_gen_orig_ #define mlx_logical_not mlx_logical_not_mlx_gen_orig_ #define mlx_logical_or mlx_logical_or_mlx_gen_orig_ +#define mlx_logical_xor mlx_logical_xor_mlx_gen_orig_ #define mlx_logsumexp_axes mlx_logsumexp_axes_mlx_gen_orig_ #define mlx_logsumexp_axis mlx_logsumexp_axis_mlx_gen_orig_ #define mlx_logsumexp mlx_logsumexp_mlx_gen_orig_ @@ -421,6 +444,8 @@ #define mlx_mean_axes mlx_mean_axes_mlx_gen_orig_ #define mlx_mean_axis mlx_mean_axis_mlx_gen_orig_ #define mlx_mean mlx_mean_mlx_gen_orig_ +#define mlx_median_axes mlx_median_axes_mlx_gen_orig_ +#define mlx_median_axis mlx_median_axis_mlx_gen_orig_ #define mlx_median mlx_median_mlx_gen_orig_ #define mlx_meshgrid mlx_meshgrid_mlx_gen_orig_ #define mlx_min_axes mlx_min_axes_mlx_gen_orig_ @@ -434,12 +459,14 @@ #define mlx_not_equal mlx_not_equal_mlx_gen_orig_ #define mlx_number_of_elements mlx_number_of_elements_mlx_gen_orig_ #define mlx_ones mlx_ones_mlx_gen_orig_ +#define mlx_ones_like_dtype mlx_ones_like_dtype_mlx_gen_orig_ #define mlx_ones_like mlx_ones_like_mlx_gen_orig_ #define mlx_outer mlx_outer_mlx_gen_orig_ #define mlx_pad mlx_pad_mlx_gen_orig_ #define mlx_pad_symmetric mlx_pad_symmetric_mlx_gen_orig_ #define mlx_partition_axis mlx_partition_axis_mlx_gen_orig_ #define mlx_partition mlx_partition_mlx_gen_orig_ +#define mlx_positive mlx_positive_mlx_gen_orig_ #define mlx_power mlx_power_mlx_gen_orig_ #define mlx_prod_axes mlx_prod_axes_mlx_gen_orig_ #define mlx_prod_axis mlx_prod_axis_mlx_gen_orig_ @@ -472,6 +499,7 @@ #define mlx_scatter_min_single mlx_scatter_min_single_mlx_gen_orig_ #define mlx_scatter_prod mlx_scatter_prod_mlx_gen_orig_ #define mlx_scatter_prod_single mlx_scatter_prod_single_mlx_gen_orig_ +#define mlx_searchsorted mlx_searchsorted_mlx_gen_orig_ #define mlx_segmented_mm mlx_segmented_mm_mlx_gen_orig_ #define mlx_sigmoid mlx_sigmoid_mlx_gen_orig_ #define mlx_sign mlx_sign_mlx_gen_orig_ @@ -519,19 +547,25 @@ #define mlx_to_fp8 mlx_to_fp8_mlx_gen_orig_ #define mlx_topk_axis mlx_topk_axis_mlx_gen_orig_ #define mlx_topk mlx_topk_mlx_gen_orig_ +#define mlx_trace_axes mlx_trace_axes_mlx_gen_orig_ #define mlx_trace mlx_trace_mlx_gen_orig_ #define mlx_transpose_axes mlx_transpose_axes_mlx_gen_orig_ #define mlx_transpose mlx_transpose_mlx_gen_orig_ #define mlx_tri mlx_tri_mlx_gen_orig_ #define mlx_tril mlx_tril_mlx_gen_orig_ #define mlx_triu mlx_triu_mlx_gen_orig_ +#define mlx_trunc mlx_trunc_mlx_gen_orig_ #define mlx_unflatten mlx_unflatten_mlx_gen_orig_ +#define mlx_unstack_axis mlx_unstack_axis_mlx_gen_orig_ +#define mlx_unstack mlx_unstack_mlx_gen_orig_ #define mlx_var_axes mlx_var_axes_mlx_gen_orig_ #define mlx_var_axis mlx_var_axis_mlx_gen_orig_ #define mlx_var mlx_var_mlx_gen_orig_ +#define mlx_vecdot mlx_vecdot_mlx_gen_orig_ #define mlx_view mlx_view_mlx_gen_orig_ #define mlx_where mlx_where_mlx_gen_orig_ #define mlx_zeros mlx_zeros_mlx_gen_orig_ +#define mlx_zeros_like_dtype mlx_zeros_like_dtype_mlx_gen_orig_ #define mlx_zeros_like mlx_zeros_like_mlx_gen_orig_ #define mlx_random_bernoulli mlx_random_bernoulli_mlx_gen_orig_ #define mlx_random_bits mlx_random_bits_mlx_gen_orig_ @@ -554,6 +588,7 @@ #define mlx_random_uniform mlx_random_uniform_mlx_gen_orig_ #define mlx_stream_new mlx_stream_new_mlx_gen_orig_ #define mlx_stream_new_device mlx_stream_new_device_mlx_gen_orig_ +#define mlx_stream_new_thread_unsafe mlx_stream_new_thread_unsafe_mlx_gen_orig_ #define mlx_stream_set mlx_stream_set_mlx_gen_orig_ #define mlx_stream_free mlx_stream_free_mlx_gen_orig_ #define mlx_stream_tostring mlx_stream_tostring_mlx_gen_orig_ @@ -730,8 +765,11 @@ #undef mlx_closure_custom_vmap_new_func_payload #undef mlx_closure_custom_vmap_set #undef mlx_closure_custom_vmap_apply +#undef mlx_compile_cache_new +#undef mlx_compile_cache_free #undef mlx_compile #undef mlx_detail_compile +#undef mlx_detail_compile_cache #undef mlx_detail_compile_clear_cache #undef mlx_detail_compile_erase #undef mlx_disable_compile @@ -785,6 +823,7 @@ #undef mlx_imported_function_free #undef mlx_imported_function_apply #undef mlx_imported_function_apply_kwargs +#undef mlx_fast_cross_entropy #undef mlx_fast_cuda_kernel_config_new #undef mlx_fast_cuda_kernel_config_free #undef mlx_fast_cuda_kernel_config_add_output_arg @@ -873,6 +912,7 @@ #undef mlx_linalg_cholesky #undef mlx_linalg_cholesky_inv #undef mlx_linalg_cross +#undef mlx_linalg_det #undef mlx_linalg_eig #undef mlx_linalg_eigh #undef mlx_linalg_eigvals @@ -885,6 +925,7 @@ #undef mlx_linalg_norm_l2 #undef mlx_linalg_pinv #undef mlx_linalg_qr +#undef mlx_linalg_slogdet #undef mlx_linalg_solve #undef mlx_linalg_solve_triangular #undef mlx_linalg_svd @@ -914,7 +955,9 @@ #undef mlx_set_cache_limit #undef mlx_set_memory_limit #undef mlx_set_wired_limit +#undef mlx_metal_get_metallib_path #undef mlx_metal_is_available +#undef mlx_metal_set_metallib_path #undef mlx_metal_start_capture #undef mlx_metal_stop_capture #undef mlx_abs @@ -974,15 +1017,23 @@ #undef mlx_copy #undef mlx_cos #undef mlx_cosh +#undef mlx_count_nonzero_axis +#undef mlx_count_nonzero_axes +#undef mlx_count_nonzero +#undef mlx_cummax_axis #undef mlx_cummax +#undef mlx_cummin_axis #undef mlx_cummin +#undef mlx_cumprod_axis #undef mlx_cumprod +#undef mlx_cumsum_axis #undef mlx_cumsum #undef mlx_degrees #undef mlx_depends #undef mlx_dequantize #undef mlx_diag #undef mlx_diagonal +#undef mlx_diff #undef mlx_divide #undef mlx_divmod #undef mlx_einsum @@ -995,6 +1046,9 @@ #undef mlx_expm1 #undef mlx_eye #undef mlx_flatten +#undef mlx_flip_axes +#undef mlx_flip_axis +#undef mlx_flip #undef mlx_floor #undef mlx_floor_divide #undef mlx_from_fp8 @@ -1004,6 +1058,7 @@ #undef mlx_gather_single #undef mlx_gather_mm #undef mlx_gather_qmm +#undef mlx_gather_qqmm #undef mlx_greater #undef mlx_greater_equal #undef mlx_hadamard_transform @@ -1022,16 +1077,19 @@ #undef mlx_left_shift #undef mlx_less #undef mlx_less_equal +#undef mlx_linspace_endpoint #undef mlx_linspace #undef mlx_log #undef mlx_log10 #undef mlx_log1p #undef mlx_log2 #undef mlx_logaddexp +#undef mlx_logcumsumexp_axis #undef mlx_logcumsumexp #undef mlx_logical_and #undef mlx_logical_not #undef mlx_logical_or +#undef mlx_logical_xor #undef mlx_logsumexp_axes #undef mlx_logsumexp_axis #undef mlx_logsumexp @@ -1044,6 +1102,8 @@ #undef mlx_mean_axes #undef mlx_mean_axis #undef mlx_mean +#undef mlx_median_axes +#undef mlx_median_axis #undef mlx_median #undef mlx_meshgrid #undef mlx_min_axes @@ -1057,12 +1117,14 @@ #undef mlx_not_equal #undef mlx_number_of_elements #undef mlx_ones +#undef mlx_ones_like_dtype #undef mlx_ones_like #undef mlx_outer #undef mlx_pad #undef mlx_pad_symmetric #undef mlx_partition_axis #undef mlx_partition +#undef mlx_positive #undef mlx_power #undef mlx_prod_axes #undef mlx_prod_axis @@ -1095,6 +1157,7 @@ #undef mlx_scatter_min_single #undef mlx_scatter_prod #undef mlx_scatter_prod_single +#undef mlx_searchsorted #undef mlx_segmented_mm #undef mlx_sigmoid #undef mlx_sign @@ -1142,19 +1205,25 @@ #undef mlx_to_fp8 #undef mlx_topk_axis #undef mlx_topk +#undef mlx_trace_axes #undef mlx_trace #undef mlx_transpose_axes #undef mlx_transpose #undef mlx_tri #undef mlx_tril #undef mlx_triu +#undef mlx_trunc #undef mlx_unflatten +#undef mlx_unstack_axis +#undef mlx_unstack #undef mlx_var_axes #undef mlx_var_axis #undef mlx_var +#undef mlx_vecdot #undef mlx_view #undef mlx_where #undef mlx_zeros +#undef mlx_zeros_like_dtype #undef mlx_zeros_like #undef mlx_random_bernoulli #undef mlx_random_bits @@ -1177,6 +1246,7 @@ #undef mlx_random_uniform #undef mlx_stream_new #undef mlx_stream_new_device +#undef mlx_stream_new_thread_unsafe #undef mlx_stream_set #undef mlx_stream_free #undef mlx_stream_tostring @@ -1475,6 +1545,8 @@ extern int (*mlx_closure_custom_vmap_apply_)( const mlx_vector_array input_0, const int* input_1, size_t input_1_num); +extern mlx_compile_cache (*mlx_compile_cache_new_)(); +extern int (*mlx_compile_cache_free_)(mlx_compile_cache cache); extern int (*mlx_compile_)(mlx_closure* res, const mlx_closure fun, bool shapeless); extern int (*mlx_detail_compile_)( mlx_closure* res, @@ -1483,8 +1555,9 @@ extern int (*mlx_detail_compile_)( bool shapeless, const uint64_t* constants, size_t constants_num); -extern int (*mlx_detail_compile_clear_cache_)(void); -extern int (*mlx_detail_compile_erase_)(uintptr_t fun_id); +extern int (*mlx_detail_compile_cache_)(mlx_compile_cache* res); +extern int (*mlx_detail_compile_clear_cache_)(const mlx_compile_cache cache); +extern int (*mlx_detail_compile_erase_)(const mlx_compile_cache cache, uintptr_t fun_id); extern int (*mlx_disable_compile_)(void); extern int (*mlx_enable_compile_)(void); extern int (*mlx_set_compile_mode_)(mlx_compile_mode mode); @@ -1619,6 +1692,11 @@ extern int (*mlx_imported_function_apply_kwargs_)( const mlx_imported_function xfunc, const mlx_vector_array args, const mlx_map_string_to_array kwargs); +extern int (*mlx_fast_cross_entropy_)( + mlx_array* res, + const mlx_array logits, + const mlx_array targets, + const mlx_stream s); extern mlx_fast_cuda_kernel_config (*mlx_fast_cuda_kernel_config_new_)(void); extern void (*mlx_fast_cuda_kernel_config_free_)(mlx_fast_cuda_kernel_config cls); extern int (*mlx_fast_cuda_kernel_config_add_output_arg_)( @@ -1984,6 +2062,7 @@ extern int (*mlx_linalg_cross_)( const mlx_array b, int axis, const mlx_stream s); +extern int (*mlx_linalg_det_)(mlx_array* res, const mlx_array a, const mlx_stream s); extern int (*mlx_linalg_eig_)( mlx_array* res_0, mlx_array* res_1, @@ -2037,6 +2116,11 @@ extern int (*mlx_linalg_qr_)( mlx_array* res_1, const mlx_array a, const mlx_stream s); +extern int (*mlx_linalg_slogdet_)( + mlx_array* res_0, + mlx_array* res_1, + const mlx_array a, + const mlx_stream s); extern int (*mlx_linalg_solve_)( mlx_array* res, const mlx_array a, @@ -2108,7 +2192,9 @@ extern int (*mlx_reset_peak_memory_)(void); extern int (*mlx_set_cache_limit_)(size_t* res, size_t limit); extern int (*mlx_set_memory_limit_)(size_t* res, size_t limit); extern int (*mlx_set_wired_limit_)(size_t* res, size_t limit); +extern int (*mlx_metal_get_metallib_path_)(mlx_string* res); extern int (*mlx_metal_is_available_)(bool* res); +extern int (*mlx_metal_set_metallib_path_)(const char* path); extern int (*mlx_metal_start_capture_)(const char* path); extern int (*mlx_metal_stop_capture_)(void); extern int (*mlx_abs_)(mlx_array* res, const mlx_array a, const mlx_stream s); @@ -2406,7 +2492,38 @@ extern int (*mlx_conv_transpose3d_)( extern int (*mlx_copy_)(mlx_array* res, const mlx_array a, const mlx_stream s); extern int (*mlx_cos_)(mlx_array* res, const mlx_array a, const mlx_stream s); extern int (*mlx_cosh_)(mlx_array* res, const mlx_array a, const mlx_stream s); +extern int (*mlx_count_nonzero_axis_)( + mlx_array* res, + const mlx_array a, + int axis, + bool keepdims, + const mlx_stream s); +extern int (*mlx_count_nonzero_axes_)( + mlx_array* res, + const mlx_array a, + const int* axes, + size_t axes_num, + bool keepdims, + const mlx_stream s); +extern int (*mlx_count_nonzero_)( + mlx_array* res, + const mlx_array a, + bool keepdims, + const mlx_stream s); +extern int (*mlx_cummax_axis_)( + mlx_array* res, + const mlx_array a, + int axis, + bool reverse, + bool inclusive, + const mlx_stream s); extern int (*mlx_cummax_)( + mlx_array* res, + const mlx_array a, + bool reverse, + bool inclusive, + const mlx_stream s); +extern int (*mlx_cummin_axis_)( mlx_array* res, const mlx_array a, int axis, @@ -2414,25 +2531,40 @@ extern int (*mlx_cummax_)( bool inclusive, const mlx_stream s); extern int (*mlx_cummin_)( + mlx_array* res, + const mlx_array a, + bool reverse, + bool inclusive, + const mlx_stream s); +extern int (*mlx_cumprod_axis_)( mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, + mlx_optional_dtype dtype, const mlx_stream s); extern int (*mlx_cumprod_)( mlx_array* res, const mlx_array a, - int axis, bool reverse, bool inclusive, + mlx_optional_dtype dtype, const mlx_stream s); -extern int (*mlx_cumsum_)( +extern int (*mlx_cumsum_axis_)( mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, + mlx_optional_dtype dtype, + const mlx_stream s); +extern int (*mlx_cumsum_)( + mlx_array* res, + const mlx_array a, + bool reverse, + bool inclusive, + mlx_optional_dtype dtype, const mlx_stream s); extern int (*mlx_degrees_)(mlx_array* res, const mlx_array a, const mlx_stream s); extern int (*mlx_depends_)( @@ -2458,6 +2590,12 @@ extern int (*mlx_diagonal_)( int axis1, int axis2, const mlx_stream s); +extern int (*mlx_diff_)( + mlx_array* res, + const mlx_array a, + int n, + int axis, + const mlx_stream s); extern int (*mlx_divide_)( mlx_array* res, const mlx_array a, @@ -2506,6 +2644,18 @@ extern int (*mlx_flatten_)( int start_axis, int end_axis, const mlx_stream s); +extern int (*mlx_flip_axes_)( + mlx_array* res, + const mlx_array a, + const int* axes, + size_t axes_num, + const mlx_stream s); +extern int (*mlx_flip_axis_)( + mlx_array* res, + const mlx_array a, + int axis, + const mlx_stream s); +extern int (*mlx_flip_)(mlx_array* res, const mlx_array a, const mlx_stream s); extern int (*mlx_floor_)(mlx_array* res, const mlx_array a, const mlx_stream s); extern int (*mlx_floor_divide_)( mlx_array* res, @@ -2569,6 +2719,20 @@ extern int (*mlx_gather_qmm_)( const char* mode, bool sorted_indices, const mlx_stream s); +extern int (*mlx_gather_qqmm_)( + mlx_array* res, + const mlx_array x, + const mlx_array w, + const mlx_array scales_w /* may be null */, + const mlx_array lhs_indices /* may be null */, + const mlx_array rhs_indices /* may be null */, + mlx_optional_int group_size, + mlx_optional_int bits, + const char* mode, + const mlx_array global_scale_x /* may be null */, + const mlx_array global_scale_w /* may be null */, + bool sorted_indices, + const mlx_stream s); extern int (*mlx_greater_)( mlx_array* res, const mlx_array a, @@ -2626,6 +2790,14 @@ extern int (*mlx_less_equal_)( const mlx_array a, const mlx_array b, const mlx_stream s); +extern int (*mlx_linspace_endpoint_)( + mlx_array* res, + double start, + double stop, + int num, + bool endpoint, + mlx_dtype dtype, + const mlx_stream s); extern int (*mlx_linspace_)( mlx_array* res, double start, @@ -2642,13 +2814,19 @@ extern int (*mlx_logaddexp_)( const mlx_array a, const mlx_array b, const mlx_stream s); -extern int (*mlx_logcumsumexp_)( +extern int (*mlx_logcumsumexp_axis_)( mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, const mlx_stream s); +extern int (*mlx_logcumsumexp_)( + mlx_array* res, + const mlx_array a, + bool reverse, + bool inclusive, + const mlx_stream s); extern int (*mlx_logical_and_)( mlx_array* res, const mlx_array a, @@ -2660,6 +2838,11 @@ extern int (*mlx_logical_or_)( const mlx_array a, const mlx_array b, const mlx_stream s); +extern int (*mlx_logical_xor_)( + mlx_array* res, + const mlx_array a, + const mlx_array b, + const mlx_stream s); extern int (*mlx_logsumexp_axes_)( mlx_array* res, const mlx_array a, @@ -2730,13 +2913,24 @@ extern int (*mlx_mean_)( const mlx_array a, bool keepdims, const mlx_stream s); -extern int (*mlx_median_)( +extern int (*mlx_median_axes_)( mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s); +extern int (*mlx_median_axis_)( + mlx_array* res, + const mlx_array a, + int axis, + bool keepdims, + const mlx_stream s); +extern int (*mlx_median_)( + mlx_array* res, + const mlx_array a, + bool keepdims, + const mlx_stream s); extern int (*mlx_meshgrid_)( mlx_vector_array* res, const mlx_vector_array arrays, @@ -2804,6 +2998,11 @@ extern int (*mlx_ones_)( size_t shape_num, mlx_dtype dtype, const mlx_stream s); +extern int (*mlx_ones_like_dtype_)( + mlx_array* res, + const mlx_array a, + mlx_dtype dtype, + const mlx_stream s); extern int (*mlx_ones_like_)(mlx_array* res, const mlx_array a, const mlx_stream s); extern int (*mlx_outer_)( mlx_array* res, @@ -2840,6 +3039,7 @@ extern int (*mlx_partition_)( const mlx_array a, int kth, const mlx_stream s); +extern int (*mlx_positive_)(mlx_array* res, const mlx_array a, const mlx_stream s); extern int (*mlx_power_)( mlx_array* res, const mlx_array a, @@ -3039,6 +3239,12 @@ extern int (*mlx_scatter_prod_single_)( const mlx_array updates, int axis, const mlx_stream s); +extern int (*mlx_searchsorted_)( + mlx_array* res, + const mlx_array sorted_sequence, + const mlx_array values, + const char* side, + const mlx_stream s); extern int (*mlx_segmented_mm_)( mlx_array* res, const mlx_array a, @@ -3290,7 +3496,7 @@ extern int (*mlx_topk_axis_)( int axis, const mlx_stream s); extern int (*mlx_topk_)(mlx_array* res, const mlx_array a, int k, const mlx_stream s); -extern int (*mlx_trace_)( +extern int (*mlx_trace_axes_)( mlx_array* res, const mlx_array a, int offset, @@ -3298,6 +3504,7 @@ extern int (*mlx_trace_)( int axis2, mlx_dtype dtype, const mlx_stream s); +extern int (*mlx_trace_)(mlx_array* res, const mlx_array a, const mlx_stream s); extern int (*mlx_transpose_axes_)( mlx_array* res, const mlx_array a, @@ -3314,6 +3521,7 @@ extern int (*mlx_tri_)( const mlx_stream s); extern int (*mlx_tril_)(mlx_array* res, const mlx_array x, int k, const mlx_stream s); extern int (*mlx_triu_)(mlx_array* res, const mlx_array x, int k, const mlx_stream s); +extern int (*mlx_trunc_)(mlx_array* res, const mlx_array a, const mlx_stream s); extern int (*mlx_unflatten_)( mlx_array* res, const mlx_array a, @@ -3321,6 +3529,12 @@ extern int (*mlx_unflatten_)( const int* shape, size_t shape_num, const mlx_stream s); +extern int (*mlx_unstack_axis_)( + mlx_vector_array* res, + const mlx_array a, + int axis, + const mlx_stream s); +extern int (*mlx_unstack_)(mlx_vector_array* res, const mlx_array a, const mlx_stream s); extern int (*mlx_var_axes_)( mlx_array* res, const mlx_array a, @@ -3342,6 +3556,12 @@ extern int (*mlx_var_)( bool keepdims, int ddof, const mlx_stream s); +extern int (*mlx_vecdot_)( + mlx_array* res, + const mlx_array a, + const mlx_array b, + int axis, + const mlx_stream s); extern int (*mlx_view_)( mlx_array* res, const mlx_array a, @@ -3359,6 +3579,11 @@ extern int (*mlx_zeros_)( size_t shape_num, mlx_dtype dtype, const mlx_stream s); +extern int (*mlx_zeros_like_dtype_)( + mlx_array* res, + const mlx_array a, + mlx_dtype dtype, + const mlx_stream s); extern int (*mlx_zeros_like_)(mlx_array* res, const mlx_array a, const mlx_stream s); extern int (*mlx_random_bernoulli_)( mlx_array* res, @@ -3490,6 +3715,7 @@ extern int (*mlx_random_uniform_)( const mlx_stream s); extern mlx_stream (*mlx_stream_new_)(void); extern mlx_stream (*mlx_stream_new_device_)(mlx_device dev); +extern mlx_stream (*mlx_stream_new_thread_unsafe_)(mlx_device dev); extern int (*mlx_stream_set_)(mlx_stream* stream, const mlx_stream src); extern int (*mlx_stream_free_)(mlx_stream stream); extern int (*mlx_stream_tostring_)(mlx_string* str, mlx_stream stream); @@ -4063,6 +4289,12 @@ static inline int mlx_closure_custom_vmap_apply( size_t input_1_num) { return mlx_closure_custom_vmap_apply_(res_0, res_1, cls, input_0, input_1, input_1_num); } +static inline mlx_compile_cache mlx_compile_cache_new() { + return mlx_compile_cache_new_(); +} +static inline int mlx_compile_cache_free(mlx_compile_cache cache) { + return mlx_compile_cache_free_(cache); +} static inline int mlx_compile(mlx_closure* res, const mlx_closure fun, bool shapeless) { return mlx_compile_(res, fun, shapeless); } @@ -4075,11 +4307,14 @@ static inline int mlx_detail_compile( size_t constants_num) { return mlx_detail_compile_(res, fun, fun_id, shapeless, constants, constants_num); } -static inline int mlx_detail_compile_clear_cache(void) { - return mlx_detail_compile_clear_cache_(); +static inline int mlx_detail_compile_cache(mlx_compile_cache* res) { + return mlx_detail_compile_cache_(res); } -static inline int mlx_detail_compile_erase(uintptr_t fun_id) { - return mlx_detail_compile_erase_(fun_id); +static inline int mlx_detail_compile_clear_cache(const mlx_compile_cache cache) { + return mlx_detail_compile_clear_cache_(cache); +} +static inline int mlx_detail_compile_erase(const mlx_compile_cache cache, uintptr_t fun_id) { + return mlx_detail_compile_erase_(cache, fun_id); } static inline int mlx_disable_compile(void) { return mlx_disable_compile_(); @@ -4317,6 +4552,13 @@ static inline int mlx_imported_function_apply_kwargs( const mlx_map_string_to_array kwargs) { return mlx_imported_function_apply_kwargs_(res, xfunc, args, kwargs); } +static inline int mlx_fast_cross_entropy( + mlx_array* res, + const mlx_array logits, + const mlx_array targets, + const mlx_stream s) { + return mlx_fast_cross_entropy_(res, logits, targets, s); +} static inline mlx_fast_cuda_kernel_config mlx_fast_cuda_kernel_config_new(void) { return mlx_fast_cuda_kernel_config_new_(); } @@ -4858,6 +5100,9 @@ static inline int mlx_linalg_cross( const mlx_stream s) { return mlx_linalg_cross_(res, a, b, axis, s); } +static inline int mlx_linalg_det(mlx_array* res, const mlx_array a, const mlx_stream s) { + return mlx_linalg_det_(res, a, s); +} static inline int mlx_linalg_eig( mlx_array* res_0, mlx_array* res_1, @@ -4935,6 +5180,13 @@ static inline int mlx_linalg_qr( const mlx_stream s) { return mlx_linalg_qr_(res_0, res_1, a, s); } +static inline int mlx_linalg_slogdet( + mlx_array* res_0, + mlx_array* res_1, + const mlx_array a, + const mlx_stream s) { + return mlx_linalg_slogdet_(res_0, res_1, a, s); +} static inline int mlx_linalg_solve( mlx_array* res, const mlx_array a, @@ -5064,9 +5316,15 @@ static inline int mlx_set_memory_limit(size_t* res, size_t limit) { static inline int mlx_set_wired_limit(size_t* res, size_t limit) { return mlx_set_wired_limit_(res, limit); } +static inline int mlx_metal_get_metallib_path(mlx_string* res) { + return mlx_metal_get_metallib_path_(res); +} static inline int mlx_metal_is_available(bool* res) { return mlx_metal_is_available_(res); } +static inline int mlx_metal_set_metallib_path(const char* path) { + return mlx_metal_set_metallib_path_(path); +} static inline int mlx_metal_start_capture(const char* path) { return mlx_metal_start_capture_(path); } @@ -5482,41 +5740,101 @@ static inline int mlx_cos(mlx_array* res, const mlx_array a, const mlx_stream s) static inline int mlx_cosh(mlx_array* res, const mlx_array a, const mlx_stream s) { return mlx_cosh_(res, a, s); } -static inline int mlx_cummax( +static inline int mlx_count_nonzero_axis( + mlx_array* res, + const mlx_array a, + int axis, + bool keepdims, + const mlx_stream s) { + return mlx_count_nonzero_axis_(res, a, axis, keepdims, s); +} +static inline int mlx_count_nonzero_axes( + mlx_array* res, + const mlx_array a, + const int* axes, + size_t axes_num, + bool keepdims, + const mlx_stream s) { + return mlx_count_nonzero_axes_(res, a, axes, axes_num, keepdims, s); +} +static inline int mlx_count_nonzero( + mlx_array* res, + const mlx_array a, + bool keepdims, + const mlx_stream s) { + return mlx_count_nonzero_(res, a, keepdims, s); +} +static inline int mlx_cummax_axis( mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, const mlx_stream s) { - return mlx_cummax_(res, a, axis, reverse, inclusive, s); + return mlx_cummax_axis_(res, a, axis, reverse, inclusive, s); +} +static inline int mlx_cummax( + mlx_array* res, + const mlx_array a, + bool reverse, + bool inclusive, + const mlx_stream s) { + return mlx_cummax_(res, a, reverse, inclusive, s); +} +static inline int mlx_cummin_axis( + mlx_array* res, + const mlx_array a, + int axis, + bool reverse, + bool inclusive, + const mlx_stream s) { + return mlx_cummin_axis_(res, a, axis, reverse, inclusive, s); } static inline int mlx_cummin( mlx_array* res, const mlx_array a, - int axis, bool reverse, bool inclusive, const mlx_stream s) { - return mlx_cummin_(res, a, axis, reverse, inclusive, s); + return mlx_cummin_(res, a, reverse, inclusive, s); +} +static inline int mlx_cumprod_axis( + mlx_array* res, + const mlx_array a, + int axis, + bool reverse, + bool inclusive, + mlx_optional_dtype dtype, + const mlx_stream s) { + return mlx_cumprod_axis_(res, a, axis, reverse, inclusive, dtype, s); } static inline int mlx_cumprod( mlx_array* res, const mlx_array a, - int axis, bool reverse, bool inclusive, + mlx_optional_dtype dtype, const mlx_stream s) { - return mlx_cumprod_(res, a, axis, reverse, inclusive, s); + return mlx_cumprod_(res, a, reverse, inclusive, dtype, s); } -static inline int mlx_cumsum( +static inline int mlx_cumsum_axis( mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, + mlx_optional_dtype dtype, const mlx_stream s) { - return mlx_cumsum_(res, a, axis, reverse, inclusive, s); + return mlx_cumsum_axis_(res, a, axis, reverse, inclusive, dtype, s); +} +static inline int mlx_cumsum( + mlx_array* res, + const mlx_array a, + bool reverse, + bool inclusive, + mlx_optional_dtype dtype, + const mlx_stream s) { + return mlx_cumsum_(res, a, reverse, inclusive, dtype, s); } static inline int mlx_degrees(mlx_array* res, const mlx_array a, const mlx_stream s) { return mlx_degrees_(res, a, s); @@ -5552,6 +5870,14 @@ static inline int mlx_diagonal( const mlx_stream s) { return mlx_diagonal_(res, a, offset, axis1, axis2, s); } +static inline int mlx_diff( + mlx_array* res, + const mlx_array a, + int n, + int axis, + const mlx_stream s) { + return mlx_diff_(res, a, n, axis, s); +} static inline int mlx_divide( mlx_array* res, const mlx_array a, @@ -5624,6 +5950,24 @@ static inline int mlx_flatten( const mlx_stream s) { return mlx_flatten_(res, a, start_axis, end_axis, s); } +static inline int mlx_flip_axes( + mlx_array* res, + const mlx_array a, + const int* axes, + size_t axes_num, + const mlx_stream s) { + return mlx_flip_axes_(res, a, axes, axes_num, s); +} +static inline int mlx_flip_axis( + mlx_array* res, + const mlx_array a, + int axis, + const mlx_stream s) { + return mlx_flip_axis_(res, a, axis, s); +} +static inline int mlx_flip(mlx_array* res, const mlx_array a, const mlx_stream s) { + return mlx_flip_(res, a, s); +} static inline int mlx_floor(mlx_array* res, const mlx_array a, const mlx_stream s) { return mlx_floor_(res, a, s); } @@ -5705,6 +6049,22 @@ static inline int mlx_gather_qmm( const mlx_stream s) { return mlx_gather_qmm_(res, x, w, scales, biases, lhs_indices, rhs_indices, transpose, group_size, bits, mode, sorted_indices, s); } +static inline int mlx_gather_qqmm( + mlx_array* res, + const mlx_array x, + const mlx_array w, + const mlx_array scales_w /* may be null */, + const mlx_array lhs_indices /* may be null */, + const mlx_array rhs_indices /* may be null */, + mlx_optional_int group_size, + mlx_optional_int bits, + const char* mode, + const mlx_array global_scale_x /* may be null */, + const mlx_array global_scale_w /* may be null */, + bool sorted_indices, + const mlx_stream s) { + return mlx_gather_qqmm_(res, x, w, scales_w, lhs_indices, rhs_indices, group_size, bits, mode, global_scale_x, global_scale_w, sorted_indices, s); +} static inline int mlx_greater( mlx_array* res, const mlx_array a, @@ -5798,6 +6158,16 @@ static inline int mlx_less_equal( const mlx_stream s) { return mlx_less_equal_(res, a, b, s); } +static inline int mlx_linspace_endpoint( + mlx_array* res, + double start, + double stop, + int num, + bool endpoint, + mlx_dtype dtype, + const mlx_stream s) { + return mlx_linspace_endpoint_(res, start, stop, num, endpoint, dtype, s); +} static inline int mlx_linspace( mlx_array* res, double start, @@ -5826,14 +6196,22 @@ static inline int mlx_logaddexp( const mlx_stream s) { return mlx_logaddexp_(res, a, b, s); } -static inline int mlx_logcumsumexp( +static inline int mlx_logcumsumexp_axis( mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, const mlx_stream s) { - return mlx_logcumsumexp_(res, a, axis, reverse, inclusive, s); + return mlx_logcumsumexp_axis_(res, a, axis, reverse, inclusive, s); +} +static inline int mlx_logcumsumexp( + mlx_array* res, + const mlx_array a, + bool reverse, + bool inclusive, + const mlx_stream s) { + return mlx_logcumsumexp_(res, a, reverse, inclusive, s); } static inline int mlx_logical_and( mlx_array* res, @@ -5852,6 +6230,13 @@ static inline int mlx_logical_or( const mlx_stream s) { return mlx_logical_or_(res, a, b, s); } +static inline int mlx_logical_xor( + mlx_array* res, + const mlx_array a, + const mlx_array b, + const mlx_stream s) { + return mlx_logical_xor_(res, a, b, s); +} static inline int mlx_logsumexp_axes( mlx_array* res, const mlx_array a, @@ -5946,14 +6331,29 @@ static inline int mlx_mean( const mlx_stream s) { return mlx_mean_(res, a, keepdims, s); } -static inline int mlx_median( +static inline int mlx_median_axes( mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s) { - return mlx_median_(res, a, axes, axes_num, keepdims, s); + return mlx_median_axes_(res, a, axes, axes_num, keepdims, s); +} +static inline int mlx_median_axis( + mlx_array* res, + const mlx_array a, + int axis, + bool keepdims, + const mlx_stream s) { + return mlx_median_axis_(res, a, axis, keepdims, s); +} +static inline int mlx_median( + mlx_array* res, + const mlx_array a, + bool keepdims, + const mlx_stream s) { + return mlx_median_(res, a, keepdims, s); } static inline int mlx_meshgrid( mlx_vector_array* res, @@ -6046,6 +6446,13 @@ static inline int mlx_ones( const mlx_stream s) { return mlx_ones_(res, shape, shape_num, dtype, s); } +static inline int mlx_ones_like_dtype( + mlx_array* res, + const mlx_array a, + mlx_dtype dtype, + const mlx_stream s) { + return mlx_ones_like_dtype_(res, a, dtype, s); +} static inline int mlx_ones_like(mlx_array* res, const mlx_array a, const mlx_stream s) { return mlx_ones_like_(res, a, s); } @@ -6094,6 +6501,9 @@ static inline int mlx_partition( const mlx_stream s) { return mlx_partition_(res, a, kth, s); } +static inline int mlx_positive(mlx_array* res, const mlx_array a, const mlx_stream s) { + return mlx_positive_(res, a, s); +} static inline int mlx_power( mlx_array* res, const mlx_array a, @@ -6357,6 +6767,14 @@ static inline int mlx_scatter_prod_single( const mlx_stream s) { return mlx_scatter_prod_single_(res, a, indices, updates, axis, s); } +static inline int mlx_searchsorted( + mlx_array* res, + const mlx_array sorted_sequence, + const mlx_array values, + const char* side, + const mlx_stream s) { + return mlx_searchsorted_(res, sorted_sequence, values, side, s); +} static inline int mlx_segmented_mm( mlx_array* res, const mlx_array a, @@ -6702,7 +7120,7 @@ static inline int mlx_topk_axis( static inline int mlx_topk(mlx_array* res, const mlx_array a, int k, const mlx_stream s) { return mlx_topk_(res, a, k, s); } -static inline int mlx_trace( +static inline int mlx_trace_axes( mlx_array* res, const mlx_array a, int offset, @@ -6710,7 +7128,10 @@ static inline int mlx_trace( int axis2, mlx_dtype dtype, const mlx_stream s) { - return mlx_trace_(res, a, offset, axis1, axis2, dtype, s); + return mlx_trace_axes_(res, a, offset, axis1, axis2, dtype, s); +} +static inline int mlx_trace(mlx_array* res, const mlx_array a, const mlx_stream s) { + return mlx_trace_(res, a, s); } static inline int mlx_transpose_axes( mlx_array* res, @@ -6738,6 +7159,9 @@ static inline int mlx_tril(mlx_array* res, const mlx_array x, int k, const mlx_s static inline int mlx_triu(mlx_array* res, const mlx_array x, int k, const mlx_stream s) { return mlx_triu_(res, x, k, s); } +static inline int mlx_trunc(mlx_array* res, const mlx_array a, const mlx_stream s) { + return mlx_trunc_(res, a, s); +} static inline int mlx_unflatten( mlx_array* res, const mlx_array a, @@ -6747,6 +7171,16 @@ static inline int mlx_unflatten( const mlx_stream s) { return mlx_unflatten_(res, a, axis, shape, shape_num, s); } +static inline int mlx_unstack_axis( + mlx_vector_array* res, + const mlx_array a, + int axis, + const mlx_stream s) { + return mlx_unstack_axis_(res, a, axis, s); +} +static inline int mlx_unstack(mlx_vector_array* res, const mlx_array a, const mlx_stream s) { + return mlx_unstack_(res, a, s); +} static inline int mlx_var_axes( mlx_array* res, const mlx_array a, @@ -6774,6 +7208,14 @@ static inline int mlx_var( const mlx_stream s) { return mlx_var_(res, a, keepdims, ddof, s); } +static inline int mlx_vecdot( + mlx_array* res, + const mlx_array a, + const mlx_array b, + int axis, + const mlx_stream s) { + return mlx_vecdot_(res, a, b, axis, s); +} static inline int mlx_view( mlx_array* res, const mlx_array a, @@ -6797,6 +7239,13 @@ static inline int mlx_zeros( const mlx_stream s) { return mlx_zeros_(res, shape, shape_num, dtype, s); } +static inline int mlx_zeros_like_dtype( + mlx_array* res, + const mlx_array a, + mlx_dtype dtype, + const mlx_stream s) { + return mlx_zeros_like_dtype_(res, a, dtype, s); +} static inline int mlx_zeros_like(mlx_array* res, const mlx_array a, const mlx_stream s) { return mlx_zeros_like_(res, a, s); } @@ -6972,6 +7421,9 @@ static inline mlx_stream mlx_stream_new(void) { static inline mlx_stream mlx_stream_new_device(mlx_device dev) { return mlx_stream_new_device_(dev); } +static inline mlx_stream mlx_stream_new_thread_unsafe(mlx_device dev) { + return mlx_stream_new_thread_unsafe_(dev); +} static inline int mlx_stream_set(mlx_stream* stream, const mlx_stream src) { return mlx_stream_set_(stream, src); } diff --git a/x/mlxrunner/mlx/include/mlx/c/compile.h b/x/mlxrunner/mlx/include/mlx/c/compile.h index ff3abc4ee..68605ce5f 100644 --- a/x/mlxrunner/mlx/include/mlx/c/compile.h +++ b/x/mlxrunner/mlx/include/mlx/c/compile.h @@ -35,6 +35,12 @@ typedef enum mlx_compile_mode_ { MLX_COMPILE_MODE_ENABLED } mlx_compile_mode; +typedef struct mlx_compile_cache_ { + void* ctx; +} mlx_compile_cache; +mlx_compile_cache mlx_compile_cache_new(); +int mlx_compile_cache_free(mlx_compile_cache cache); + int mlx_compile(mlx_closure* res, const mlx_closure fun, bool shapeless); int mlx_detail_compile( mlx_closure* res, @@ -43,10 +49,9 @@ int mlx_detail_compile( bool shapeless, const uint64_t* constants, size_t constants_num); -int mlx_detail_compile_clear_cache(void); - -int mlx_detail_compile_erase(uintptr_t fun_id); - +int mlx_detail_compile_cache(mlx_compile_cache* res); +int mlx_detail_compile_clear_cache(const mlx_compile_cache cache); +int mlx_detail_compile_erase(const mlx_compile_cache cache, uintptr_t fun_id); int mlx_disable_compile(void); int mlx_enable_compile(void); int mlx_set_compile_mode(mlx_compile_mode mode); diff --git a/x/mlxrunner/mlx/include/mlx/c/fast.h b/x/mlxrunner/mlx/include/mlx/c/fast.h index e6c5bdf62..914fd2a28 100644 --- a/x/mlxrunner/mlx/include/mlx/c/fast.h +++ b/x/mlxrunner/mlx/include/mlx/c/fast.h @@ -28,6 +28,12 @@ extern "C" { */ /**@{*/ +int mlx_fast_cross_entropy( + mlx_array* res, + const mlx_array logits, + const mlx_array targets, + const mlx_stream s); + typedef struct mlx_fast_cuda_kernel_config_ { void* ctx; } mlx_fast_cuda_kernel_config; diff --git a/x/mlxrunner/mlx/include/mlx/c/linalg.h b/x/mlxrunner/mlx/include/mlx/c/linalg.h index 91d5d661e..6d5b5ae15 100644 --- a/x/mlxrunner/mlx/include/mlx/c/linalg.h +++ b/x/mlxrunner/mlx/include/mlx/c/linalg.h @@ -44,6 +44,7 @@ int mlx_linalg_cross( const mlx_array b, int axis, const mlx_stream s); +int mlx_linalg_det(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_linalg_eig( mlx_array* res_0, mlx_array* res_1, @@ -97,6 +98,11 @@ int mlx_linalg_qr( mlx_array* res_1, const mlx_array a, const mlx_stream s); +int mlx_linalg_slogdet( + mlx_array* res_0, + mlx_array* res_1, + const mlx_array a, + const mlx_stream s); int mlx_linalg_solve( mlx_array* res, const mlx_array a, diff --git a/x/mlxrunner/mlx/include/mlx/c/metal.h b/x/mlxrunner/mlx/include/mlx/c/metal.h index 5877b224b..f65d15b62 100644 --- a/x/mlxrunner/mlx/include/mlx/c/metal.h +++ b/x/mlxrunner/mlx/include/mlx/c/metal.h @@ -28,7 +28,9 @@ extern "C" { */ /**@{*/ +int mlx_metal_get_metallib_path(mlx_string* res); int mlx_metal_is_available(bool* res); +int mlx_metal_set_metallib_path(const char* path); int mlx_metal_start_capture(const char* path); int mlx_metal_stop_capture(void); diff --git a/x/mlxrunner/mlx/include/mlx/c/ops.h b/x/mlxrunner/mlx/include/mlx/c/ops.h index 44fc09c4d..741778570 100644 --- a/x/mlxrunner/mlx/include/mlx/c/ops.h +++ b/x/mlxrunner/mlx/include/mlx/c/ops.h @@ -323,7 +323,38 @@ int mlx_conv_transpose3d( int mlx_copy(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_cos(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_cosh(mlx_array* res, const mlx_array a, const mlx_stream s); +int mlx_count_nonzero_axis( + mlx_array* res, + const mlx_array a, + int axis, + bool keepdims, + const mlx_stream s); +int mlx_count_nonzero_axes( + mlx_array* res, + const mlx_array a, + const int* axes, + size_t axes_num, + bool keepdims, + const mlx_stream s); +int mlx_count_nonzero( + mlx_array* res, + const mlx_array a, + bool keepdims, + const mlx_stream s); +int mlx_cummax_axis( + mlx_array* res, + const mlx_array a, + int axis, + bool reverse, + bool inclusive, + const mlx_stream s); int mlx_cummax( + mlx_array* res, + const mlx_array a, + bool reverse, + bool inclusive, + const mlx_stream s); +int mlx_cummin_axis( mlx_array* res, const mlx_array a, int axis, @@ -331,25 +362,40 @@ int mlx_cummax( bool inclusive, const mlx_stream s); int mlx_cummin( + mlx_array* res, + const mlx_array a, + bool reverse, + bool inclusive, + const mlx_stream s); +int mlx_cumprod_axis( mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, + mlx_optional_dtype dtype, const mlx_stream s); int mlx_cumprod( mlx_array* res, const mlx_array a, - int axis, bool reverse, bool inclusive, + mlx_optional_dtype dtype, const mlx_stream s); -int mlx_cumsum( +int mlx_cumsum_axis( mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, + mlx_optional_dtype dtype, + const mlx_stream s); +int mlx_cumsum( + mlx_array* res, + const mlx_array a, + bool reverse, + bool inclusive, + mlx_optional_dtype dtype, const mlx_stream s); int mlx_degrees(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_depends( @@ -375,6 +421,12 @@ int mlx_diagonal( int axis1, int axis2, const mlx_stream s); +int mlx_diff( + mlx_array* res, + const mlx_array a, + int n, + int axis, + const mlx_stream s); int mlx_divide( mlx_array* res, const mlx_array a, @@ -423,6 +475,18 @@ int mlx_flatten( int start_axis, int end_axis, const mlx_stream s); +int mlx_flip_axes( + mlx_array* res, + const mlx_array a, + const int* axes, + size_t axes_num, + const mlx_stream s); +int mlx_flip_axis( + mlx_array* res, + const mlx_array a, + int axis, + const mlx_stream s); +int mlx_flip(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_floor(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_floor_divide( mlx_array* res, @@ -486,6 +550,20 @@ int mlx_gather_qmm( const char* mode, bool sorted_indices, const mlx_stream s); +int mlx_gather_qqmm( + mlx_array* res, + const mlx_array x, + const mlx_array w, + const mlx_array scales_w /* may be null */, + const mlx_array lhs_indices /* may be null */, + const mlx_array rhs_indices /* may be null */, + mlx_optional_int group_size, + mlx_optional_int bits, + const char* mode, + const mlx_array global_scale_x /* may be null */, + const mlx_array global_scale_w /* may be null */, + bool sorted_indices, + const mlx_stream s); int mlx_greater( mlx_array* res, const mlx_array a, @@ -543,6 +621,14 @@ int mlx_less_equal( const mlx_array a, const mlx_array b, const mlx_stream s); +int mlx_linspace_endpoint( + mlx_array* res, + double start, + double stop, + int num, + bool endpoint, + mlx_dtype dtype, + const mlx_stream s); int mlx_linspace( mlx_array* res, double start, @@ -559,13 +645,19 @@ int mlx_logaddexp( const mlx_array a, const mlx_array b, const mlx_stream s); -int mlx_logcumsumexp( +int mlx_logcumsumexp_axis( mlx_array* res, const mlx_array a, int axis, bool reverse, bool inclusive, const mlx_stream s); +int mlx_logcumsumexp( + mlx_array* res, + const mlx_array a, + bool reverse, + bool inclusive, + const mlx_stream s); int mlx_logical_and( mlx_array* res, const mlx_array a, @@ -577,6 +669,11 @@ int mlx_logical_or( const mlx_array a, const mlx_array b, const mlx_stream s); +int mlx_logical_xor( + mlx_array* res, + const mlx_array a, + const mlx_array b, + const mlx_stream s); int mlx_logsumexp_axes( mlx_array* res, const mlx_array a, @@ -647,13 +744,24 @@ int mlx_mean( const mlx_array a, bool keepdims, const mlx_stream s); -int mlx_median( +int mlx_median_axes( mlx_array* res, const mlx_array a, const int* axes, size_t axes_num, bool keepdims, const mlx_stream s); +int mlx_median_axis( + mlx_array* res, + const mlx_array a, + int axis, + bool keepdims, + const mlx_stream s); +int mlx_median( + mlx_array* res, + const mlx_array a, + bool keepdims, + const mlx_stream s); int mlx_meshgrid( mlx_vector_array* res, const mlx_vector_array arrays, @@ -721,6 +829,11 @@ int mlx_ones( size_t shape_num, mlx_dtype dtype, const mlx_stream s); +int mlx_ones_like_dtype( + mlx_array* res, + const mlx_array a, + mlx_dtype dtype, + const mlx_stream s); int mlx_ones_like(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_outer( mlx_array* res, @@ -757,6 +870,7 @@ int mlx_partition( const mlx_array a, int kth, const mlx_stream s); +int mlx_positive(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_power( mlx_array* res, const mlx_array a, @@ -956,6 +1070,12 @@ int mlx_scatter_prod_single( const mlx_array updates, int axis, const mlx_stream s); +int mlx_searchsorted( + mlx_array* res, + const mlx_array sorted_sequence, + const mlx_array values, + const char* side, + const mlx_stream s); int mlx_segmented_mm( mlx_array* res, const mlx_array a, @@ -1207,7 +1327,7 @@ int mlx_topk_axis( int axis, const mlx_stream s); int mlx_topk(mlx_array* res, const mlx_array a, int k, const mlx_stream s); -int mlx_trace( +int mlx_trace_axes( mlx_array* res, const mlx_array a, int offset, @@ -1215,6 +1335,7 @@ int mlx_trace( int axis2, mlx_dtype dtype, const mlx_stream s); +int mlx_trace(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_transpose_axes( mlx_array* res, const mlx_array a, @@ -1231,6 +1352,7 @@ int mlx_tri( const mlx_stream s); int mlx_tril(mlx_array* res, const mlx_array x, int k, const mlx_stream s); int mlx_triu(mlx_array* res, const mlx_array x, int k, const mlx_stream s); +int mlx_trunc(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_unflatten( mlx_array* res, const mlx_array a, @@ -1238,6 +1360,12 @@ int mlx_unflatten( const int* shape, size_t shape_num, const mlx_stream s); +int mlx_unstack_axis( + mlx_vector_array* res, + const mlx_array a, + int axis, + const mlx_stream s); +int mlx_unstack(mlx_vector_array* res, const mlx_array a, const mlx_stream s); int mlx_var_axes( mlx_array* res, const mlx_array a, @@ -1259,6 +1387,12 @@ int mlx_var( bool keepdims, int ddof, const mlx_stream s); +int mlx_vecdot( + mlx_array* res, + const mlx_array a, + const mlx_array b, + int axis, + const mlx_stream s); int mlx_view( mlx_array* res, const mlx_array a, @@ -1276,6 +1410,11 @@ int mlx_zeros( size_t shape_num, mlx_dtype dtype, const mlx_stream s); +int mlx_zeros_like_dtype( + mlx_array* res, + const mlx_array a, + mlx_dtype dtype, + const mlx_stream s); int mlx_zeros_like(mlx_array* res, const mlx_array a, const mlx_stream s); /**@}*/ diff --git a/x/mlxrunner/mlx/include/mlx/c/stream.h b/x/mlxrunner/mlx/include/mlx/c/stream.h index d5865b806..abd26701b 100644 --- a/x/mlxrunner/mlx/include/mlx/c/stream.h +++ b/x/mlxrunner/mlx/include/mlx/c/stream.h @@ -33,6 +33,18 @@ mlx_stream mlx_stream_new(void); * Returns a new stream on a device. */ mlx_stream mlx_stream_new_device(mlx_device dev); +/** + * Returns a new stream on a device that can be used from any thread. + * + * Streams are otherwise thread affine: a stream's GPU command encoder is + * registered per thread, so evaluating on a stream from a thread other than + * the one that created it fails. Streams returned here are registered + * globally instead. + * + * MLX applies no synchronization to these streams -- it is the caller's + * responsibility to ensure there are no data races on them. + */ +mlx_stream mlx_stream_new_thread_unsafe(mlx_device dev); /** * Set stream to provided src stream. */ diff --git a/x/mlxrunner/mlx/ops.go b/x/mlxrunner/mlx/ops.go index 060075aa1..304cdbb40 100644 --- a/x/mlxrunner/mlx/ops.go +++ b/x/mlxrunner/mlx/ops.go @@ -103,7 +103,8 @@ func (t *Array) Concatenate(axis int, others ...*Array) *Array { func (t *Array) Cumsum(axis int, reverse, inclusive bool) *Array { out := New("CUMSUM") - C.mlx_cumsum(&out.ctx, t.ctx, C.int(axis), C.bool(reverse), C.bool(inclusive), DefaultStream().ctx) + optDtype := C.mlx_optional_dtype{has_value: false} + C.mlx_cumsum_axis(&out.ctx, t.ctx, C.int(axis), C.bool(reverse), C.bool(inclusive), optDtype, DefaultStream().ctx) return out }