CERBERUS/ollama
1
0
Fork 0
mirror of https://github.com/ollama/ollama.git synced 2026-10-20 18:41:17 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

MLX, MLX-C: version bump (#18080)

Browse source 
* MLX, MLX-C: version bump

* resolve comments

* address comments
This commit is contained in:
Daniel Hiltgen 2026-09-01 16:24:04 -07:00 • committed by GitHub
parent f348c7e3f5
commit 3ba380d0be
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
17 changed files with 920 additions and 152 deletions
Download patch file Download diff file Expand all files Collapse all files

5
.github/workflows/release.yaml vendored
View file

@ -171,6 +171,11 @@ jobs:
- '"cufft_dev"'
- '"nvrtc"'
- '"nvrtc_dev"'
- '"cusolver"'
- '"cusolver_dev"'
- '"cusparse"'
- '"cusparse_dev"'
- '"nvjitlink"'
- '"crt"'
- '"nvvm"'
- '"nvptxcompiler"'

5
.github/workflows/test.yaml vendored
View file

@ -231,6 +231,11 @@ jobs:
- '"cufft_dev"'
- '"nvrtc"'
- '"nvrtc_dev"'
- '"cusolver"'
- '"cusolver_dev"'
- '"cusparse"'
- '"cusparse_dev"'
- '"nvjitlink"'
- '"crt"'
- '"nvvm"'
- '"nvptxcompiler"'

1
Dockerfile
View file

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

2
MLX_C_VERSION
View file

@ -1 +1 @@
fba4470b89073180056c9ea46c443051375f7399
c74db5307cc8ce122f48d97ef951b30578674e7f

2
MLX_VERSION
View file

@ -1 +1 @@
c793734eb715dbcfdb1ced58e348ec53c2d7ed85
37c26e5755da637255d57ea34b4879196a485301

15
cmake/local.cmake
View file

@ -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

56
cmake/mlx/CMakeLists.txt
View file

@ -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}")

88
mlx/compat/0001-mlx-c-regen-0.32.1.patch
View file

@ -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);
/**@}*/

205
x/mlxrunner/mlx/generated.c
View file

@ -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);

502
x/mlxrunner/mlx/generated.h
View file

File diff suppressed because it is too large Load diff

13
x/mlxrunner/mlx/include/mlx/c/compile.h
View file

@ -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);

6
x/mlxrunner/mlx/include/mlx/c/fast.h
View file

@ -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;

6
x/mlxrunner/mlx/include/mlx/c/linalg.h
View file

@ -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,

2
x/mlxrunner/mlx/include/mlx/c/metal.h
View file

@ -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);

149
x/mlxrunner/mlx/include/mlx/c/ops.h
View file

@ -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);
/**@}*/

12
x/mlxrunner/mlx/include/mlx/c/stream.h
View file

@ -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.
*/

3
x/mlxrunner/mlx/ops.go
View file

@ -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
}