From 687508dd98a9cca8569effc190d36037404ee1ff Mon Sep 17 00:00:00 2001 From: Daniel Hiltgen Date: Thu, 22 Jan 2026 17:16:22 -0800 Subject: [PATCH] win: symbol exports and minor fixes (#3024) Co-authored-by: Cheng --- CMakeLists.txt | 5 +- benchmarks/cpp/CMakeLists.txt | 8 + docs/Doxyfile | 1 + examples/cpp/CMakeLists.txt | 8 + mlx/CMakeLists.txt | 22 +- mlx/allocator.h | 8 +- mlx/api.h | 31 + mlx/array.h | 7 +- mlx/backend/cpu/jit_compiler.cpp | 12 +- mlx/backend/cuda/cuda.h | 4 +- mlx/backend/gpu/available.h | 4 +- mlx/backend/metal/metal.h | 13 +- mlx/backend/no_gpu/allocator.cpp | 1 + mlx/compile.h | 11 +- mlx/compile_impl.h | 9 +- mlx/device.h | 14 +- mlx/distributed/distributed.h | 9 +- mlx/distributed/nccl/nccl.cpp | 22 +- mlx/distributed/nccl/nccl_stub/nccl_stubs.cpp | 2 + mlx/distributed/ops.h | 17 +- mlx/dtype.h | 15 +- mlx/einsum.h | 5 +- mlx/export.h | 27 +- mlx/export_impl.h | 19 +- mlx/fast.h | 17 +- mlx/fft.h | 54 +- mlx/graph_utils.h | 7 +- mlx/io.h | 28 +- mlx/io/load.cpp | 3 +- mlx/linalg.h | 50 +- mlx/memory.h | 20 +- mlx/ops.h | 673 +++++++++--------- mlx/primitives.h | 5 +- mlx/random.h | 47 +- mlx/scheduler.h | 3 +- mlx/stream.h | 15 +- mlx/transforms.h | 28 +- mlx/transforms_impl.h | 6 +- mlx/utils.h | 33 +- mlx/version.cpp | 2 + mlx/version.h | 4 +- tests/CMakeLists.txt | 11 + tests/tests.cpp | 4 +- 43 files changed, 701 insertions(+), 583 deletions(-) create mode 100644 mlx/api.h diff --git a/CMakeLists.txt b/CMakeLists.txt index 2e2c05da..66731c4c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -257,9 +257,12 @@ if(MLX_BUILD_CPU) GIT_REPOSITORY https://github.com/OpenMathLib/OpenBLAS.git GIT_TAG v0.3.28 EXCLUDE_FROM_ALL) - set(BUILD_STATIC_LIBS ON) # link statically + block(PROPAGATE openblas_SOURCE_DIR) + set(BUILD_SHARED_LIBS OFF) # link statically + set(BUILD_STATIC_LIBS ON) set(NOFORTRAN ON) # msvc has no fortran compiler FetchContent_MakeAvailable(openblas) + endblock() target_link_libraries(mlx PRIVATE openblas) target_include_directories( mlx PRIVATE "${openblas_SOURCE_DIR}/lapack-netlib/LAPACKE/include" diff --git a/benchmarks/cpp/CMakeLists.txt b/benchmarks/cpp/CMakeLists.txt index 82d5ffce..50b689db 100644 --- a/benchmarks/cpp/CMakeLists.txt +++ b/benchmarks/cpp/CMakeLists.txt @@ -3,6 +3,14 @@ function(build_benchmark SRCFILE) set(target "${src_name}") add_executable(${target} ${SRCFILE}) target_link_libraries(${target} PRIVATE mlx) + # On Windows, copy the mlx DLL to the executable directory for runtime loading + if(WIN32 AND BUILD_SHARED_LIBS) + add_custom_command( + TARGET ${target} + POST_BUILD + COMMAND ${CMAKE_COMMAND} -E copy_if_different $ + $) + endif() endfunction(build_benchmark) build_benchmark(single_ops.cpp) diff --git a/docs/Doxyfile b/docs/Doxyfile index e47712d4..d24a895b 100644 --- a/docs/Doxyfile +++ b/docs/Doxyfile @@ -26,6 +26,7 @@ ENABLE_PREPROCESSING = YES MACRO_EXPANSION = YES EXPAND_ONLY_PREDEF = NO SKIP_FUNCTION_MACROS = NO +PREDEFINED = MLX_API= ################################################################################ # Compound extraction control. # diff --git a/examples/cpp/CMakeLists.txt b/examples/cpp/CMakeLists.txt index 1b3969c9..433b8494 100644 --- a/examples/cpp/CMakeLists.txt +++ b/examples/cpp/CMakeLists.txt @@ -3,6 +3,14 @@ function(build_example SRCFILE) set(target "${src_name}") add_executable(${target} ${SRCFILE}) target_link_libraries(${target} PRIVATE mlx) + # On Windows, copy the mlx DLL to the executable directory for runtime loading + if(WIN32 AND BUILD_SHARED_LIBS) + add_custom_command( + TARGET ${target} + POST_BUILD + COMMAND ${CMAKE_COMMAND} -E copy_if_different $ + $) + endif() endfunction(build_example) build_example(tutorial.cpp) diff --git a/mlx/CMakeLists.txt b/mlx/CMakeLists.txt index 9dc58e78..578078cc 100644 --- a/mlx/CMakeLists.txt +++ b/mlx/CMakeLists.txt @@ -22,11 +22,18 @@ target_sources( # Define MLX_VERSION only in the version.cpp file. add_library(mlx_version OBJECT ${CMAKE_CURRENT_SOURCE_DIR}/version.cpp) target_compile_definitions(mlx_version PRIVATE MLX_VERSION="${MLX_VERSION}") +# mlx_version needs access to api.h for MLX_API export macro +target_include_directories(mlx_version PRIVATE ${PROJECT_SOURCE_DIR}) +# On Windows shared lib builds, mlx_version also needs MLX_EXPORT for proper DLL +# linkage +if(WIN32 AND BUILD_SHARED_LIBS) + target_compile_definitions(mlx_version PRIVATE MLX_EXPORT) +endif() target_link_libraries(mlx PRIVATE $) if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU") # Supress warnings: note: parameter passing for argument of type - # ‘std::pair’ when C++17 is enabled changed to match C++14 in + # 'std::pair' when C++17 is enabled changed to match C++14 in # GCC 10.1 target_compile_options(mlx PRIVATE -Wno-psabi) endif() @@ -39,11 +46,14 @@ if(MSVC) # expression to only apply to C/CXX, not CUDA (NVCC doesn't understand /bigobj # directly). target_compile_options(mlx PRIVATE $<$:/bigobj>) -endif() - -if(WIN32) - # Export symbols by default to behave like macOS/linux. - set_target_properties(mlx PROPERTIES WINDOWS_EXPORT_ALL_SYMBOLS TRUE) + # Windows DLLs have a 65535 symbol export limit. We use explicit exports via + # MLX_API macro (__declspec(dllexport)) on public API functions only. This + # avoids exporting internal template instantiations. + if(BUILD_SHARED_LIBS) + target_compile_definitions(mlx PRIVATE MLX_EXPORT) + else() + target_compile_definitions(mlx PUBLIC MLX_STATIC) + endif() endif() add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backend/common) diff --git a/mlx/allocator.h b/mlx/allocator.h index cd6a78e7..824deac2 100644 --- a/mlx/allocator.h +++ b/mlx/allocator.h @@ -4,12 +4,14 @@ #include +#include "mlx/api.h" + namespace mlx::core::allocator { // Simple wrapper around buffer pointers // WARNING: Only Buffer objects constructed from and those that wrap // raw pointers from mlx::allocator are supported. -class Buffer { +class MLX_API Buffer { private: void* ptr_; @@ -28,7 +30,7 @@ class Buffer { }; }; -class Allocator { +class MLX_API Allocator { /** Abstract base class for a memory allocator. */ public: virtual Buffer malloc(size_t size) = 0; @@ -47,7 +49,7 @@ class Allocator { virtual ~Allocator() = default; }; -Allocator& allocator(); +MLX_API Allocator& allocator(); inline Buffer malloc(size_t size) { return allocator().malloc(size); diff --git a/mlx/api.h b/mlx/api.h new file mode 100644 index 00000000..e2a701d0 --- /dev/null +++ b/mlx/api.h @@ -0,0 +1,31 @@ +// Copyright © 2024 Apple Inc. + +#pragma once + +// MLX_API macro for controlling symbol visibility +// +// On Windows, DLL symbols must be explicitly exported. This header provides +// the MLX_API macro which: +// - Expands to __declspec(dllexport) when building the MLX DLL (MLX_EXPORT +// defined) +// - Expands to __declspec(dllimport) when using the MLX DLL +// - Expands to nothing for static library builds or non-Windows platforms +// +// Usage: +// MLX_API array arange(double start, double stop, ...); +// class MLX_API SomeClass { ... }; + +#if defined(_WIN32) || defined(_WIN64) +#if defined(MLX_STATIC) +// Static library build - no import/export decorations needed +#define MLX_API +#elif defined(MLX_EXPORT) +#define MLX_API __declspec(dllexport) +#else +#define MLX_API __declspec(dllimport) +#endif +#else +// On non-Windows platforms, symbols are visible by default +// Could use __attribute__((visibility("default"))) if needed +#define MLX_API +#endif diff --git a/mlx/array.h b/mlx/array.h index 1b7f7ed0..9b0f6575 100644 --- a/mlx/array.h +++ b/mlx/array.h @@ -8,6 +8,7 @@ #include #include "mlx/allocator.h" +#include "mlx/api.h" #include "mlx/dtype.h" #include "mlx/event.h" #include "mlx/small_vector.h" @@ -22,7 +23,7 @@ using ShapeElem = int32_t; using Shape = SmallVector; using Strides = SmallVector; -class array { +class MLX_API array { /* An array is really a node in a graph. It contains a shared ArrayDesc * object */ @@ -153,7 +154,7 @@ class array { template T item() const; - struct ArrayIterator { + struct MLX_API ArrayIterator { using iterator_category = std::random_access_iterator_tag; using difference_type = size_t; using value_type = const array; @@ -464,7 +465,7 @@ class array { template void init(const It src); - struct ArrayDesc { + struct MLX_API ArrayDesc { Shape shape; Strides strides; size_t size; diff --git a/mlx/backend/cpu/jit_compiler.cpp b/mlx/backend/cpu/jit_compiler.cpp index 01bab52f..fee31f4c 100644 --- a/mlx/backend/cpu/jit_compiler.cpp +++ b/mlx/backend/cpu/jit_compiler.cpp @@ -34,13 +34,23 @@ struct VisualStudioInfo { arch = "x64"; #endif // Get path of Visual Studio. + // Use -latest to get only the most recent installation when multiple + // versions are installed, avoiding path concatenation issues. std::string vs_path = JitCompiler::exec(fmt::format( "\"{0}\\Microsoft Visual Studio\\Installer\\vswhere.exe\"" - " -property installationPath", + " -latest -property installationPath", std::getenv("ProgramFiles(x86)"))); if (vs_path.empty()) { throw std::runtime_error("Can not find Visual Studio."); } + // Trim any trailing whitespace/newlines from the path + vs_path.erase( + std::find_if( + vs_path.rbegin(), + vs_path.rend(), + [](unsigned char ch) { return !std::isspace(ch); }) + .base(), + vs_path.end()); // Read the envs from vcvarsall. std::string envs = JitCompiler::exec(fmt::format( "\"{0}\\VC\\Auxiliary\\Build\\vcvarsall.bat\" {1} >NUL && set", diff --git a/mlx/backend/cuda/cuda.h b/mlx/backend/cuda/cuda.h index 2c6a5c72..f5495965 100644 --- a/mlx/backend/cuda/cuda.h +++ b/mlx/backend/cuda/cuda.h @@ -2,9 +2,11 @@ #pragma once +#include "mlx/api.h" + namespace mlx::core::cu { /* Check if the CUDA backend is available. */ -bool is_available(); +MLX_API bool is_available(); } // namespace mlx::core::cu diff --git a/mlx/backend/gpu/available.h b/mlx/backend/gpu/available.h index 476c7acf..5d0ef7f8 100644 --- a/mlx/backend/gpu/available.h +++ b/mlx/backend/gpu/available.h @@ -2,8 +2,10 @@ #pragma once +#include "mlx/api.h" + namespace mlx::core::gpu { -bool is_available(); +MLX_API bool is_available(); } // namespace mlx::core::gpu diff --git a/mlx/backend/metal/metal.h b/mlx/backend/metal/metal.h index af2995b6..6662e21e 100644 --- a/mlx/backend/metal/metal.h +++ b/mlx/backend/metal/metal.h @@ -6,17 +6,20 @@ #include #include +#include "mlx/api.h" + namespace mlx::core::metal { /* Check if the Metal backend is available. */ -bool is_available(); +MLX_API bool is_available(); /** Capture a GPU trace, saving it to an absolute file `path` */ -void start_capture(std::string path = ""); -void stop_capture(); +MLX_API void start_capture(std::string path = ""); +MLX_API void stop_capture(); /** Get information about the GPU and system settings. */ -const std::unordered_map>& -device_info(); +MLX_API const + std::unordered_map>& + device_info(); } // namespace mlx::core::metal diff --git a/mlx/backend/no_gpu/allocator.cpp b/mlx/backend/no_gpu/allocator.cpp index 76d008e6..abb83e50 100644 --- a/mlx/backend/no_gpu/allocator.cpp +++ b/mlx/backend/no_gpu/allocator.cpp @@ -4,6 +4,7 @@ #include #include "mlx/allocator.h" +#include "mlx/memory.h" #ifdef __APPLE__ #include "mlx/backend/no_gpu/apple_memory.h" diff --git a/mlx/compile.h b/mlx/compile.h index a076cfbc..eba0983e 100644 --- a/mlx/compile.h +++ b/mlx/compile.h @@ -2,6 +2,7 @@ #pragma once +#include "mlx/api.h" #include "mlx/array.h" namespace mlx::core { @@ -9,11 +10,11 @@ namespace mlx::core { enum class CompileMode { disabled, no_simplify, no_fuse, enabled }; /** Compile takes a function and returns a compiled function. */ -std::function(const std::vector&)> compile( +MLX_API std::function(const std::vector&)> compile( std::function(const std::vector&)> fun, bool shapeless = false); -std::function(const std::vector&)> compile( +MLX_API std::function(const std::vector&)> compile( std::vector (*fun)(const std::vector&), bool shapeless = false); @@ -32,13 +33,13 @@ std::function(const std::vector&)> compile( * Setting the environment variable ``MLX_DISABLE_COMPILE`` can also * be used to disable compilation. */ -void disable_compile(); +MLX_API void disable_compile(); /** Globally enable compilation. * This will override the environment variable ``MLX_DISABLE_COMPILE``. */ -void enable_compile(); +MLX_API void enable_compile(); /** Set the compiler mode to the given value. */ -void set_compile_mode(CompileMode mode); +MLX_API void set_compile_mode(CompileMode mode); } // namespace mlx::core diff --git a/mlx/compile_impl.h b/mlx/compile_impl.h index ae8e26b9..238a8b94 100644 --- a/mlx/compile_impl.h +++ b/mlx/compile_impl.h @@ -4,6 +4,7 @@ #include +#include "mlx/api.h" #include "mlx/array.h" namespace mlx::core::detail { @@ -14,24 +15,24 @@ using ArrayFnWithExtra = // This is not part of the general C++ API as calling with a bad id is a bad // idea. -std::function(const std::vector&)> compile( +MLX_API std::function(const std::vector&)> compile( std::function(const std::vector&)> fun, std::uintptr_t fun_id, bool shapeless = false, std::vector constants = {}); -ArrayFnWithExtra compile( +MLX_API ArrayFnWithExtra compile( ArrayFnWithExtra fun, std::uintptr_t fun_id, bool shapeless, std::vector constants); // Erase cached compile functions -void compile_erase(std::uintptr_t fun_id); +MLX_API void compile_erase(std::uintptr_t fun_id); // Clear the compiler cache causing a recompilation of all compiled functions // when called again. -void compile_clear_cache(); +MLX_API void compile_clear_cache(); bool compile_available_for_device(const Device& device); diff --git a/mlx/device.h b/mlx/device.h index 80c624c1..fb14b917 100644 --- a/mlx/device.h +++ b/mlx/device.h @@ -2,9 +2,11 @@ #pragma once +#include "mlx/api.h" + namespace mlx::core { -struct Device { +struct MLX_API Device { enum class DeviceType { cpu, gpu, @@ -19,13 +21,13 @@ struct Device { int index; }; -const Device& default_device(); +MLX_API const Device& default_device(); -void set_default_device(const Device& d); +MLX_API void set_default_device(const Device& d); -bool operator==(const Device& lhs, const Device& rhs); -bool operator!=(const Device& lhs, const Device& rhs); +MLX_API bool operator==(const Device& lhs, const Device& rhs); +MLX_API bool operator!=(const Device& lhs, const Device& rhs); -bool is_available(const Device& d); +MLX_API bool is_available(const Device& d); } // namespace mlx::core diff --git a/mlx/distributed/distributed.h b/mlx/distributed/distributed.h index a6971dd9..00c7a80e 100644 --- a/mlx/distributed/distributed.h +++ b/mlx/distributed/distributed.h @@ -4,6 +4,7 @@ #include +#include "mlx/api.h" #include "mlx/array.h" #include "mlx/utils.h" @@ -15,15 +16,15 @@ class GroupImpl; }; /* Check if a communication backend is available */ -bool is_available(); -bool is_available(const std::string& bk); +MLX_API bool is_available(); +MLX_API bool is_available(const std::string& bk); /** * A distributed::Group represents a group of independent mlx processes that * can communicate. We must also be able to create sub-groups from a group in * order to define more granular communication. */ -struct Group { +struct MLX_API Group { Group(std::shared_ptr group) : group_(std::move(group)) {} int rank() const; @@ -55,6 +56,6 @@ struct Group { * distributed subsystem. Otherwise simply return a singleton group which will * render communication operations as no-op. */ -Group init(bool strict = false, const std::string& bk = "any"); +MLX_API Group init(bool strict = false, const std::string& bk = "any"); } // namespace mlx::core::distributed diff --git a/mlx/distributed/nccl/nccl.cpp b/mlx/distributed/nccl/nccl.cpp index bb9340ba..d8244bf9 100644 --- a/mlx/distributed/nccl/nccl.cpp +++ b/mlx/distributed/nccl/nccl.cpp @@ -1,9 +1,14 @@ +// NCCL distributed support currently requires Unix socket APIs +// TODO: Add Windows Winsock2 support for Windows builds +#ifndef _WIN32 #include -#include -#include #include #include #include +#endif + +#include +#include #include #include #include @@ -93,6 +98,7 @@ void dispatch_dtype(const array& arr, F&& f) { }); } +#ifndef _WIN32 inline void sendAll(int sock, const void* buf, size_t len) { const char* ptr = reinterpret_cast(buf); while (len > 0) { @@ -118,7 +124,9 @@ inline void recvAll(int sock, void* buf, size_t len) { len -= rec; } } +#endif // _WIN32 +#ifndef _WIN32 inline void bootstrap_unique_id( ncclUniqueId& id, int rank, @@ -246,6 +254,16 @@ inline void bootstrap_unique_id( close(sock); } } +#else // _WIN32 +inline void bootstrap_unique_id( + ncclUniqueId& id, + int rank, + int size, + const std::string& initMethod) { + throw std::runtime_error( + "[nccl] Distributed NCCL is not yet supported on Windows"); +} +#endif // _WIN32 } // namespace detail diff --git a/mlx/distributed/nccl/nccl_stub/nccl_stubs.cpp b/mlx/distributed/nccl/nccl_stub/nccl_stubs.cpp index 43593623..cd1b5222 100644 --- a/mlx/distributed/nccl/nccl_stub/nccl_stubs.cpp +++ b/mlx/distributed/nccl/nccl_stub/nccl_stubs.cpp @@ -1,4 +1,6 @@ +#if !defined(_WIN32) #include +#endif #include #include diff --git a/mlx/distributed/ops.h b/mlx/distributed/ops.h index 7688a5f1..e223c5be 100644 --- a/mlx/distributed/ops.h +++ b/mlx/distributed/ops.h @@ -4,51 +4,52 @@ #include +#include "mlx/api.h" #include "mlx/distributed/distributed.h" #include "mlx/utils.h" namespace mlx::core::distributed { -array all_sum( +MLX_API array all_sum( const array& x, std::optional group = std::nullopt, StreamOrDevice s = {}); -array all_gather( +MLX_API array all_gather( const array& x, std::optional group = std::nullopt, StreamOrDevice S = {}); -array send( +MLX_API array send( const array& x, int dst, std::optional group = std::nullopt, StreamOrDevice s = {}); -array recv( +MLX_API array recv( Shape shape, Dtype dtype, int src, std::optional group = std::nullopt, StreamOrDevice s = {}); -array recv_like( +MLX_API array recv_like( const array& x, int src, std::optional group = std::nullopt, StreamOrDevice s = {}); -array all_max( +MLX_API array all_max( const array& x, std::optional group = std::nullopt, StreamOrDevice s = {}); -array all_min( +MLX_API array all_min( const array& x, std::optional group = std::nullopt, StreamOrDevice s = {}); -array sum_scatter( +MLX_API array sum_scatter( const array& x, std::optional group = std::nullopt, StreamOrDevice s = {}); diff --git a/mlx/dtype.h b/mlx/dtype.h index e02b6ca3..744ca587 100644 --- a/mlx/dtype.h +++ b/mlx/dtype.h @@ -5,6 +5,7 @@ #include #include +#include "mlx/api.h" #include "mlx/types/complex.h" #include "mlx/types/half_types.h" @@ -94,21 +95,21 @@ inline constexpr Dtype::Category integer = Dtype::Category::integer; inline constexpr Dtype::Category number = Dtype::Category::number; inline constexpr Dtype::Category generic = Dtype::Category::generic; -bool issubdtype(const Dtype& a, const Dtype& b); -bool issubdtype(const Dtype::Category& a, const Dtype& b); -bool issubdtype(const Dtype& a, const Dtype::Category& b); -bool issubdtype(const Dtype::Category& a, const Dtype::Category& b); +MLX_API bool issubdtype(const Dtype& a, const Dtype& b); +MLX_API bool issubdtype(const Dtype::Category& a, const Dtype& b); +MLX_API bool issubdtype(const Dtype& a, const Dtype::Category& b); +MLX_API bool issubdtype(const Dtype::Category& a, const Dtype::Category& b); -Dtype promote_types(const Dtype& t1, const Dtype& t2); +MLX_API Dtype promote_types(const Dtype& t1, const Dtype& t2); inline uint8_t size_of(const Dtype& t) { return t.size(); } -Dtype::Kind kindof(const Dtype& t); +MLX_API Dtype::Kind kindof(const Dtype& t); template -struct TypeToDtype { +struct MLX_API TypeToDtype { operator Dtype(); }; diff --git a/mlx/einsum.h b/mlx/einsum.h index f57e9a77..05588f88 100644 --- a/mlx/einsum.h +++ b/mlx/einsum.h @@ -5,16 +5,17 @@ #include #include +#include "mlx/api.h" #include "mlx/array.h" #include "mlx/utils.h" namespace mlx::core { -std::pair>, std::string> einsum_path( +MLX_API std::pair>, std::string> einsum_path( const std::string& subscripts, const std::vector& operands); -array einsum( +MLX_API array einsum( const std::string& subscripts, const std::vector& operands, StreamOrDevice s = {}); diff --git a/mlx/export.h b/mlx/export.h index 0a8e9fb0..5532f7c8 100644 --- a/mlx/export.h +++ b/mlx/export.h @@ -6,6 +6,7 @@ #include #include #include +#include "mlx/api.h" #include "mlx/array.h" namespace mlx::core { @@ -46,17 +47,17 @@ struct FunctionExporter; * Make an exporter to save multiple traces of a given function to * the same file. */ -FunctionExporter exporter( +MLX_API FunctionExporter exporter( const std::string& file, const std::function(const Args&)>& fun, bool shapeless = false); -FunctionExporter exporter( +MLX_API FunctionExporter exporter( const std::string& file, const std::function(const Kwargs&)>& fun, bool shapeless = false); -FunctionExporter exporter( +MLX_API FunctionExporter exporter( const std::string& path, const std::function(const Args&, const Kwargs&)>& fun, bool shapeless = false); @@ -64,19 +65,19 @@ FunctionExporter exporter( /** * Export a function to a file. */ -void export_function( +MLX_API void export_function( const std::string& file, const std::function(const Args&)>& fun, const Args& args, bool shapeless = false); -void export_function( +MLX_API void export_function( const std::string& file, const std::function(const Kwargs&)>& fun, const Kwargs& kwargs, bool shapeless = false); -void export_function( +MLX_API void export_function( const std::string& file, const std::function(const Args&, const Kwargs&)>& fun, const Args& args, @@ -88,23 +89,23 @@ struct ImportedFunction; /** * Import a function from a file. */ -ImportedFunction import_function(const std::string& file); +MLX_API ImportedFunction import_function(const std::string& file); /** * Make an exporter to export multiple traces of a given function with the same * callback. */ -FunctionExporter exporter( +MLX_API FunctionExporter exporter( const ExportCallback& callback, const std::function(const Args&)>& fun, bool shapeless = false); -FunctionExporter exporter( +MLX_API FunctionExporter exporter( const ExportCallback& callback, const std::function(const Kwargs&)>& fun, bool shapeless = false); -FunctionExporter exporter( +MLX_API FunctionExporter exporter( const ExportCallback& callback, const std::function(const Args&, const Kwargs&)>& fun, bool shapeless = false); @@ -112,19 +113,19 @@ FunctionExporter exporter( /** * Export a function with a callback. */ -void export_function( +MLX_API void export_function( const ExportCallback& callback, const std::function(const Args&)>& fun, const Args& args, bool shapeless = false); -void export_function( +MLX_API void export_function( const ExportCallback& callback, const std::function(const Kwargs&)>& fun, const Kwargs& kwargs, bool shapeless = false); -void export_function( +MLX_API void export_function( const ExportCallback& callback, const std::function(const Args&, const Kwargs&)>& fun, const Args& args, diff --git a/mlx/export_impl.h b/mlx/export_impl.h index be215aaa..467a5f0d 100644 --- a/mlx/export_impl.h +++ b/mlx/export_impl.h @@ -1,5 +1,6 @@ // Copyright © 2024 Apple Inc. +#include "mlx/api.h" #include "mlx/io/load.h" #pragma once @@ -8,7 +9,7 @@ namespace mlx::core { struct FunctionTable; -struct FunctionExporter { +struct MLX_API FunctionExporter { void operator()(const std::initializer_list& args) { this->operator()(Args(args)); } @@ -23,32 +24,32 @@ struct FunctionExporter { FunctionExporter(FunctionExporter&& other) = default; private: - friend FunctionExporter exporter( + friend MLX_API FunctionExporter exporter( const std::string&, const std::function(const Args&)>&, bool shapeless); - friend FunctionExporter exporter( + friend MLX_API FunctionExporter exporter( const std::string&, const std::function(const Kwargs&)>&, bool shapeless); - friend FunctionExporter exporter( + friend MLX_API FunctionExporter exporter( const std::string&, const std::function(const Args&, const Kwargs&)>&, bool shapeless); - friend FunctionExporter exporter( + friend MLX_API FunctionExporter exporter( const ExportCallback&, const std::function(const Args&)>&, bool shapeless); - friend FunctionExporter exporter( + friend MLX_API FunctionExporter exporter( const ExportCallback&, const std::function(const Kwargs&)>&, bool shapeless); - friend FunctionExporter exporter( + friend MLX_API FunctionExporter exporter( const ExportCallback&, const std::function(const Args&, const Kwargs&)>&, bool shapeless); @@ -78,7 +79,7 @@ struct FunctionExporter { std::shared_ptr ftable; }; -struct ImportedFunction { +struct MLX_API ImportedFunction { std::vector operator()( const std::initializer_list& args) const { return this->operator()(Args(args)); @@ -89,7 +90,7 @@ struct ImportedFunction { private: ImportedFunction(const std::string& file); - friend ImportedFunction import_function(const std::string&); + friend MLX_API ImportedFunction import_function(const std::string&); ImportedFunction(); std::shared_ptr ftable; diff --git a/mlx/fast.h b/mlx/fast.h index 0884bac7..1183aba8 100644 --- a/mlx/fast.h +++ b/mlx/fast.h @@ -5,24 +5,25 @@ #include #include +#include "mlx/api.h" #include "mlx/utils.h" namespace mlx::core::fast { -array rms_norm( +MLX_API array rms_norm( const array& x, const std::optional& weight, float eps, StreamOrDevice s = {}); -array layer_norm( +MLX_API array layer_norm( const array& x, const std::optional& weight, const std::optional& bias, float eps, StreamOrDevice s = {}); -array rope( +MLX_API array rope( const array& x, int dims, bool traditional, @@ -32,7 +33,7 @@ array rope( const std::optional& freqs = std::nullopt, StreamOrDevice s = {}); -array rope( +MLX_API array rope( const array& x, int dims, bool traditional, @@ -43,7 +44,7 @@ array rope( StreamOrDevice s = {}); /** Computes: O = softmax(Q @ K.T) @ V **/ -array scaled_dot_product_attention( +MLX_API array scaled_dot_product_attention( const array& queries, const array& keys, const array& values, @@ -67,7 +68,7 @@ using CustomKernelFunction = std::function( bool, StreamOrDevice)>; -CustomKernelFunction metal_kernel( +MLX_API CustomKernelFunction metal_kernel( const std::string& name, const std::vector& input_names, const std::vector& output_names, @@ -76,7 +77,7 @@ CustomKernelFunction metal_kernel( bool ensure_row_contiguous = true, bool atomic_outputs = false); -CustomKernelFunction cuda_kernel( +MLX_API CustomKernelFunction cuda_kernel( const std::string& name, const std::vector& input_names, const std::vector& output_names, @@ -85,7 +86,7 @@ CustomKernelFunction cuda_kernel( bool ensure_row_contiguous = true, int shared_memory = 0); -std::vector precompiled_cuda_kernel( +MLX_API std::vector precompiled_cuda_kernel( const std::string& name, const std::string& compiled_source, const std::vector& inputs, diff --git a/mlx/fft.h b/mlx/fft.h index 163e06b8..9abf2b18 100644 --- a/mlx/fft.h +++ b/mlx/fft.h @@ -6,30 +6,30 @@ #include "array.h" #include "device.h" +#include "mlx/api.h" #include "utils.h" namespace mlx::core::fft { /** Compute the n-dimensional Fourier Transform. */ -array fftn( +MLX_API array fftn( const array& a, const Shape& n, const std::vector& axes, StreamOrDevice s = {}); -array fftn(const array& a, const std::vector& axes, StreamOrDevice s = {}); -array fftn(const array& a, StreamOrDevice s = {}); +MLX_API array +fftn(const array& a, const std::vector& axes, StreamOrDevice s = {}); +MLX_API array fftn(const array& a, StreamOrDevice s = {}); /** Compute the n-dimensional inverse Fourier Transform. */ -array ifftn( +MLX_API array ifftn( const array& a, const Shape& n, const std::vector& axes, StreamOrDevice s = {}); -array ifftn( - const array& a, - const std::vector& axes, - StreamOrDevice s = {}); -array ifftn(const array& a, StreamOrDevice s = {}); +MLX_API array +ifftn(const array& a, const std::vector& axes, StreamOrDevice s = {}); +MLX_API array ifftn(const array& a, StreamOrDevice s = {}); /** Compute the one-dimensional Fourier Transform. */ inline array fft(const array& a, int n, int axis, StreamOrDevice s = {}) { @@ -78,28 +78,24 @@ inline array ifft2( } /** Compute the n-dimensional Fourier Transform on a real input. */ -array rfftn( +MLX_API array rfftn( const array& a, const Shape& n, const std::vector& axes, StreamOrDevice s = {}); -array rfftn( - const array& a, - const std::vector& axes, - StreamOrDevice s = {}); -array rfftn(const array& a, StreamOrDevice s = {}); +MLX_API array +rfftn(const array& a, const std::vector& axes, StreamOrDevice s = {}); +MLX_API array rfftn(const array& a, StreamOrDevice s = {}); /** Compute the n-dimensional inverse of `rfftn`. */ -array irfftn( +MLX_API array irfftn( const array& a, const Shape& n, const std::vector& axes, StreamOrDevice s = {}); -array irfftn( - const array& a, - const std::vector& axes, - StreamOrDevice s = {}); -array irfftn(const array& a, StreamOrDevice s = {}); +MLX_API array +irfftn(const array& a, const std::vector& axes, StreamOrDevice s = {}); +MLX_API array irfftn(const array& a, StreamOrDevice s = {}); /** Compute the one-dimensional Fourier Transform on a real input. */ inline array rfft(const array& a, int n, int axis, StreamOrDevice s = {}) { @@ -146,22 +142,18 @@ inline array irfft2( return irfftn(a, axes, s); } /** Shift the zero-frequency component to the center of the spectrum. */ -array fftshift(const array& a, StreamOrDevice s = {}); +MLX_API array fftshift(const array& a, StreamOrDevice s = {}); /** Shift the zero-frequency component to the center of the spectrum along * specified axes. */ -array fftshift( - const array& a, - const std::vector& axes, - StreamOrDevice s = {}); +MLX_API array +fftshift(const array& a, const std::vector& axes, StreamOrDevice s = {}); /** The inverse of fftshift. */ -array ifftshift(const array& a, StreamOrDevice s = {}); +MLX_API array ifftshift(const array& a, StreamOrDevice s = {}); /** The inverse of fftshift along specified axes. */ -array ifftshift( - const array& a, - const std::vector& axes, - StreamOrDevice s = {}); +MLX_API array +ifftshift(const array& a, const std::vector& axes, StreamOrDevice s = {}); } // namespace mlx::core::fft diff --git a/mlx/graph_utils.h b/mlx/graph_utils.h index fcbeef15..54297c2c 100644 --- a/mlx/graph_utils.h +++ b/mlx/graph_utils.h @@ -4,18 +4,19 @@ #include +#include "mlx/api.h" #include "mlx/array.h" namespace mlx::core { -struct NodeNamer { +struct MLX_API NodeNamer { std::unordered_map names; const std::string& get_name(const array& x); void set_name(const array& x, std::string n); }; -void print_graph( +MLX_API void print_graph( std::ostream& os, NodeNamer namer, const std::vector& outputs); @@ -39,7 +40,7 @@ print_graph(std::ostream& os, NodeNamer namer, Arrays&&... outputs) { std::vector{std::forward(outputs)...}); } -void export_to_dot( +MLX_API void export_to_dot( std::ostream& os, NodeNamer namer, const std::vector& outputs); diff --git a/mlx/io.h b/mlx/io.h index 23380b23..760f2985 100644 --- a/mlx/io.h +++ b/mlx/io.h @@ -5,6 +5,7 @@ #include #include +#include "mlx/api.h" #include "mlx/array.h" #include "mlx/io/load.h" #include "mlx/stream.h" @@ -21,39 +22,38 @@ using SafetensorsLoad = std::pair< std::unordered_map>; /** Save array to out stream in .npy format */ -void save(std::shared_ptr out_stream, array a); +MLX_API void save(std::shared_ptr out_stream, array a); /** Save array to file in .npy format */ -void save(std::string file, array a); +MLX_API void save(std::string file, array a); /** Load array from reader in .npy format */ -array load(std::shared_ptr in_stream, StreamOrDevice s = {}); +MLX_API array +load(std::shared_ptr in_stream, StreamOrDevice s = {}); /** Load array from file in .npy format */ -array load(std::string file, StreamOrDevice s = {}); +MLX_API array load(std::string file, StreamOrDevice s = {}); /** Load array map from .safetensors file format */ -SafetensorsLoad load_safetensors( - std::shared_ptr in_stream, - StreamOrDevice s = {}); -SafetensorsLoad load_safetensors( - const std::string& file, - StreamOrDevice s = {}); +MLX_API SafetensorsLoad +load_safetensors(std::shared_ptr in_stream, StreamOrDevice s = {}); +MLX_API SafetensorsLoad +load_safetensors(const std::string& file, StreamOrDevice s = {}); -void save_safetensors( +MLX_API void save_safetensors( std::shared_ptr in_stream, std::unordered_map, std::unordered_map metadata = {}); -void save_safetensors( +MLX_API void save_safetensors( std::string file, std::unordered_map, std::unordered_map metadata = {}); /** Load array map and metadata from .gguf file format */ -GGUFLoad load_gguf(const std::string& file, StreamOrDevice s = {}); +MLX_API GGUFLoad load_gguf(const std::string& file, StreamOrDevice s = {}); -void save_gguf( +MLX_API void save_gguf( std::string file, std::unordered_map array_map, std::unordered_map meta_data = {}); diff --git a/mlx/io/load.cpp b/mlx/io/load.cpp index 37680cb4..878d6d1a 100644 --- a/mlx/io/load.cpp +++ b/mlx/io/load.cpp @@ -7,13 +7,12 @@ // Used by pread implementation. #ifdef _WIN32 -#ifdef _MSC_VER #define NOMINMAX -#endif #include #endif // _WIN32 #include "mlx/backend/cuda/cuda.h" +#include "mlx/io.h" #include "mlx/io/load.h" #include "mlx/ops.h" #include "mlx/primitives.h" diff --git a/mlx/linalg.h b/mlx/linalg.h index 0690fba9..fe3f83c2 100644 --- a/mlx/linalg.h +++ b/mlx/linalg.h @@ -4,6 +4,7 @@ #include +#include "mlx/api.h" #include "mlx/array.h" #include "mlx/device.h" #include "mlx/ops.h" @@ -22,7 +23,7 @@ namespace mlx::core::linalg { * - If both axis and ord are provided, then the corresponding matrix or vector * norm is computed. At most 2 axes can be specified. */ -array norm( +MLX_API array norm( const array& a, const double ord, const std::optional>& axis = std::nullopt, @@ -36,7 +37,7 @@ inline array norm( StreamOrDevice s = {}) { return norm(a, ord, std::vector{axis}, keepdims, s); } -array norm( +MLX_API array norm( const array& a, const std::string& ord, const std::optional>& axis = std::nullopt, @@ -50,7 +51,7 @@ inline array norm( StreamOrDevice s = {}) { return norm(a, ord, std::vector{axis}, keepdims, s); } -array norm( +MLX_API array norm( const array& a, const std::optional>& axis = std::nullopt, bool keepdims = false, @@ -60,31 +61,36 @@ norm(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {}) { return norm(a, std::vector{axis}, keepdims, s); } -std::pair qr(const array& a, StreamOrDevice s = {}); +MLX_API std::pair qr(const array& a, StreamOrDevice s = {}); -std::vector +MLX_API std::vector svd(const array& a, bool compute_uv, StreamOrDevice s /* = {} */); inline std::vector svd(const array& a, StreamOrDevice s = {}) { return svd(a, true, s); } -array inv(const array& a, StreamOrDevice s = {}); +MLX_API array inv(const array& a, StreamOrDevice s = {}); -array tri_inv(const array& a, bool upper = false, StreamOrDevice s = {}); +MLX_API array +tri_inv(const array& a, bool upper = false, StreamOrDevice s = {}); -array cholesky(const array& a, bool upper = false, StreamOrDevice s = {}); +MLX_API array +cholesky(const array& a, bool upper = false, StreamOrDevice s = {}); -array pinv(const array& a, StreamOrDevice s = {}); +MLX_API array pinv(const array& a, StreamOrDevice s = {}); -array cholesky_inv(const array& a, bool upper = false, StreamOrDevice s = {}); +MLX_API array +cholesky_inv(const array& a, bool upper = false, StreamOrDevice s = {}); -std::vector lu(const array& a, StreamOrDevice s = {}); +MLX_API std::vector lu(const array& a, StreamOrDevice s = {}); -std::pair lu_factor(const array& a, StreamOrDevice s = {}); +MLX_API std::pair lu_factor( + const array& a, + StreamOrDevice s = {}); -array solve(const array& a, const array& b, StreamOrDevice s = {}); +MLX_API array solve(const array& a, const array& b, StreamOrDevice s = {}); -array solve_triangular( +MLX_API array solve_triangular( const array& a, const array& b, bool upper = false, @@ -93,19 +99,17 @@ array solve_triangular( /** * Compute the cross product of two arrays along the given axis. */ -array cross( - const array& a, - const array& b, - int axis = -1, - StreamOrDevice s = {}); +MLX_API array +cross(const array& a, const array& b, int axis = -1, StreamOrDevice s = {}); -std::pair eig(const array& a, StreamOrDevice s = {}); +MLX_API std::pair eig(const array& a, StreamOrDevice s = {}); -array eigvals(const array& a, StreamOrDevice s = {}); +MLX_API array eigvals(const array& a, StreamOrDevice s = {}); -array eigvalsh(const array& a, std::string UPLO = "L", StreamOrDevice s = {}); +MLX_API array +eigvalsh(const array& a, std::string UPLO = "L", StreamOrDevice s = {}); -std::pair +MLX_API std::pair eigh(const array& a, std::string UPLO = "L", StreamOrDevice s = {}); } // namespace mlx::core::linalg diff --git a/mlx/memory.h b/mlx/memory.h index 8a264734..f4eabc99 100644 --- a/mlx/memory.h +++ b/mlx/memory.h @@ -4,6 +4,8 @@ #include +#include "mlx/api.h" + namespace mlx::core { /* Get the actively used memory in bytes. @@ -11,25 +13,25 @@ namespace mlx::core { * Note, this will not always match memory use reported by the system because * it does not include cached memory buffers. * */ -size_t get_active_memory(); +MLX_API size_t get_active_memory(); /* Get the peak amount of used memory in bytes. * * The maximum memory used recorded from the beginning of the program * execution or since the last call to reset_peak_memory. * */ -size_t get_peak_memory(); +MLX_API size_t get_peak_memory(); /* Reset the peak memory to zero. * */ -void reset_peak_memory(); +MLX_API void reset_peak_memory(); /* Get the cache size in bytes. * * The cache includes memory not currently used that has not been returned * to the system allocator. * */ -size_t get_cache_memory(); +MLX_API size_t get_cache_memory(); /* Set the memory limit. * The memory limit is a guideline for the maximum amount of memory to use @@ -42,10 +44,10 @@ size_t get_cache_memory(); * * Returns the previous memory limit. * */ -size_t set_memory_limit(size_t limit); +MLX_API size_t set_memory_limit(size_t limit); /* Get the current memory limit. */ -size_t get_memory_limit(); +MLX_API size_t get_memory_limit(); /* Set the cache limit. * If using more than the given limit, free memory will be reclaimed @@ -56,10 +58,10 @@ size_t get_memory_limit(); * * Returns the previous cache limit. * */ -size_t set_cache_limit(size_t limit); +MLX_API size_t set_cache_limit(size_t limit); /* Clear the memory cache. */ -void clear_cache(); +MLX_API void clear_cache(); /* Set the wired size limit. * @@ -73,6 +75,6 @@ void clear_cache(); * * Returns the previous wired limit. * */ -size_t set_wired_limit(size_t limit); +MLX_API size_t set_wired_limit(size_t limit); } // namespace mlx::core diff --git a/mlx/ops.h b/mlx/ops.h index ff92cbe9..1ff3bbfa 100644 --- a/mlx/ops.h +++ b/mlx/ops.h @@ -4,6 +4,7 @@ #include +#include "mlx/api.h" #include "mlx/array.h" #include "mlx/device.h" #include "mlx/stream.h" @@ -19,24 +20,26 @@ namespace mlx::core { /** * A 1D array of numbers starting at `start` (optional), * stopping at stop, stepping by `step` (optional). */ -array arange( +MLX_API array arange( double start, double stop, double step, Dtype dtype, StreamOrDevice s = {}); -array arange(double start, double stop, double step, StreamOrDevice s = {}); -array arange(double start, double stop, Dtype dtype, StreamOrDevice s = {}); -array arange(double start, double stop, StreamOrDevice s = {}); -array arange(double stop, Dtype dtype, StreamOrDevice s = {}); -array arange(double stop, StreamOrDevice s = {}); +MLX_API array +arange(double start, double stop, double step, StreamOrDevice s = {}); +MLX_API array +arange(double start, double stop, Dtype dtype, StreamOrDevice s = {}); +MLX_API array arange(double start, double stop, StreamOrDevice s = {}); +MLX_API array arange(double stop, Dtype dtype, StreamOrDevice s = {}); +MLX_API array arange(double stop, StreamOrDevice s = {}); -array arange(int start, int stop, int step, StreamOrDevice s = {}); -array arange(int start, int stop, StreamOrDevice s = {}); -array arange(int stop, StreamOrDevice s = {}); +MLX_API array arange(int start, int stop, int step, StreamOrDevice s = {}); +MLX_API array arange(int start, int stop, StreamOrDevice s = {}); +MLX_API array arange(int stop, StreamOrDevice s = {}); /** A 1D array of `num` evenly spaced numbers in the range `[start, stop]` */ -array linspace( +MLX_API array linspace( double start, double stop, int num = 50, @@ -44,10 +47,10 @@ array linspace( StreamOrDevice s = {}); /** Convert an array to the given data type. */ -array astype(array a, Dtype dtype, StreamOrDevice s = {}); +MLX_API array astype(array a, Dtype dtype, StreamOrDevice s = {}); /** Create a view of an array with the given shape and strides. */ -array as_strided( +MLX_API array as_strided( array a, Shape shape, Strides strides, @@ -55,11 +58,11 @@ array as_strided( StreamOrDevice s = {}); /** Copy another array. */ -array copy(array a, StreamOrDevice s = {}); +MLX_API array copy(array a, StreamOrDevice s = {}); /** Fill an array of the given shape with the given value(s). */ -array full(Shape shape, array vals, Dtype dtype, StreamOrDevice s = {}); -array full(Shape shape, array vals, StreamOrDevice s = {}); +MLX_API array full(Shape shape, array vals, Dtype dtype, StreamOrDevice s = {}); +MLX_API array full(Shape shape, array vals, StreamOrDevice s = {}); template array full(Shape shape, T val, Dtype dtype, StreamOrDevice s = {}) { return full(std::move(shape), array(val, dtype), to_stream(s)); @@ -69,8 +72,9 @@ array full(Shape shape, T val, StreamOrDevice s = {}) { return full(std::move(shape), array(val), to_stream(s)); } -array full_like(const array& a, array vals, Dtype dtype, StreamOrDevice s = {}); -array full_like(const array& a, array vals, StreamOrDevice s = {}); +MLX_API array +full_like(const array& a, array vals, Dtype dtype, StreamOrDevice s = {}); +MLX_API array full_like(const array& a, array vals, StreamOrDevice s = {}); template array full_like(const array& a, T val, Dtype dtype, StreamOrDevice s = {}) { return full_like(a, array(val, dtype), dtype, to_stream(s)); @@ -81,22 +85,22 @@ array full_like(const array& a, T val, StreamOrDevice s = {}) { } /** Fill an array of the given shape with zeros. */ -array zeros(const Shape& shape, Dtype dtype, StreamOrDevice s = {}); +MLX_API array zeros(const Shape& shape, Dtype dtype, StreamOrDevice s = {}); inline array zeros(const Shape& shape, StreamOrDevice s = {}) { return zeros(shape, float32, s); } -array zeros_like(const array& a, StreamOrDevice s = {}); +MLX_API array zeros_like(const array& a, StreamOrDevice s = {}); /** Fill an array of the given shape with ones. */ -array ones(const Shape& shape, Dtype dtype, StreamOrDevice s = {}); +MLX_API array ones(const Shape& shape, Dtype dtype, StreamOrDevice s = {}); inline array ones(const Shape& shape, StreamOrDevice s = {}) { return ones(shape, float32, s); } -array ones_like(const array& a, StreamOrDevice s = {}); +MLX_API array ones_like(const array& a, StreamOrDevice s = {}); /** Fill an array of the given shape (n,m) with ones in the specified diagonal * k, and zeros everywhere else. */ -array eye(int n, int m, int k, Dtype dtype, StreamOrDevice s = {}); +MLX_API array eye(int n, int m, int k, Dtype dtype, StreamOrDevice s = {}); inline array eye(int n, Dtype dtype, StreamOrDevice s = {}) { return eye(n, n, 0, dtype, s); } @@ -112,64 +116,63 @@ inline array eye(int n, StreamOrDevice s = {}) { /** Create a square matrix of shape (n,n) of zeros, and ones in the major * diagonal. */ -array identity(int n, Dtype dtype, StreamOrDevice s = {}); +MLX_API array identity(int n, Dtype dtype, StreamOrDevice s = {}); inline array identity(int n, StreamOrDevice s = {}) { return identity(n, float32, s); } -array tri(int n, int m, int k, Dtype type, StreamOrDevice s = {}); +MLX_API array tri(int n, int m, int k, Dtype type, StreamOrDevice s = {}); inline array tri(int n, Dtype type, StreamOrDevice s = {}) { return tri(n, n, 0, type, s); } -array tril(array x, int k = 0, StreamOrDevice s = {}); -array triu(array x, int k = 0, StreamOrDevice s = {}); +MLX_API array tril(array x, int k = 0, StreamOrDevice s = {}); +MLX_API array triu(array x, int k = 0, StreamOrDevice s = {}); /** Reshape an array to the given shape. */ -array reshape(const array& a, Shape shape, StreamOrDevice s = {}); +MLX_API array reshape(const array& a, Shape shape, StreamOrDevice s = {}); /** Unflatten the axis to the given shape. */ -array unflatten(const array& a, int axis, Shape shape, StreamOrDevice s = {}); +MLX_API array +unflatten(const array& a, int axis, Shape shape, StreamOrDevice s = {}); /** Flatten the dimensions in the range `[start_axis, end_axis]` . */ -array flatten( +MLX_API array flatten( const array& a, int start_axis, int end_axis = -1, StreamOrDevice s = {}); /** Flatten the array to 1D. */ -array flatten(const array& a, StreamOrDevice s = {}); +MLX_API array flatten(const array& a, StreamOrDevice s = {}); /** Multiply the array by the Hadamard matrix of corresponding size. */ -array hadamard_transform( +MLX_API array hadamard_transform( const array& a, std::optional scale = std::nullopt, StreamOrDevice s = {}); /** Remove singleton dimensions at the given axes. */ -array squeeze( - const array& a, - const std::vector& axes, - StreamOrDevice s = {}); +MLX_API array +squeeze(const array& a, const std::vector& axes, StreamOrDevice s = {}); /** Remove singleton dimensions at the given axis. */ -array squeeze(const array& a, int axis, StreamOrDevice s = {}); +MLX_API array squeeze(const array& a, int axis, StreamOrDevice s = {}); /** Remove all singleton dimensions. */ -array squeeze(const array& a, StreamOrDevice s = {}); +MLX_API array squeeze(const array& a, StreamOrDevice s = {}); /** Add a singleton dimension at the given axes. */ -array expand_dims( +MLX_API array expand_dims( const array& a, const std::vector& axes, StreamOrDevice s = {}); /** Add a singleton dimension at the given axis. */ -array expand_dims(const array& a, int axis, StreamOrDevice s = {}); +MLX_API array expand_dims(const array& a, int axis, StreamOrDevice s = {}); /** Slice an array. */ -array slice( +MLX_API array slice( const array& a, Shape start, Shape stop, @@ -185,10 +188,11 @@ inline array slice( } /** Slice an array with a stride of 1 in each dimension. */ -array slice(const array& a, Shape start, Shape stop, StreamOrDevice s = {}); +MLX_API array +slice(const array& a, Shape start, Shape stop, StreamOrDevice s = {}); /** Slice an array with dynamic starting indices. */ -array slice( +MLX_API array slice( const array& a, const array& start, std::vector axes, @@ -196,7 +200,7 @@ array slice( StreamOrDevice s = {}); /** Update a slice from the source array. */ -array slice_update( +MLX_API array slice_update( const array& src, const array& update, Shape start, @@ -205,7 +209,7 @@ array slice_update( StreamOrDevice s = {}); /** Update a slice from the source array with stride 1 in each dimension. */ -array slice_update( +MLX_API array slice_update( const array& src, const array& update, Shape start, @@ -213,7 +217,7 @@ array slice_update( StreamOrDevice s = {}); /** Update a slice from the source array with dynamic starting indices. */ -array slice_update( +MLX_API array slice_update( const array& src, const array& update, const array& start, @@ -221,16 +225,17 @@ array slice_update( StreamOrDevice s = {}); /** Split an array into sub-arrays along a given axis. */ -std::vector +MLX_API std::vector split(const array& a, int num_splits, int axis, StreamOrDevice s = {}); -std::vector split(const array& a, int num_splits, StreamOrDevice s = {}); -std::vector +MLX_API std::vector +split(const array& a, int num_splits, StreamOrDevice s = {}); +MLX_API std::vector split(const array& a, const Shape& indices, int axis, StreamOrDevice s = {}); -std::vector +MLX_API std::vector split(const array& a, const Shape& indices, StreamOrDevice s = {}); /** A vector of coordinate arrays from coordinate vectors. */ -std::vector meshgrid( +MLX_API std::vector meshgrid( const std::vector& arrays, bool sparse = false, const std::string& indexing = "xy", @@ -239,28 +244,33 @@ std::vector meshgrid( /** * Clip (limit) the values in an array. */ -array clip( +MLX_API array clip( const array& a, const std::optional& a_min = std::nullopt, const std::optional& a_max = std::nullopt, StreamOrDevice s = {}); /** Concatenate arrays along a given axis. */ -array concatenate(std::vector arrays, int axis, StreamOrDevice s = {}); -array concatenate(std::vector arrays, StreamOrDevice s = {}); +MLX_API array +concatenate(std::vector arrays, int axis, StreamOrDevice s = {}); +MLX_API array concatenate(std::vector arrays, StreamOrDevice s = {}); /** Stack arrays along a new axis. */ -array stack(const std::vector& arrays, int axis, StreamOrDevice s = {}); -array stack(const std::vector& arrays, StreamOrDevice s = {}); +MLX_API array +stack(const std::vector& arrays, int axis, StreamOrDevice s = {}); +MLX_API array stack(const std::vector& arrays, StreamOrDevice s = {}); /** Repeat an array along an axis. */ -array repeat(const array& arr, int repeats, int axis, StreamOrDevice s = {}); -array repeat(const array& arr, int repeats, StreamOrDevice s = {}); +MLX_API array +repeat(const array& arr, int repeats, int axis, StreamOrDevice s = {}); +MLX_API array repeat(const array& arr, int repeats, StreamOrDevice s = {}); -array tile(const array& arr, std::vector reps, StreamOrDevice s = {}); +MLX_API array +tile(const array& arr, std::vector reps, StreamOrDevice s = {}); /** Permutes the dimensions according to the given axes. */ -array transpose(const array& a, std::vector axes, StreamOrDevice s = {}); +MLX_API array +transpose(const array& a, std::vector axes, StreamOrDevice s = {}); inline array transpose( const array& a, std::initializer_list axes, @@ -269,18 +279,16 @@ inline array transpose( } /** Swap two axes of an array. */ -array swapaxes(const array& a, int axis1, int axis2, StreamOrDevice s = {}); +MLX_API array +swapaxes(const array& a, int axis1, int axis2, StreamOrDevice s = {}); /** Move an axis of an array. */ -array moveaxis( - const array& a, - int source, - int destination, - StreamOrDevice s = {}); +MLX_API array +moveaxis(const array& a, int source, int destination, StreamOrDevice s = {}); /** Pad an array with a constant value */ -array pad( - const array& a, +MLX_API array +pad(const array& a, const std::vector& axes, const Shape& low_pad_size, const Shape& high_pad_size, @@ -289,38 +297,39 @@ array pad( StreamOrDevice s = {}); /** Pad an array with a constant value along all axes */ -array pad( - const array& a, +MLX_API array +pad(const array& a, const std::vector>& pad_width, const array& pad_value = array(0), const std::string& mode = "constant", StreamOrDevice s = {}); -array pad( - const array& a, +MLX_API array +pad(const array& a, const std::pair& pad_width, const array& pad_value = array(0), const std::string& mode = "constant", StreamOrDevice s = {}); -array pad( - const array& a, +MLX_API array +pad(const array& a, int pad_width, const array& pad_value = array(0), const std::string& mode = "constant", StreamOrDevice s = {}); /** Permutes the dimensions in reverse order. */ -array transpose(const array& a, StreamOrDevice s = {}); +MLX_API array transpose(const array& a, StreamOrDevice s = {}); /** Broadcast an array to a given shape. */ -array broadcast_to(const array& a, const Shape& shape, StreamOrDevice s = {}); +MLX_API array +broadcast_to(const array& a, const Shape& shape, StreamOrDevice s = {}); /** Broadcast a vector of arrays against one another. */ -std::vector broadcast_arrays( +MLX_API std::vector broadcast_arrays( const std::vector& inputs, StreamOrDevice s = {}); /** Returns the bool array with (a == b) element-wise. */ -array equal(const array& a, const array& b, StreamOrDevice s = {}); +MLX_API array equal(const array& a, const array& b, StreamOrDevice s = {}); inline array operator==(const array& a, const array& b) { return equal(a, b); } @@ -334,7 +343,7 @@ array operator==(const array& a, T b) { } /** Returns the bool array with (a != b) element-wise. */ -array not_equal(const array& a, const array& b, StreamOrDevice s = {}); +MLX_API array not_equal(const array& a, const array& b, StreamOrDevice s = {}); inline array operator!=(const array& a, const array& b) { return not_equal(a, b); } @@ -348,7 +357,7 @@ array operator!=(const array& a, T b) { } /** Returns bool array with (a > b) element-wise. */ -array greater(const array& a, const array& b, StreamOrDevice s = {}); +MLX_API array greater(const array& a, const array& b, StreamOrDevice s = {}); inline array operator>(const array& a, const array& b) { return greater(a, b); } @@ -362,7 +371,8 @@ array operator>(const array& a, T b) { } /** Returns bool array with (a >= b) element-wise. */ -array greater_equal(const array& a, const array& b, StreamOrDevice s = {}); +MLX_API array +greater_equal(const array& a, const array& b, StreamOrDevice s = {}); inline array operator>=(const array& a, const array& b) { return greater_equal(a, b); } @@ -376,7 +386,7 @@ array operator>=(const array& a, T b) { } /** Returns bool array with (a < b) element-wise. */ -array less(const array& a, const array& b, StreamOrDevice s = {}); +MLX_API array less(const array& a, const array& b, StreamOrDevice s = {}); inline array operator<(const array& a, const array& b) { return less(a, b); } @@ -390,7 +400,7 @@ array operator<(const array& a, T b) { } /** Returns bool array with (a <= b) element-wise. */ -array less_equal(const array& a, const array& b, StreamOrDevice s = {}); +MLX_API array less_equal(const array& a, const array& b, StreamOrDevice s = {}); inline array operator<=(const array& a, const array& b) { return less_equal(a, b); } @@ -404,7 +414,7 @@ array operator<=(const array& a, T b) { } /** True if two arrays have the same shape and elements. */ -array array_equal( +MLX_API array array_equal( const array& a, const array& b, bool equal_nan, @@ -414,25 +424,25 @@ array_equal(const array& a, const array& b, StreamOrDevice s = {}) { return array_equal(a, b, false, s); } -array isnan(const array& a, StreamOrDevice s = {}); +MLX_API array isnan(const array& a, StreamOrDevice s = {}); -array isinf(const array& a, StreamOrDevice s = {}); +MLX_API array isinf(const array& a, StreamOrDevice s = {}); -array isfinite(const array& a, StreamOrDevice s = {}); +MLX_API array isfinite(const array& a, StreamOrDevice s = {}); -array isposinf(const array& a, StreamOrDevice s = {}); +MLX_API array isposinf(const array& a, StreamOrDevice s = {}); -array isneginf(const array& a, StreamOrDevice s = {}); +MLX_API array isneginf(const array& a, StreamOrDevice s = {}); /** Select from x or y depending on condition. */ -array where( +MLX_API array where( const array& condition, const array& x, const array& y, StreamOrDevice s = {}); /** Replace NaN and infinities with finite numbers. */ -array nan_to_num( +MLX_API array nan_to_num( const array& a, float nan = 0.0f, const std::optional posinf = std::nullopt, @@ -440,13 +450,13 @@ array nan_to_num( StreamOrDevice s = {}); /** True if all elements in the array are true (or non-zero). **/ -array all(const array& a, bool keepdims, StreamOrDevice s = {}); +MLX_API array all(const array& a, bool keepdims, StreamOrDevice s = {}); inline array all(const array& a, StreamOrDevice s = {}) { return all(a, false, to_stream(s)); } /** True if the two arrays are equal within the specified tolerance. */ -array allclose( +MLX_API array allclose( const array& a, const array& b, double rtol = 1e-5, @@ -456,7 +466,7 @@ array allclose( /** Returns a boolean array where two arrays are element-wise equal within the * specified tolerance. */ -array isclose( +MLX_API array isclose( const array& a, const array& b, double rtol = 1e-5, @@ -468,8 +478,8 @@ array isclose( * Reduces the input along the given axes. An output value is true * if all the corresponding inputs are true. **/ -array all( - const array& a, +MLX_API array +all(const array& a, const std::vector& axes, bool keepdims = false, StreamOrDevice s = {}); @@ -478,14 +488,11 @@ array all( * Reduces the input along the given axis. An output value is true * if all the corresponding inputs are true. **/ -array all( - const array& a, - int axis, - bool keepdims = false, - StreamOrDevice s = {}); +MLX_API array +all(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {}); /** True if any elements in the array are true (or non-zero). **/ -array any(const array& a, bool keepdims, StreamOrDevice s = {}); +MLX_API array any(const array& a, bool keepdims, StreamOrDevice s = {}); inline array any(const array& a, StreamOrDevice s = {}) { return any(a, false, to_stream(s)); } @@ -494,8 +501,8 @@ inline array any(const array& a, StreamOrDevice s = {}) { * Reduces the input along the given axes. An output value is true * if any of the corresponding inputs are true. **/ -array any( - const array& a, +MLX_API array +any(const array& a, const std::vector& axes, bool keepdims = false, StreamOrDevice s = {}); @@ -504,82 +511,71 @@ array any( * Reduces the input along the given axis. An output value is true * if any of the corresponding inputs are true. **/ -array any( - const array& a, - int axis, - bool keepdims = false, - StreamOrDevice s = {}); +MLX_API array +any(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {}); /** Sums the elements of an array. */ -array sum(const array& a, bool keepdims, StreamOrDevice s = {}); +MLX_API array sum(const array& a, bool keepdims, StreamOrDevice s = {}); inline array sum(const array& a, StreamOrDevice s = {}) { return sum(a, false, to_stream(s)); } /** Sums the elements of an array along the given axes. */ -array sum( - const array& a, +MLX_API array +sum(const array& a, const std::vector& axes, bool keepdims = false, StreamOrDevice s = {}); /** Sums the elements of an array along the given axis. */ -array sum( - const array& a, - int axis, - bool keepdims = false, - StreamOrDevice s = {}); +MLX_API array +sum(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {}); /** Computes the mean of the elements of an array. */ -array mean(const array& a, bool keepdims, StreamOrDevice s = {}); +MLX_API array mean(const array& a, bool keepdims, StreamOrDevice s = {}); inline array mean(const array& a, StreamOrDevice s = {}) { return mean(a, false, to_stream(s)); } /** Computes the mean of the elements of an array along the given axes */ -array mean( +MLX_API array mean( const array& a, const std::vector& axes, bool keepdims = false, StreamOrDevice s = {}); /** Computes the mean of the elements of an array along the given axis */ -array mean( - const array& a, - int axis, - bool keepdims = false, - StreamOrDevice s = {}); +MLX_API array +mean(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {}); /** Computes the median of the elements of an array. */ -array median(const array& a, bool keepdims, StreamOrDevice s = {}); +MLX_API array median(const array& a, bool keepdims, StreamOrDevice s = {}); inline array median(const array& a, StreamOrDevice s = {}) { return median(a, false, to_stream(s)); } /** Computes the median of the elements of an array along the given axes */ -array median( +MLX_API array median( const array& a, const std::vector& axes, bool keepdims = false, StreamOrDevice s = {}); /** Computes the median of the elements of an array along the given axis */ -array median( - const array& a, - int axis, - bool keepdims = false, - StreamOrDevice s = {}); +MLX_API array +median(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {}); /** Computes the variance of the elements of an array. */ -array var(const array& a, bool keepdims, int ddof = 0, StreamOrDevice s = {}); +MLX_API array +var(const array& a, bool keepdims, int ddof = 0, StreamOrDevice s = {}); inline array var(const array& a, StreamOrDevice s = {}) { return var(a, false, 0, to_stream(s)); } /** Computes the variance of the elements of an array along the given * axes */ -array var( - const array& a, +MLX_API array +var(const array& a, const std::vector& axes, bool keepdims = false, int ddof = 0, @@ -587,23 +583,24 @@ array var( /** Computes the variance of the elements of an array along the given * axis */ -array var( - const array& a, +MLX_API array +var(const array& a, int axis, bool keepdims = false, int ddof = 0, StreamOrDevice s = {}); /** Computes the standard deviation of the elements of an array. */ -array std(const array& a, bool keepdims, int ddof = 0, StreamOrDevice s = {}); +MLX_API array +std(const array& a, bool keepdims, int ddof = 0, StreamOrDevice s = {}); inline array std(const array& a, StreamOrDevice s = {}) { return std(a, false, 0, to_stream(s)); } /** Computes the standard deviation of the elements of an array along the given * axes */ -array std( - const array& a, +MLX_API array +std(const array& a, const std::vector& axes, bool keepdims = false, int ddof = 0, @@ -611,150 +608,137 @@ array std( /** Computes the standard deviation of the elements of an array along the given * axis */ -array std( - const array& a, +MLX_API array +std(const array& a, int axis, bool keepdims = false, int ddof = 0, StreamOrDevice s = {}); /** The product of all elements of the array. */ -array prod(const array& a, bool keepdims, StreamOrDevice s = {}); +MLX_API array prod(const array& a, bool keepdims, StreamOrDevice s = {}); inline array prod(const array& a, StreamOrDevice s = {}) { return prod(a, false, to_stream(s)); } /** The product of the elements of an array along the given axes. */ -array prod( +MLX_API array prod( const array& a, const std::vector& axes, bool keepdims = false, StreamOrDevice s = {}); /** The product of the elements of an array along the given axis. */ -array prod( - const array& a, - int axis, - bool keepdims = false, - StreamOrDevice s = {}); +MLX_API array +prod(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {}); /** The maximum of all elements of the array. */ -array max(const array& a, bool keepdims, StreamOrDevice s = {}); +MLX_API array max(const array& a, bool keepdims, StreamOrDevice s = {}); inline array max(const array& a, StreamOrDevice s = {}) { return max(a, false, to_stream(s)); } /** The maximum of the elements of an array along the given axes. */ -array max( - const array& a, +MLX_API array +max(const array& a, const std::vector& axes, bool keepdims = false, StreamOrDevice s = {}); /** The maximum of the elements of an array along the given axis. */ -array max( - const array& a, - int axis, - bool keepdims = false, - StreamOrDevice s = {}); +MLX_API array +max(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {}); /** The minimum of all elements of the array. */ -array min(const array& a, bool keepdims, StreamOrDevice s = {}); +MLX_API array min(const array& a, bool keepdims, StreamOrDevice s = {}); inline array min(const array& a, StreamOrDevice s = {}) { return min(a, false, to_stream(s)); } /** The minimum of the elements of an array along the given axes. */ -array min( - const array& a, +MLX_API array +min(const array& a, const std::vector& axes, bool keepdims = false, StreamOrDevice s = {}); /** The minimum of the elements of an array along the given axis. */ -array min( - const array& a, - int axis, - bool keepdims = false, - StreamOrDevice s = {}); +MLX_API array +min(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {}); /** Returns the index of the minimum value in the array. */ -array argmin(const array& a, bool keepdims, StreamOrDevice s = {}); +MLX_API array argmin(const array& a, bool keepdims, StreamOrDevice s = {}); inline array argmin(const array& a, StreamOrDevice s = {}) { return argmin(a, false, s); } /** Returns the indices of the minimum values along a given axis. */ -array argmin( - const array& a, - int axis, - bool keepdims = false, - StreamOrDevice s = {}); +MLX_API array +argmin(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {}); /** Returns the index of the maximum value in the array. */ -array argmax(const array& a, bool keepdims, StreamOrDevice s = {}); +MLX_API array argmax(const array& a, bool keepdims, StreamOrDevice s = {}); inline array argmax(const array& a, StreamOrDevice s = {}) { return argmax(a, false, s); } /** Returns the indices of the maximum values along a given axis. */ -array argmax( - const array& a, - int axis, - bool keepdims = false, - StreamOrDevice s = {}); +MLX_API array +argmax(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {}); /** Returns a sorted copy of the flattened array. */ -array sort(const array& a, StreamOrDevice s = {}); +MLX_API array sort(const array& a, StreamOrDevice s = {}); /** Returns a sorted copy of the array along a given axis. */ -array sort(const array& a, int axis, StreamOrDevice s = {}); +MLX_API array sort(const array& a, int axis, StreamOrDevice s = {}); /** Returns indices that sort the flattened array. */ -array argsort(const array& a, StreamOrDevice s = {}); +MLX_API array argsort(const array& a, StreamOrDevice s = {}); /** Returns indices that sort the array along a given axis. */ -array argsort(const array& a, int axis, StreamOrDevice s = {}); +MLX_API array argsort(const array& a, int axis, StreamOrDevice s = {}); /** * Returns a partitioned copy of the flattened array * such that the smaller kth elements are first. **/ -array partition(const array& a, int kth, StreamOrDevice s = {}); +MLX_API array partition(const array& a, int kth, StreamOrDevice s = {}); /** * Returns a partitioned copy of the array along a given axis * such that the smaller kth elements are first. **/ -array partition(const array& a, int kth, int axis, StreamOrDevice s = {}); +MLX_API array +partition(const array& a, int kth, int axis, StreamOrDevice s = {}); /** * Returns indices that partition the flattened array * such that the smaller kth elements are first. **/ -array argpartition(const array& a, int kth, StreamOrDevice s = {}); +MLX_API array argpartition(const array& a, int kth, StreamOrDevice s = {}); /** * Returns indices that partition the array along a given axis * such that the smaller kth elements are first. **/ -array argpartition(const array& a, int kth, int axis, StreamOrDevice s = {}); +MLX_API array +argpartition(const array& a, int kth, int axis, StreamOrDevice s = {}); /** Returns topk elements of the flattened array. */ -array topk(const array& a, int k, StreamOrDevice s = {}); +MLX_API array topk(const array& a, int k, StreamOrDevice s = {}); /** Returns topk elements of the array along a given axis. */ -array topk(const array& a, int k, int axis, StreamOrDevice s = {}); +MLX_API array topk(const array& a, int k, int axis, StreamOrDevice s = {}); /** Cumulative logsumexp of an array. */ -array logcumsumexp( +MLX_API array logcumsumexp( const array& a, bool reverse = false, bool inclusive = true, StreamOrDevice s = {}); /** Cumulative logsumexp of an array along the given axis. */ -array logcumsumexp( +MLX_API array logcumsumexp( const array& a, int axis, bool reverse = false, @@ -762,52 +746,53 @@ array logcumsumexp( StreamOrDevice s = {}); /** The logsumexp of all elements of the array. */ -array logsumexp(const array& a, bool keepdims, StreamOrDevice s = {}); +MLX_API array logsumexp(const array& a, bool keepdims, StreamOrDevice s = {}); inline array logsumexp(const array& a, StreamOrDevice s = {}) { return logsumexp(a, false, to_stream(s)); } /** The logsumexp of the elements of an array along the given axes. */ -array logsumexp( +MLX_API array logsumexp( const array& a, const std::vector& axes, bool keepdims = false, StreamOrDevice s = {}); /** The logsumexp of the elements of an array along the given axis. */ -array logsumexp( +MLX_API array logsumexp( const array& a, int axis, bool keepdims = false, StreamOrDevice s = {}); /** Absolute value of elements in an array. */ -array abs(const array& a, StreamOrDevice s = {}); +MLX_API array abs(const array& a, StreamOrDevice s = {}); /** Negate an array. */ -array negative(const array& a, StreamOrDevice s = {}); -array operator-(const array& a); +MLX_API array negative(const array& a, StreamOrDevice s = {}); +MLX_API array operator-(const array& a); /** The sign of the elements in an array. */ -array sign(const array& a, StreamOrDevice s = {}); +MLX_API array sign(const array& a, StreamOrDevice s = {}); /** Logical not of an array */ -array logical_not(const array& a, StreamOrDevice s = {}); +MLX_API array logical_not(const array& a, StreamOrDevice s = {}); /** Logical and of two arrays */ -array logical_and(const array& a, const array& b, StreamOrDevice s = {}); -array operator&&(const array& a, const array& b); +MLX_API array +logical_and(const array& a, const array& b, StreamOrDevice s = {}); +MLX_API array operator&&(const array& a, const array& b); /** Logical or of two arrays */ -array logical_or(const array& a, const array& b, StreamOrDevice s = {}); -array operator||(const array& a, const array& b); +MLX_API array logical_or(const array& a, const array& b, StreamOrDevice s = {}); +MLX_API array operator||(const array& a, const array& b); /** The reciprocal (1/x) of the elements in an array. */ -array reciprocal(const array& a, StreamOrDevice s = {}); +MLX_API array reciprocal(const array& a, StreamOrDevice s = {}); /** Add two arrays. */ -array add(const array& a, const array& b, StreamOrDevice s = {}); -array operator+(const array& a, const array& b); +MLX_API array add(const array& a, const array& b, StreamOrDevice s = {}); +MLX_API array operator+(const array& a, const array& b); template array operator+(T a, const array& b) { return add(array(a), b); @@ -818,8 +803,8 @@ array operator+(const array& a, T b) { } /** Subtract two arrays. */ -array subtract(const array& a, const array& b, StreamOrDevice s = {}); -array operator-(const array& a, const array& b); +MLX_API array subtract(const array& a, const array& b, StreamOrDevice s = {}); +MLX_API array operator-(const array& a, const array& b); template array operator-(T a, const array& b) { return subtract(array(a), b); @@ -830,8 +815,8 @@ array operator-(const array& a, T b) { } /** Multiply two arrays. */ -array multiply(const array& a, const array& b, StreamOrDevice s = {}); -array operator*(const array& a, const array& b); +MLX_API array multiply(const array& a, const array& b, StreamOrDevice s = {}); +MLX_API array operator*(const array& a, const array& b); template array operator*(T a, const array& b) { return multiply(array(a), b); @@ -842,21 +827,22 @@ array operator*(const array& a, T b) { } /** Divide two arrays. */ -array divide(const array& a, const array& b, StreamOrDevice s = {}); -array operator/(const array& a, const array& b); -array operator/(double a, const array& b); -array operator/(const array& a, double b); +MLX_API array divide(const array& a, const array& b, StreamOrDevice s = {}); +MLX_API array operator/(const array& a, const array& b); +MLX_API array operator/(double a, const array& b); +MLX_API array operator/(const array& a, double b); /** Compute the element-wise quotient and remainder. */ -std::vector +MLX_API std::vector divmod(const array& a, const array& b, StreamOrDevice s = {}); /** Compute integer division. Equivalent to doing floor(a / x). */ -array floor_divide(const array& a, const array& b, StreamOrDevice s = {}); +MLX_API array +floor_divide(const array& a, const array& b, StreamOrDevice s = {}); /** Compute the element-wise remainder of division */ -array remainder(const array& a, const array& b, StreamOrDevice s = {}); -array operator%(const array& a, const array& b); +MLX_API array remainder(const array& a, const array& b, StreamOrDevice s = {}); +MLX_API array operator%(const array& a, const array& b); template array operator%(T a, const array& b) { return remainder(array(a), b); @@ -867,109 +853,109 @@ array operator%(const array& a, T b) { } /** Element-wise maximum between two arrays. */ -array maximum(const array& a, const array& b, StreamOrDevice s = {}); +MLX_API array maximum(const array& a, const array& b, StreamOrDevice s = {}); /** Element-wise minimum between two arrays. */ -array minimum(const array& a, const array& b, StreamOrDevice s = {}); +MLX_API array minimum(const array& a, const array& b, StreamOrDevice s = {}); /** Floor the element of an array. **/ -array floor(const array& a, StreamOrDevice s = {}); +MLX_API array floor(const array& a, StreamOrDevice s = {}); /** Ceil the element of an array. **/ -array ceil(const array& a, StreamOrDevice s = {}); +MLX_API array ceil(const array& a, StreamOrDevice s = {}); /** Square the elements of an array. */ -array square(const array& a, StreamOrDevice s = {}); +MLX_API array square(const array& a, StreamOrDevice s = {}); /** Exponential of the elements of an array. */ -array exp(const array& a, StreamOrDevice s = {}); +MLX_API array exp(const array& a, StreamOrDevice s = {}); /** Sine of the elements of an array */ -array sin(const array& a, StreamOrDevice s = {}); +MLX_API array sin(const array& a, StreamOrDevice s = {}); /** Cosine of the elements of an array */ -array cos(const array& a, StreamOrDevice s = {}); +MLX_API array cos(const array& a, StreamOrDevice s = {}); /** Tangent of the elements of an array */ -array tan(const array& a, StreamOrDevice s = {}); +MLX_API array tan(const array& a, StreamOrDevice s = {}); /** Arc Sine of the elements of an array */ -array arcsin(const array& a, StreamOrDevice s = {}); +MLX_API array arcsin(const array& a, StreamOrDevice s = {}); /** Arc Cosine of the elements of an array */ -array arccos(const array& a, StreamOrDevice s = {}); +MLX_API array arccos(const array& a, StreamOrDevice s = {}); /** Arc Tangent of the elements of an array */ -array arctan(const array& a, StreamOrDevice s = {}); +MLX_API array arctan(const array& a, StreamOrDevice s = {}); /** Inverse tangent of the ratio of two arrays */ -array arctan2(const array& a, const array& b, StreamOrDevice s = {}); +MLX_API array arctan2(const array& a, const array& b, StreamOrDevice s = {}); /** Hyperbolic Sine of the elements of an array */ -array sinh(const array& a, StreamOrDevice s = {}); +MLX_API array sinh(const array& a, StreamOrDevice s = {}); /** Hyperbolic Cosine of the elements of an array */ -array cosh(const array& a, StreamOrDevice s = {}); +MLX_API array cosh(const array& a, StreamOrDevice s = {}); /** Hyperbolic Tangent of the elements of an array */ -array tanh(const array& a, StreamOrDevice s = {}); +MLX_API array tanh(const array& a, StreamOrDevice s = {}); /** Inverse Hyperbolic Sine of the elements of an array */ -array arcsinh(const array& a, StreamOrDevice s = {}); +MLX_API array arcsinh(const array& a, StreamOrDevice s = {}); /** Inverse Hyperbolic Cosine of the elements of an array */ -array arccosh(const array& a, StreamOrDevice s = {}); +MLX_API array arccosh(const array& a, StreamOrDevice s = {}); /** Inverse Hyperbolic Tangent of the elements of an array */ -array arctanh(const array& a, StreamOrDevice s = {}); +MLX_API array arctanh(const array& a, StreamOrDevice s = {}); /** Convert the elements of an array from Radians to Degrees **/ -array degrees(const array& a, StreamOrDevice s = {}); +MLX_API array degrees(const array& a, StreamOrDevice s = {}); /** Convert the elements of an array from Degrees to Radians **/ -array radians(const array& a, StreamOrDevice s = {}); +MLX_API array radians(const array& a, StreamOrDevice s = {}); /** Natural logarithm of the elements of an array. */ -array log(const array& a, StreamOrDevice s = {}); +MLX_API array log(const array& a, StreamOrDevice s = {}); /** Log base 2 of the elements of an array. */ -array log2(const array& a, StreamOrDevice s = {}); +MLX_API array log2(const array& a, StreamOrDevice s = {}); /** Log base 10 of the elements of an array. */ -array log10(const array& a, StreamOrDevice s = {}); +MLX_API array log10(const array& a, StreamOrDevice s = {}); /** Natural logarithm of one plus elements in the array: `log(1 + a)`. */ -array log1p(const array& a, StreamOrDevice s = {}); +MLX_API array log1p(const array& a, StreamOrDevice s = {}); /** Log-add-exp of one elements in the array: `log(exp(a) + exp(b))`. */ -array logaddexp(const array& a, const array& b, StreamOrDevice s = {}); +MLX_API array logaddexp(const array& a, const array& b, StreamOrDevice s = {}); /** Element-wise logistic sigmoid of the array: `1 / (1 + exp(-x)`. */ -array sigmoid(const array& a, StreamOrDevice s = {}); +MLX_API array sigmoid(const array& a, StreamOrDevice s = {}); /** Computes the error function of the elements of an array. */ -array erf(const array& a, StreamOrDevice s = {}); +MLX_API array erf(const array& a, StreamOrDevice s = {}); /** Computes the inverse error function of the elements of an array. */ -array erfinv(const array& a, StreamOrDevice s = {}); +MLX_API array erfinv(const array& a, StreamOrDevice s = {}); /** Computes the expm1 function of the elements of an array. */ -array expm1(const array& a, StreamOrDevice s = {}); +MLX_API array expm1(const array& a, StreamOrDevice s = {}); /** Stop the flow of gradients. */ -array stop_gradient(const array& a, StreamOrDevice s = {}); +MLX_API array stop_gradient(const array& a, StreamOrDevice s = {}); /** Round a floating point number */ -array round(const array& a, int decimals, StreamOrDevice s = {}); +MLX_API array round(const array& a, int decimals, StreamOrDevice s = {}); inline array round(const array& a, StreamOrDevice s = {}) { return round(a, 0, s); } /** Matrix-matrix multiplication. */ -array matmul(const array& a, const array& b, StreamOrDevice s = {}); +MLX_API array matmul(const array& a, const array& b, StreamOrDevice s = {}); /** Gather array entries given indices and slices */ -array gather( +MLX_API array gather( const array& a, const std::vector& indices, const std::vector& axes, @@ -985,29 +971,26 @@ inline array gather( } /** Compute the Kronecker product of two arrays. */ -array kron(const array& a, const array& b, StreamOrDevice s = {}); +MLX_API array kron(const array& a, const array& b, StreamOrDevice s = {}); /** Take array slices at the given indices of the specified axis. */ -array take( - const array& a, - const array& indices, - int axis, - StreamOrDevice s = {}); -array take(const array& a, int index, int axis, StreamOrDevice s = {}); +MLX_API array +take(const array& a, const array& indices, int axis, StreamOrDevice s = {}); +MLX_API array take(const array& a, int index, int axis, StreamOrDevice s = {}); /** Take array entries at the given indices treating the array as flattened. */ -array take(const array& a, const array& indices, StreamOrDevice s = {}); -array take(const array& a, int index, StreamOrDevice s = {}); +MLX_API array take(const array& a, const array& indices, StreamOrDevice s = {}); +MLX_API array take(const array& a, int index, StreamOrDevice s = {}); /** Take array entries given indices along the axis */ -array take_along_axis( +MLX_API array take_along_axis( const array& a, const array& indices, int axis, StreamOrDevice s = {}); /** Put the values into the array at the given indices along the axis */ -array put_along_axis( +MLX_API array put_along_axis( const array& a, const array& indices, const array& values, @@ -1015,7 +998,7 @@ array put_along_axis( StreamOrDevice s = {}); /** Add the values into the array at the given indices along the axis */ -array scatter_add_axis( +MLX_API array scatter_add_axis( const array& a, const array& indices, const array& values, @@ -1120,7 +1103,7 @@ array scatter_add_axis( * updates. Out-of-bounds accesses on ``a`` are undefined and typically result * in unintended or invalid memory writes. */ -array scatter( +MLX_API array scatter( const array& a, const std::vector& indices, const array& updates, @@ -1136,7 +1119,7 @@ inline array scatter( } /** Scatter and add updates to given indices */ -array scatter_add( +MLX_API array scatter_add( const array& a, const std::vector& indices, const array& updates, @@ -1152,7 +1135,7 @@ inline array scatter_add( } /** Scatter and prod updates to given indices */ -array scatter_prod( +MLX_API array scatter_prod( const array& a, const std::vector& indices, const array& updates, @@ -1168,7 +1151,7 @@ inline array scatter_prod( } /** Scatter and max updates to given linear indices */ -array scatter_max( +MLX_API array scatter_max( const array& a, const std::vector& indices, const array& updates, @@ -1183,7 +1166,7 @@ inline array scatter_max( return scatter_max(a, {indices}, updates, std::vector{axis}, s); } /** Scatter and min updates to given linear indices */ -array scatter_min( +MLX_API array scatter_min( const array& a, const std::vector& indices, const array& updates, @@ -1198,27 +1181,28 @@ inline array scatter_min( return scatter_min(a, {indices}, updates, std::vector{axis}, s); } -array masked_scatter( +MLX_API array masked_scatter( const array& a, const array& mask, const array& src, StreamOrDevice s = {}); /** Square root the elements of an array. */ -array sqrt(const array& a, StreamOrDevice s = {}); +MLX_API array sqrt(const array& a, StreamOrDevice s = {}); /** Square root and reciprocal the elements of an array. */ -array rsqrt(const array& a, StreamOrDevice s = {}); +MLX_API array rsqrt(const array& a, StreamOrDevice s = {}); /** Softmax of an array. */ -array softmax( +MLX_API array softmax( const array& a, const std::vector& axes, bool precise = false, StreamOrDevice s = {}); /** Softmax of an array. */ -array softmax(const array& a, bool precise = false, StreamOrDevice s = {}); +MLX_API array +softmax(const array& a, bool precise = false, StreamOrDevice s = {}); /** Softmax of an array. */ inline array @@ -1227,17 +1211,17 @@ softmax(const array& a, int axis, bool precise = false, StreamOrDevice s = {}) { } /** Raise elements of a to the power of b element-wise */ -array power(const array& a, const array& b, StreamOrDevice s = {}); +MLX_API array power(const array& a, const array& b, StreamOrDevice s = {}); /** Cumulative sum of an array. */ -array cumsum( +MLX_API array cumsum( const array& a, bool reverse = false, bool inclusive = true, StreamOrDevice s = {}); /** Cumulative sum of an array along the given axis. */ -array cumsum( +MLX_API array cumsum( const array& a, int axis, bool reverse = false, @@ -1245,14 +1229,14 @@ array cumsum( StreamOrDevice s = {}); /** Cumulative product of an array. */ -array cumprod( +MLX_API array cumprod( const array& a, bool reverse = false, bool inclusive = true, StreamOrDevice s = {}); /** Cumulative product of an array along the given axis. */ -array cumprod( +MLX_API array cumprod( const array& a, int axis, bool reverse = false, @@ -1260,14 +1244,14 @@ array cumprod( StreamOrDevice s = {}); /** Cumulative max of an array. */ -array cummax( +MLX_API array cummax( const array& a, bool reverse = false, bool inclusive = true, StreamOrDevice s = {}); /** Cumulative max of an array along the given axis. */ -array cummax( +MLX_API array cummax( const array& a, int axis, bool reverse = false, @@ -1275,14 +1259,14 @@ array cummax( StreamOrDevice s = {}); /** Cumulative min of an array. */ -array cummin( +MLX_API array cummin( const array& a, bool reverse = false, bool inclusive = true, StreamOrDevice s = {}); /** Cumulative min of an array along the given axis. */ -array cummin( +MLX_API array cummin( const array& a, int axis, bool reverse = false, @@ -1290,7 +1274,7 @@ array cummin( StreamOrDevice s = {}); /** General convolution with a filter */ -array conv_general( +MLX_API array conv_general( array input, array weight, std::vector stride = {}, @@ -1327,7 +1311,7 @@ inline array conv_general( } /** 1D convolution with a filter */ -array conv1d( +MLX_API array conv1d( const array& input, const array& weight, int stride = 1, @@ -1337,7 +1321,7 @@ array conv1d( StreamOrDevice s = {}); /** 2D convolution with a filter */ -array conv2d( +MLX_API array conv2d( const array& input, const array& weight, const std::pair& stride = {1, 1}, @@ -1347,7 +1331,7 @@ array conv2d( StreamOrDevice s = {}); /** 3D convolution with a filter */ -array conv3d( +MLX_API array conv3d( const array& input, const array& weight, const std::tuple& stride = {1, 1, 1}, @@ -1357,7 +1341,7 @@ array conv3d( StreamOrDevice s = {}); /** 1D transposed convolution with a filter */ -array conv_transpose1d( +MLX_API array conv_transpose1d( const array& input, const array& weight, int stride = 1, @@ -1368,7 +1352,7 @@ array conv_transpose1d( StreamOrDevice s = {}); /** 2D transposed convolution with a filter */ -array conv_transpose2d( +MLX_API array conv_transpose2d( const array& input, const array& weight, const std::pair& stride = {1, 1}, @@ -1379,7 +1363,7 @@ array conv_transpose2d( StreamOrDevice s = {}); /** 3D transposed convolution with a filter */ -array conv_transpose3d( +MLX_API array conv_transpose3d( const array& input, const array& weight, const std::tuple& stride = {1, 1, 1}, @@ -1390,7 +1374,7 @@ array conv_transpose3d( StreamOrDevice s = {}); /** Quantized matmul multiplies x with a quantized matrix w*/ -array quantized_matmul( +MLX_API array quantized_matmul( array x, array w, array scales, @@ -1402,7 +1386,7 @@ array quantized_matmul( StreamOrDevice s = {}); /** Quantize a matrix along its last axis */ -std::vector quantize( +MLX_API std::vector quantize( const array& w, std::optional group_size = std::nullopt, std::optional bits = std::nullopt, @@ -1410,7 +1394,7 @@ std::vector quantize( StreamOrDevice s = {}); /** Dequantize a matrix produced by quantize() */ -array dequantize( +MLX_API array dequantize( const array& w, const array& scales, const std::optional& biases = std::nullopt, @@ -1420,7 +1404,7 @@ array dequantize( std::optional dtype = std::nullopt, StreamOrDevice s = {}); -array qqmm( +MLX_API array qqmm( array x, // input activations array w, // maybe quantized weights std::optional w_scales = std::nullopt, // optional scales if w is @@ -1431,13 +1415,13 @@ array qqmm( StreamOrDevice s = {}); /** Convert an E4M3 float8 to the given floating point dtype. */ -array from_fp8(array x, Dtype dtype, StreamOrDevice s = {}); +MLX_API array from_fp8(array x, Dtype dtype, StreamOrDevice s = {}); /** Convert a floating point matrix to E4M3 float8. */ -array to_fp8(array x, StreamOrDevice s = {}); +MLX_API array to_fp8(array x, StreamOrDevice s = {}); /** Compute matrix products with matrix-level gather. */ -array gather_qmm( +MLX_API array gather_qmm( const array& x, const array& w, const array& scales, @@ -1452,13 +1436,13 @@ array gather_qmm( StreamOrDevice s = {}); /** Returns a contraction of a and b over multiple dimensions. */ -array tensordot( +MLX_API array tensordot( const array& a, const array& b, const int axis = 2, StreamOrDevice s = {}); -array tensordot( +MLX_API array tensordot( const array& a, const array& b, const std::vector& axes_a, @@ -1466,13 +1450,13 @@ array tensordot( StreamOrDevice s = {}); /** Compute the outer product of two vectors. */ -array outer(const array& a, const array& b, StreamOrDevice s = {}); +MLX_API array outer(const array& a, const array& b, StreamOrDevice s = {}); /** Compute the inner product of two vectors. */ -array inner(const array& a, const array& b, StreamOrDevice s = {}); +MLX_API array inner(const array& a, const array& b, StreamOrDevice s = {}); /** Compute D = beta * C + alpha * (A @ B) */ -array addmm( +MLX_API array addmm( array c, array a, array b, @@ -1481,7 +1465,7 @@ array addmm( StreamOrDevice s = {}); /** Compute matrix product with block masking */ -array block_masked_mm( +MLX_API array block_masked_mm( array a, array b, int block_size, @@ -1491,7 +1475,7 @@ array block_masked_mm( StreamOrDevice s = {}); /** Compute matrix product with matrix-level gather */ -array gather_mm( +MLX_API array gather_mm( array a, array b, std::optional lhs_indices = std::nullopt, @@ -1503,10 +1487,11 @@ array gather_mm( * Compute a matrix product but segment the inner dimension and write the * result separately for each segment. */ -array segmented_mm(array a, array b, array segments, StreamOrDevice s = {}); +MLX_API array +segmented_mm(array a, array b, array segments, StreamOrDevice s = {}); /** Extract a diagonal or construct a diagonal array */ -array diagonal( +MLX_API array diagonal( const array& a, int offset = 0, int axis1 = 0, @@ -1514,44 +1499,40 @@ array diagonal( StreamOrDevice s = {}); /** Extract diagonal from a 2d array or create a diagonal matrix. */ -array diag(const array& a, int k = 0, StreamOrDevice s = {}); +MLX_API array diag(const array& a, int k = 0, StreamOrDevice s = {}); /** Return the sum along a specified diagonal in the given array. */ -array trace( +MLX_API array trace( const array& a, int offset, int axis1, int axis2, Dtype dtype, StreamOrDevice s = {}); -array trace( - const array& a, - int offset, - int axis1, - int axis2, - StreamOrDevice s = {}); -array trace(const array& a, StreamOrDevice s = {}); +MLX_API array +trace(const array& a, int offset, int axis1, int axis2, StreamOrDevice s = {}); +MLX_API array trace(const array& a, StreamOrDevice s = {}); /** * Implements the identity function but allows injecting dependencies to other * arrays. This ensures that these other arrays will have been computed * when the outputs of this function are computed. */ -std::vector depends( +MLX_API std::vector depends( const std::vector& inputs, const std::vector& dependencies); /** convert an array to an atleast ndim array */ -array atleast_1d(const array& a, StreamOrDevice s = {}); -std::vector atleast_1d( +MLX_API array atleast_1d(const array& a, StreamOrDevice s = {}); +MLX_API std::vector atleast_1d( const std::vector& a, StreamOrDevice s = {}); -array atleast_2d(const array& a, StreamOrDevice s = {}); -std::vector atleast_2d( +MLX_API array atleast_2d(const array& a, StreamOrDevice s = {}); +MLX_API std::vector atleast_2d( const std::vector& a, StreamOrDevice s = {}); -array atleast_3d(const array& a, StreamOrDevice s = {}); -std::vector atleast_3d( +MLX_API array atleast_3d(const array& a, StreamOrDevice s = {}); +MLX_API std::vector atleast_3d( const std::vector& a, StreamOrDevice s = {}); @@ -1559,68 +1540,70 @@ std::vector atleast_3d( * Extract the number of elements along some axes as a scalar array. Used to * allow shape dependent shapeless compilation (pun intended). */ -array number_of_elements( +MLX_API array number_of_elements( const array& a, std::vector axes, bool inverted, Dtype dtype = int32, StreamOrDevice s = {}); -array conjugate(const array& a, StreamOrDevice s = {}); +MLX_API array conjugate(const array& a, StreamOrDevice s = {}); /** Bitwise and. */ -array bitwise_and(const array& a, const array& b, StreamOrDevice s = {}); -array operator&(const array& a, const array& b); +MLX_API array +bitwise_and(const array& a, const array& b, StreamOrDevice s = {}); +MLX_API array operator&(const array& a, const array& b); /** Bitwise inclusive or. */ -array bitwise_or(const array& a, const array& b, StreamOrDevice s = {}); -array operator|(const array& a, const array& b); +MLX_API array bitwise_or(const array& a, const array& b, StreamOrDevice s = {}); +MLX_API array operator|(const array& a, const array& b); /** Bitwise exclusive or. */ -array bitwise_xor(const array& a, const array& b, StreamOrDevice s = {}); -array operator^(const array& a, const array& b); +MLX_API array +bitwise_xor(const array& a, const array& b, StreamOrDevice s = {}); +MLX_API array operator^(const array& a, const array& b); /** Shift bits to the left. */ -array left_shift(const array& a, const array& b, StreamOrDevice s = {}); -array operator<<(const array& a, const array& b); +MLX_API array left_shift(const array& a, const array& b, StreamOrDevice s = {}); +MLX_API array operator<<(const array& a, const array& b); /** Shift bits to the right. */ -array right_shift(const array& a, const array& b, StreamOrDevice s = {}); -array operator>>(const array& a, const array& b); +MLX_API array +right_shift(const array& a, const array& b, StreamOrDevice s = {}); +MLX_API array operator>>(const array& a, const array& b); /** Invert the bits. */ -array bitwise_invert(const array& a, StreamOrDevice s = {}); -array operator~(const array& a); +MLX_API array bitwise_invert(const array& a, StreamOrDevice s = {}); +MLX_API array operator~(const array& a); -array view(const array& a, const Dtype& dtype, StreamOrDevice s = {}); +MLX_API array view(const array& a, const Dtype& dtype, StreamOrDevice s = {}); /** Roll elements along an axis and introduce them on the other side */ -array roll(const array& a, int shift, StreamOrDevice s = {}); -array roll(const array& a, const Shape& shift, StreamOrDevice s = {}); -array roll(const array& a, int shift, int axis, StreamOrDevice s = {}); -array roll( +MLX_API array roll(const array& a, int shift, StreamOrDevice s = {}); +MLX_API array roll(const array& a, const Shape& shift, StreamOrDevice s = {}); +MLX_API array roll(const array& a, int shift, int axis, StreamOrDevice s = {}); +MLX_API array roll( const array& a, int shift, const std::vector& axes, StreamOrDevice s = {}); -array roll(const array& a, const Shape& shift, int axis, StreamOrDevice s = {}); -array roll( +MLX_API array +roll(const array& a, const Shape& shift, int axis, StreamOrDevice s = {}); +MLX_API array roll( const array& a, const Shape& shift, const std::vector& axes, StreamOrDevice s = {}); /* The real part of a complex array. */ -array real(const array& a, StreamOrDevice s = {}); +MLX_API array real(const array& a, StreamOrDevice s = {}); /* The imaginary part of a complex array. */ -array imag(const array& a, StreamOrDevice s = {}); +MLX_API array imag(const array& a, StreamOrDevice s = {}); /* Ensure the array's underlying memory is contiguous. */ -array contiguous( - const array& a, - bool allow_col_major = false, - StreamOrDevice s = {}); +MLX_API array +contiguous(const array& a, bool allow_col_major = false, StreamOrDevice s = {}); /** @} */ diff --git a/mlx/primitives.h b/mlx/primitives.h index c3ce00f9..4971ea9c 100644 --- a/mlx/primitives.h +++ b/mlx/primitives.h @@ -4,6 +4,7 @@ #include +#include "mlx/api.h" #include "mlx/array.h" #include "mlx/device.h" #include "mlx/io/load.h" @@ -123,7 +124,7 @@ class Primitive { Stream stream_; }; -class UnaryPrimitive : public Primitive { +class MLX_API UnaryPrimitive : public Primitive { /** * An abstract base class for a primitive with a single output. */ @@ -349,7 +350,7 @@ class ArgPartition : public UnaryPrimitive { int axis_; }; -class ArgReduce : public UnaryPrimitive { +class MLX_API ArgReduce : public UnaryPrimitive { public: enum ReduceType { ArgMin, diff --git a/mlx/random.h b/mlx/random.h index b6837418..a23c2557 100644 --- a/mlx/random.h +++ b/mlx/random.h @@ -5,13 +5,14 @@ #include #include +#include "mlx/api.h" #include "mlx/array.h" #include "mlx/stream.h" #include "mlx/utils.h" namespace mlx::core::random { -class KeySequence { +class MLX_API KeySequence { public: explicit KeySequence(uint64_t seed); @@ -35,13 +36,13 @@ class KeySequence { }; /** Get a PRNG key from a seed. */ -array key(uint64_t seed); +MLX_API array key(uint64_t seed); /** Seed the default PRNG key. */ -void seed(uint64_t seed); +MLX_API void seed(uint64_t seed); /** Generate an array with type uint32 filled with random bits. */ -array bits( +MLX_API array bits( const Shape& shape, int width, const std::optional& key = std::nullopt, @@ -54,13 +55,13 @@ inline array bits( } /** Split the rng key into a pair of keys. */ -std::pair split(const array& key, StreamOrDevice s = {}); +MLX_API std::pair split(const array& key, StreamOrDevice s = {}); /** Split the rng key into `num` keys. */ -array split(const array& key, int num, StreamOrDevice s = {}); +MLX_API array split(const array& key, int num, StreamOrDevice s = {}); /** Generate uniform random numbers between low and high. */ -array uniform( +MLX_API array uniform( const array& low, const array& high, const Shape& shape, @@ -80,7 +81,7 @@ array uniform( } /** Generate uniform random numbers between 0 and 1. */ -array uniform( +MLX_API array uniform( const Shape& shape, Dtype dtype, const std::optional& key = std::nullopt, @@ -93,7 +94,7 @@ inline array uniform( } /** Generate samples from the standard normal distribution. */ -array normal( +MLX_API array normal( const Shape& shape, Dtype dtype, const std::optional& loc, @@ -135,7 +136,7 @@ inline array normal( } /** Generate samples from a multivariate normal distribution. **/ -array multivariate_normal( +MLX_API array multivariate_normal( const array& mean, const array& cov, const Shape& shape, @@ -144,7 +145,7 @@ array multivariate_normal( StreamOrDevice s = {}); /** Generate integer samples uniformly at random */ -array randint( +MLX_API array randint( const array& low, const array& high, const Shape& shape, @@ -164,12 +165,12 @@ array randint( } /** Generate binary variables with probability to be true equal to p */ -array bernoulli( +MLX_API array bernoulli( const array& p, const Shape& shape, const std::optional& key = std::nullopt, StreamOrDevice s = {}); -array bernoulli( +MLX_API array bernoulli( const array& p, const std::optional& key = std::nullopt, StreamOrDevice s = {}); @@ -191,11 +192,11 @@ array bernoulli( return bernoulli(array(p), shape, key, s); } -array bernoulli( +MLX_API array bernoulli( const std::optional& key = std::nullopt, StreamOrDevice s = {}); -array truncated_normal( +MLX_API array truncated_normal( const array& lower, const array& upper, const Shape& shape, @@ -203,41 +204,41 @@ array truncated_normal( const std::optional& key = std::nullopt, StreamOrDevice s = {}); -array truncated_normal( +MLX_API array truncated_normal( const array& lower, const array& upper, Dtype dtype = float32, const std::optional& key = std::nullopt, StreamOrDevice s = {}); -array gumbel( +MLX_API array gumbel( const Shape& shape, Dtype dtype = float32, const std::optional& key = std::nullopt, StreamOrDevice s = {}); -array categorical( +MLX_API array categorical( const array& logits, int axis, const Shape& shape, const std::optional& key = std::nullopt, StreamOrDevice s = {}); -array categorical( +MLX_API array categorical( const array& logits_, int axis, int num_samples, const std::optional& key = std::nullopt, StreamOrDevice s = {}); -array categorical( +MLX_API array categorical( const array& logits, int axis = -1, const std::optional& key = std::nullopt, StreamOrDevice s = {}); /** Generate samples from the laplace distribution. */ -array laplace( +MLX_API array laplace( const Shape& shape, Dtype dtype, const float loc, @@ -267,14 +268,14 @@ inline array laplace( } /* Randomly permute the elements of x along the given axis. */ -array permutation( +MLX_API array permutation( const array& x, int axis = 0, const std::optional& key = std::nullopt, StreamOrDevice s = {}); /* A random permutation of `arange(x)` */ -array permutation( +MLX_API array permutation( int x, const std::optional& key = std::nullopt, StreamOrDevice s = {}); diff --git a/mlx/scheduler.h b/mlx/scheduler.h index d01d4140..4b949459 100644 --- a/mlx/scheduler.h +++ b/mlx/scheduler.h @@ -8,6 +8,7 @@ #include #include +#include "mlx/api.h" #include "mlx/backend/gpu/eval.h" #include "mlx/device.h" #include "mlx/stream.h" @@ -162,7 +163,7 @@ void Scheduler::enqueue(const Stream& stream, F&& f) { threads_[stream.index]->enqueue(std::forward(f)); } -Scheduler& scheduler(); +MLX_API Scheduler& scheduler(); template void enqueue(const Stream& stream, F&& f) { diff --git a/mlx/stream.h b/mlx/stream.h index 3ced4039..9bd14db2 100644 --- a/mlx/stream.h +++ b/mlx/stream.h @@ -2,27 +2,28 @@ #pragma once +#include "mlx/api.h" #include "mlx/device.h" namespace mlx::core { -struct Stream { +struct MLX_API Stream { int index; Device device; explicit Stream(int index, Device device) : index(index), device(device) {} }; /** Get the default stream for the given device. */ -Stream default_stream(Device d); +MLX_API Stream default_stream(Device d); /** Make the stream the default for its device. */ -void set_default_stream(Stream s); +MLX_API void set_default_stream(Stream s); /** Make a new stream on the given device. */ -Stream new_stream(Device d); +MLX_API Stream new_stream(Device d); /** Get the stream with the given index. */ -Stream get_stream(int index); +MLX_API Stream get_stream(int index); inline bool operator==(const Stream& lhs, const Stream& rhs) { return lhs.index == rhs.index; @@ -33,9 +34,9 @@ inline bool operator!=(const Stream& lhs, const Stream& rhs) { } /* Synchronize with the default stream. */ -void synchronize(); +MLX_API void synchronize(); /* Synchronize with the provided stream. */ -void synchronize(Stream); +MLX_API void synchronize(Stream); } // namespace mlx::core diff --git a/mlx/transforms.h b/mlx/transforms.h index 4afb21e2..1848be79 100644 --- a/mlx/transforms.h +++ b/mlx/transforms.h @@ -4,18 +4,19 @@ #include +#include "mlx/api.h" #include "mlx/array.h" namespace mlx::core { -void async_eval(std::vector outputs); +MLX_API void async_eval(std::vector outputs); template > void async_eval(Arrays&&... outputs) { async_eval(std::vector{std::forward(outputs)...}); } -void eval(std::vector outputs); +MLX_API void eval(std::vector outputs); template > void eval(Arrays&&... outputs) { @@ -29,7 +30,7 @@ void eval(Arrays&&... outputs) { * Jacobian of the function evaluated at the primals. Returns a pair of * vectors of output arrays and VJP arrays. **/ -std::pair, std::vector> vjp( +MLX_API std::pair, std::vector> vjp( const std::function(const std::vector&)>& fun, const std::vector& primals, const std::vector& cotangents); @@ -37,7 +38,7 @@ std::pair, std::vector> vjp( /** * Computes the output and vector-Jacobian product (VJP) of a unary function. */ -std::pair vjp( +MLX_API std::pair vjp( const std::function& fun, const array& primal, const array& cotangent); @@ -49,7 +50,7 @@ std::pair vjp( * evaluated at the primals with the vector of tangents. Returns a pair of * vectors of output arrays and JVP arrays. **/ -std::pair, std::vector> jvp( +MLX_API std::pair, std::vector> jvp( const std::function(const std::vector&)>& fun, const std::vector& primals, const std::vector& tangents); @@ -57,7 +58,7 @@ std::pair, std::vector> jvp( /** * Computes the output and Jacobian-vector product (JVP) of a unary function. */ -std::pair jvp( +MLX_API std::pair jvp( const std::function& fun, const array& primal, const array& tangent); @@ -75,7 +76,7 @@ using SimpleValueAndGradFn = std::function>( * Returns a function which computes the value and gradient of the input * function with respect to a vector of input arrays. **/ -ValueAndGradFn value_and_grad( +MLX_API ValueAndGradFn value_and_grad( const std::function(const std::vector&)>& fun, const std::vector& argnums); @@ -157,7 +158,7 @@ std::function inline grad( /** * Automatically vectorize a unary function over the requested axes. */ -std::function vmap( +MLX_API std::function vmap( const std::function& fun, int in_axis = 0, int out_axis = 0); @@ -165,7 +166,7 @@ std::function vmap( /** * Automatically vectorize a binary function over the requested axes. */ -std::function vmap( +MLX_API std::function vmap( const std::function& fun, int in_axis_a = 0, int in_axis_b = 0, @@ -180,7 +181,7 @@ std::function vmap( * Returns a vectorized function with the same signature as the input * function. */ -std::function(const std::vector&)> vmap( +MLX_API std::function(const std::vector&)> vmap( const std::function(const std::vector&)>& fun, const std::vector& in_axes = {}, const std::vector& out_axes = {}); @@ -194,7 +195,8 @@ std::function(const std::vector&)> vmap( * If any transformation is not provided, then a default one is created by * calling `vjp`, `jvp` and `vmap` on the function directly. */ -std::function(const std::vector&)> custom_function( +MLX_API std::function(const std::vector&)> +custom_function( std::function(const std::vector&)> fun, std::optional( const std::vector&, @@ -212,7 +214,7 @@ std::function(const std::vector&)> custom_function( * Return a function that behaves exactly like `fun` but if the vjp of the * results is computed `fun_vjp` will be used instead of `vjp(fun, ...)` . */ -std::function(const std::vector&)> custom_vjp( +MLX_API std::function(const std::vector&)> custom_vjp( std::function(const std::vector&)> fun, std::function( const std::vector&, @@ -223,7 +225,7 @@ std::function(const std::vector&)> custom_vjp( * Checkpoint the gradient of a function. Namely, discard all intermediate * state and recalculate it when we need to compute the gradient. */ -std::function(const std::vector&)> checkpoint( +MLX_API std::function(const std::vector&)> checkpoint( std::function(const std::vector&)> fun); } // namespace mlx::core diff --git a/mlx/transforms_impl.h b/mlx/transforms_impl.h index 46851fa3..eff458c4 100644 --- a/mlx/transforms_impl.h +++ b/mlx/transforms_impl.h @@ -2,14 +2,16 @@ #pragma once +#include "mlx/api.h" + namespace mlx::core::detail { -std::pair, std::vector> vmap_trace( +MLX_API std::pair, std::vector> vmap_trace( const std::function(const std::vector&)>& fun, const std::vector& inputs, const std::vector& in_axes); -std::vector vmap_replace( +MLX_API std::vector vmap_replace( const std::vector& inputs, const std::vector& s_inputs, const std::vector& s_outputs, diff --git a/mlx/utils.h b/mlx/utils.h index dbf79a71..bb2de466 100644 --- a/mlx/utils.h +++ b/mlx/utils.h @@ -5,6 +5,7 @@ #include #include +#include "mlx/api.h" #include "mlx/array.h" #include "mlx/device.h" #include "mlx/dtype.h" @@ -13,8 +14,8 @@ namespace mlx::core { using StreamOrDevice = std::variant; -Stream to_stream(StreamOrDevice s); -Stream to_stream(StreamOrDevice s, Device default_); +MLX_API Stream to_stream(StreamOrDevice s); +MLX_API Stream to_stream(StreamOrDevice s, Device default_); struct StreamContext { public: @@ -54,13 +55,13 @@ struct PrintFormatter { bool capitalize_bool{false}; }; -PrintFormatter& get_global_formatter(); +MLX_API PrintFormatter& get_global_formatter(); /** Print the exception and then abort. */ -void abort_with_exception(const std::exception& error); +MLX_API void abort_with_exception(const std::exception& error); /** Holds information about floating-point types. */ -struct finfo { +struct MLX_API finfo { explicit finfo(Dtype dtype); Dtype dtype; double min; @@ -69,7 +70,7 @@ struct finfo { }; /** Holds information about integral types. */ -struct iinfo { +struct MLX_API iinfo { explicit iinfo(Dtype dtype); Dtype dtype; int64_t min; @@ -83,23 +84,21 @@ inline Dtype result_type(const array& a, const array& b) { inline Dtype result_type(const array& a, const array& b, const array& c) { return promote_types(result_type(a, b), c.dtype()); } -Dtype result_type(const std::vector& arrays); +MLX_API Dtype result_type(const std::vector& arrays); -Shape broadcast_shapes(const Shape& s1, const Shape& s2); +MLX_API Shape broadcast_shapes(const Shape& s1, const Shape& s2); /** * Returns the axis normalized to be in the range [0, ndim). */ -int normalize_axis_index( - int axis, - int ndim, - const std::string& msg_prefix = ""); +MLX_API int +normalize_axis_index(int axis, int ndim, const std::string& msg_prefix = ""); -std::ostream& operator<<(std::ostream& os, const Device& d); -std::ostream& operator<<(std::ostream& os, const Stream& s); -std::ostream& operator<<(std::ostream& os, const Dtype& d); -std::ostream& operator<<(std::ostream& os, const Dtype::Kind& k); -std::ostream& operator<<(std::ostream& os, array a); +MLX_API std::ostream& operator<<(std::ostream& os, const Device& d); +MLX_API std::ostream& operator<<(std::ostream& os, const Stream& s); +MLX_API std::ostream& operator<<(std::ostream& os, const Dtype& d); +MLX_API std::ostream& operator<<(std::ostream& os, const Dtype::Kind& k); +MLX_API std::ostream& operator<<(std::ostream& os, array a); inline std::ostream& operator<<(std::ostream& os, const complex64_t& v) { return os << v.real() << (v.imag() >= 0 ? "+" : "") << v.imag() << "j"; } diff --git a/mlx/version.cpp b/mlx/version.cpp index b60168e0..79f5948d 100644 --- a/mlx/version.cpp +++ b/mlx/version.cpp @@ -1,5 +1,7 @@ // Copyright © 2025 Apple Inc. +#include "mlx/version.h" + namespace mlx::core { const char* version() { diff --git a/mlx/version.h b/mlx/version.h index d83a24e8..363bf0c0 100644 --- a/mlx/version.h +++ b/mlx/version.h @@ -2,6 +2,8 @@ #pragma once +#include "mlx/api.h" + #define MLX_VERSION_MAJOR 0 #define MLX_VERSION_MINOR 30 #define MLX_VERSION_PATCH 4 @@ -15,6 +17,6 @@ namespace mlx::core { * * For dev builds, the version will include the suffix ".devYYYYMMDD+hash" */ -const char* version(); +MLX_API const char* version(); } // namespace mlx::core diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 7cac06c9..ed15546e 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -39,5 +39,16 @@ target_sources( target_link_libraries(tests PRIVATE mlx doctest) target_compile_options(tests PRIVATE ${SANITIZER_COMPILE_FLAGS}) target_link_options(tests PRIVATE ${SANITIZER_LINK_FLAGS}) + +# On Windows, copy the mlx DLL to the test executable directory for runtime +# loading +if(WIN32 AND BUILD_SHARED_LIBS) + add_custom_command( + TARGET tests + POST_BUILD + COMMAND ${CMAKE_COMMAND} -E copy_if_different $ + $) +endif() + doctest_discover_tests(tests) add_test(NAME tests COMMAND tests) diff --git a/tests/tests.cpp b/tests/tests.cpp index edaf49c2..7c7ea3a5 100644 --- a/tests/tests.cpp +++ b/tests/tests.cpp @@ -15,7 +15,9 @@ int main(int argc, char** argv) { const char* device = std::getenv("DEVICE"); if (device != nullptr && std::string(device) == "cpu") { set_default_device(Device::cpu); - } else if (metal::is_available()) { + } else if (is_available(Device::gpu)) { + // Use generic GPU availability check (works for Metal on macOS, or CUDA on + // Linux/Windows) set_default_device(Device::gpu); }