mlx-cpu 0.30.1__py3-none-manylinux_2_35_x86_64.whl

mlx/share/cmake/MLX/MLXConfigVersion.cmake

@@ -0,0 +1,65 @@
+# This is a basic version file for the Config-mode of find_package().
+# It is used by write_basic_package_version_file() as input file for configure_file()
+# to create a version-file which can be installed along a config.cmake file.
+#
+# The created file sets PACKAGE_VERSION_EXACT if the current version string and
+# the requested version string are exactly the same and it sets
+# PACKAGE_VERSION_COMPATIBLE if the current version is >= requested version,
+# but only if the requested major version is the same as the current one.
+# The variable CVF_VERSION must be set before calling configure_file().
+
+
+set(PACKAGE_VERSION "0.30.1")
+
+if(PACKAGE_VERSION VERSION_LESS PACKAGE_FIND_VERSION)
+  set(PACKAGE_VERSION_COMPATIBLE FALSE)
+else()
+
+  if("0.30.1" MATCHES "^([0-9]+)\\.")
+    set(CVF_VERSION_MAJOR "${CMAKE_MATCH_1}")
+    if(NOT CVF_VERSION_MAJOR VERSION_EQUAL 0)
+      string(REGEX REPLACE "^0+" "" CVF_VERSION_MAJOR "${CVF_VERSION_MAJOR}")
+    endif()
+  else()
+    set(CVF_VERSION_MAJOR "0.30.1")
+  endif()
+
+  if(PACKAGE_FIND_VERSION_RANGE)
+    # both endpoints of the range must have the expected major version
+    math (EXPR CVF_VERSION_MAJOR_NEXT "${CVF_VERSION_MAJOR} + 1")
+    if (NOT PACKAGE_FIND_VERSION_MIN_MAJOR STREQUAL CVF_VERSION_MAJOR
+        OR ((PACKAGE_FIND_VERSION_RANGE_MAX STREQUAL "INCLUDE" AND NOT PACKAGE_FIND_VERSION_MAX_MAJOR STREQUAL CVF_VERSION_MAJOR)
+          OR (PACKAGE_FIND_VERSION_RANGE_MAX STREQUAL "EXCLUDE" AND NOT PACKAGE_FIND_VERSION_MAX VERSION_LESS_EQUAL CVF_VERSION_MAJOR_NEXT)))
+      set(PACKAGE_VERSION_COMPATIBLE FALSE)
+    elseif(PACKAGE_FIND_VERSION_MIN_MAJOR STREQUAL CVF_VERSION_MAJOR
+        AND ((PACKAGE_FIND_VERSION_RANGE_MAX STREQUAL "INCLUDE" AND PACKAGE_VERSION VERSION_LESS_EQUAL PACKAGE_FIND_VERSION_MAX)
+        OR (PACKAGE_FIND_VERSION_RANGE_MAX STREQUAL "EXCLUDE" AND PACKAGE_VERSION VERSION_LESS PACKAGE_FIND_VERSION_MAX)))
+      set(PACKAGE_VERSION_COMPATIBLE TRUE)
+    else()
+      set(PACKAGE_VERSION_COMPATIBLE FALSE)
+    endif()
+  else()
+    if(PACKAGE_FIND_VERSION_MAJOR STREQUAL CVF_VERSION_MAJOR)
+      set(PACKAGE_VERSION_COMPATIBLE TRUE)
+    else()
+      set(PACKAGE_VERSION_COMPATIBLE FALSE)
+    endif()
+
+    if(PACKAGE_FIND_VERSION STREQUAL PACKAGE_VERSION)
+      set(PACKAGE_VERSION_EXACT TRUE)
+    endif()
+  endif()
+endif()
+
+
+# if the installed or the using project don't have CMAKE_SIZEOF_VOID_P set, ignore it:
+if("${CMAKE_SIZEOF_VOID_P}" STREQUAL "" OR "8" STREQUAL "")
+  return()
+endif()
+
+# check that the installed version has the same 32/64bit-ness as the one which is currently searching:
+if(NOT CMAKE_SIZEOF_VOID_P STREQUAL "8")
+  math(EXPR installedBits "8 * 8")
+  set(PACKAGE_VERSION "${PACKAGE_VERSION} (${installedBits}bit)")
+  set(PACKAGE_VERSION_UNSUITABLE TRUE)
+endif()

mlx/share/cmake/MLX/MLXTargets-release.cmake

@@ -0,0 +1,19 @@
+#----------------------------------------------------------------
+# Generated CMake target import file for configuration "Release".
+#----------------------------------------------------------------
+
+# Commands may need to know the format version.
+set(CMAKE_IMPORT_FILE_VERSION 1)
+
+# Import target "mlx" for configuration "Release"
+set_property(TARGET mlx APPEND PROPERTY IMPORTED_CONFIGURATIONS RELEASE)
+set_target_properties(mlx PROPERTIES
+  IMPORTED_LOCATION_RELEASE "${_IMPORT_PREFIX}/lib/libmlx.so"
+  IMPORTED_SONAME_RELEASE "libmlx.so"
+  )
+
+list(APPEND _cmake_import_check_targets mlx )
+list(APPEND _cmake_import_check_files_for_mlx "${_IMPORT_PREFIX}/lib/libmlx.so" )
+
+# Commands beyond this point should not need to know the version.
+set(CMAKE_IMPORT_FILE_VERSION)

mlx/share/cmake/MLX/MLXTargets.cmake

@@ -0,0 +1,106 @@
+# Generated by CMake
+
+if("${CMAKE_MAJOR_VERSION}.${CMAKE_MINOR_VERSION}" LESS 2.8)
+   message(FATAL_ERROR "CMake >= 2.8.3 required")
+endif()
+if(CMAKE_VERSION VERSION_LESS "2.8.3")
+   message(FATAL_ERROR "CMake >= 2.8.3 required")
+endif()
+cmake_policy(PUSH)
+cmake_policy(VERSION 2.8.3...4.0)
+#----------------------------------------------------------------
+# Generated CMake target import file.
+#----------------------------------------------------------------
+
+# Commands may need to know the format version.
+set(CMAKE_IMPORT_FILE_VERSION 1)
+
+# Protect against multiple inclusion, which would fail when already imported targets are added once more.
+set(_cmake_targets_defined "")
+set(_cmake_targets_not_defined "")
+set(_cmake_expected_targets "")
+foreach(_cmake_expected_target IN ITEMS mlx)
+  list(APPEND _cmake_expected_targets "${_cmake_expected_target}")
+  if(TARGET "${_cmake_expected_target}")
+    list(APPEND _cmake_targets_defined "${_cmake_expected_target}")
+  else()
+    list(APPEND _cmake_targets_not_defined "${_cmake_expected_target}")
+  endif()
+endforeach()
+unset(_cmake_expected_target)
+if(_cmake_targets_defined STREQUAL _cmake_expected_targets)
+  unset(_cmake_targets_defined)
+  unset(_cmake_targets_not_defined)
+  unset(_cmake_expected_targets)
+  unset(CMAKE_IMPORT_FILE_VERSION)
+  cmake_policy(POP)
+  return()
+endif()
+if(NOT _cmake_targets_defined STREQUAL "")
+  string(REPLACE ";" ", " _cmake_targets_defined_text "${_cmake_targets_defined}")
+  string(REPLACE ";" ", " _cmake_targets_not_defined_text "${_cmake_targets_not_defined}")
+  message(FATAL_ERROR "Some (but not all) targets in this export set were already defined.\nTargets Defined: ${_cmake_targets_defined_text}\nTargets not yet defined: ${_cmake_targets_not_defined_text}\n")
+endif()
+unset(_cmake_targets_defined)
+unset(_cmake_targets_not_defined)
+unset(_cmake_expected_targets)
+
+
+# Compute the installation prefix relative to this file.
+get_filename_component(_IMPORT_PREFIX "${CMAKE_CURRENT_LIST_FILE}" PATH)
+get_filename_component(_IMPORT_PREFIX "${_IMPORT_PREFIX}" PATH)
+get_filename_component(_IMPORT_PREFIX "${_IMPORT_PREFIX}" PATH)
+get_filename_component(_IMPORT_PREFIX "${_IMPORT_PREFIX}" PATH)
+if(_IMPORT_PREFIX STREQUAL "/")
+  set(_IMPORT_PREFIX "")
+endif()
+
+# Create imported target mlx
+add_library(mlx SHARED IMPORTED)
+
+set_target_properties(mlx PROPERTIES
+  INTERFACE_INCLUDE_DIRECTORIES "${_IMPORT_PREFIX}/include;${_IMPORT_PREFIX}/include"
+)
+
+# Load information for each installed configuration.
+file(GLOB _cmake_config_files "${CMAKE_CURRENT_LIST_DIR}/MLXTargets-*.cmake")
+foreach(_cmake_config_file IN LISTS _cmake_config_files)
+  include("${_cmake_config_file}")
+endforeach()
+unset(_cmake_config_file)
+unset(_cmake_config_files)
+
+# Cleanup temporary variables.
+set(_IMPORT_PREFIX)
+
+# Loop over all imported files and verify that they actually exist
+foreach(_cmake_target IN LISTS _cmake_import_check_targets)
+  if(CMAKE_VERSION VERSION_LESS "3.28"
+      OR NOT DEFINED _cmake_import_check_xcframework_for_${_cmake_target}
+      OR NOT IS_DIRECTORY "${_cmake_import_check_xcframework_for_${_cmake_target}}")
+    foreach(_cmake_file IN LISTS "_cmake_import_check_files_for_${_cmake_target}")
+      if(NOT EXISTS "${_cmake_file}")
+        message(FATAL_ERROR "The imported target \"${_cmake_target}\" references the file
+   \"${_cmake_file}\"
+but this file does not exist.  Possible reasons include:
+* The file was deleted, renamed, or moved to another location.
+* An install or uninstall procedure did not complete successfully.
+* The installation package was faulty and contained
+   \"${CMAKE_CURRENT_LIST_FILE}\"
+but not all the files it references.
+")
+      endif()
+    endforeach()
+  endif()
+  unset(_cmake_file)
+  unset("_cmake_import_check_files_for_${_cmake_target}")
+endforeach()
+unset(_cmake_target)
+unset(_cmake_import_check_targets)
+
+# This file does not depend on other imported targets which have
+# been exported from the same project but in a separate export set.
+
+# Commands beyond this point should not need to know the version.
+set(CMAKE_IMPORT_FILE_VERSION)
+cmake_policy(POP)

mlx/share/cmake/MLX/extension.cmake

@@ -0,0 +1,50 @@
+include(CMakeParseArguments)
+
+# clang format off
+#
+# ##############################################################################
+# Build metal library
+#
+# Adds a custom target ${TARGET} to build ${OUTPUT_DIRECTORY}/{TITLE}.metallib
+# from list ${SOURCES}, including list ${INCLUDE_DIRS}, depends on list ${DEPS}
+#
+# Args: TARGET: Custom target to be added for the metal library TITLE: Name of
+# the .metallib OUTPUT_DIRECTORY: Where to place ${TITLE}.metallib SOURCES: List
+# of source files INCLUDE_DIRS: List of include dirs DEPS: List of dependency
+# files (like headers) DEBUG: Boolean, if true, enables debug compile options
+# for this specific library. If not provided, uses global MLX_METAL_DEBUG.
+#
+# clang format on
+
+macro(mlx_build_metallib)
+  # Parse args
+  set(oneValueArgs TARGET TITLE OUTPUT_DIRECTORY DEBUG)
+  set(multiValueArgs SOURCES INCLUDE_DIRS DEPS)
+  cmake_parse_arguments(MTLLIB "" "${oneValueArgs}" "${multiValueArgs}" ${ARGN})
+
+  # Set output
+  set(MTLLIB_BUILD_TARGET "${MTLLIB_OUTPUT_DIRECTORY}/${MTLLIB_TITLE}.metallib")
+
+  # Collect compile options
+  set(MTLLIB_COMPILE_OPTIONS -Wall -Wextra -fno-fast-math -Wno-c++17-extensions)
+  if(MLX_METAL_DEBUG OR MTLLIB_DEBUG)
+    set(MTLLIB_COMPILE_OPTIONS ${MTLLIB_COMPILE_OPTIONS} -gline-tables-only
+                               -frecord-sources)
+  endif()
+
+  # Prepare metallib build command
+  add_custom_command(
+    OUTPUT ${MTLLIB_BUILD_TARGET}
+    COMMAND
+      xcrun -sdk macosx metal
+      "$<LIST:TRANSFORM,${MTLLIB_INCLUDE_DIRS},PREPEND,-I>"
+      ${MTLLIB_COMPILE_OPTIONS} ${MTLLIB_SOURCES} -o ${MTLLIB_BUILD_TARGET}
+    DEPENDS ${MTLLIB_DEPS} ${MTLLIB_SOURCES}
+    COMMAND_EXPAND_LISTS
+    COMMENT "Building ${MTLLIB_TITLE}.metallib"
+    VERBATIM)
+
+  # Add metallib custom target
+  add_custom_target(${MTLLIB_TARGET} DEPENDS ${MTLLIB_BUILD_TARGET})
+
+endmacro(mlx_build_metallib)

mlx/utils.py

@@ -0,0 +1,325 @@
+# Copyright © 2023 Apple Inc.
+from collections import defaultdict
+from itertools import zip_longest
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+
+def tree_map(
+    fn: Callable, tree: Any, *rest: Any, is_leaf: Optional[Callable] = None
+) -> Any:
+    """Applies ``fn`` to the leaves of the Python tree ``tree`` and
+    returns a new collection with the results.
+
+    If ``rest`` is provided, every item is assumed to be a superset of ``tree``
+    and the corresponding leaves are provided as extra positional arguments to
+    ``fn``. In that respect, :meth:`tree_map` is closer to :func:`itertools.starmap`
+    than to :func:`map`.
+
+    The keyword argument ``is_leaf`` decides what constitutes a leaf from
+    ``tree`` similar to :func:`tree_flatten`.
+
+    .. code-block:: python
+
+        import mlx.nn as nn
+        from mlx.utils import tree_map
+
+        model = nn.Linear(10, 10)
+        print(model.parameters().keys())
+        # dict_keys(['weight', 'bias'])
+
+        # square the parameters
+        model.update(tree_map(lambda x: x*x, model.parameters()))
+
+    Args:
+        fn (callable): The function that processes the leaves of the tree.
+        tree (Any): The main Python tree that will be iterated upon.
+        rest (tuple[Any]): Extra trees to be iterated together with ``tree``.
+        is_leaf (callable, optional): An optional callable that returns ``True``
+           if the passed object is considered a leaf or ``False`` otherwise.
+
+    Returns:
+        A Python tree with the new values returned by ``fn``.
+    """
+    if is_leaf is not None and is_leaf(tree):
+        return fn(tree, *rest)
+    elif isinstance(tree, (list, tuple)):
+        TreeType = type(tree)
+        subtrees = (
+            tree_map(fn, child, *(r[i] for r in rest), is_leaf=is_leaf)
+            for i, child in enumerate(tree)
+        )
+        return TreeType(*subtrees) if hasattr(tree, "_fields") else TreeType(subtrees)
+    elif isinstance(tree, dict):
+        return {
+            k: tree_map(fn, child, *(r[k] for r in rest), is_leaf=is_leaf)
+            for k, child in tree.items()
+        }
+    else:
+        return fn(tree, *rest)
+
+
+def tree_map_with_path(
+    fn: Callable,
+    tree: Any,
+    *rest: Any,
+    is_leaf: Optional[Callable] = None,
+    path: Optional[Any] = None,
+) -> Any:
+    """Applies ``fn`` to the path and leaves of the Python tree ``tree`` and
+    returns a new collection with the results.
+
+    This function is the same :func:`tree_map` but the ``fn`` takes the path as
+    the first argument followed by the remaining tree nodes.
+
+    Args:
+        fn (callable): The function that processes the leaves of the tree.
+        tree (Any): The main Python tree that will be iterated upon.
+        rest (tuple[Any]): Extra trees to be iterated together with ``tree``.
+        is_leaf (Optional[Callable]): An optional callable that returns ``True``
+           if the passed object is considered a leaf or ``False`` otherwise.
+        path (Optional[Any]): Prefix will be added to the result.
+
+    Returns:
+        A Python tree with the new values returned by ``fn``.
+
+    Example:
+        >>> from mlx.utils import tree_map_with_path
+        >>> tree = {"model": [{"w": 0, "b": 1}, {"w": 0, "b": 1}]}
+        >>> new_tree = tree_map_with_path(lambda path, _: print(path), tree)
+        model.0.w
+        model.0.b
+        model.1.w
+        model.1.b
+    """
+    if is_leaf is not None and is_leaf(tree):
+        return fn(path, tree, *rest)
+    elif isinstance(tree, (list, tuple)):
+        prefix = f"{path}." if path else ""
+        TreeType = type(tree)
+        return TreeType(
+            tree_map_with_path(
+                fn, child, *(r[i] for r in rest), is_leaf=is_leaf, path=f"{prefix}{i}"
+            )
+            for i, child in enumerate(tree)
+        )
+    elif isinstance(tree, dict):
+        prefix = f"{path}." if path else ""
+        return {
+            k: tree_map_with_path(
+                fn, child, *(r[k] for r in rest), is_leaf=is_leaf, path=f"{prefix}{k}"
+            )
+            for k, child in tree.items()
+        }
+    else:
+        return fn(path, tree, *rest)
+
+
+def tree_flatten(
+    tree: Any,
+    prefix: str = "",
+    is_leaf: Optional[Callable] = None,
+    destination: Optional[Union[List[Tuple[str, Any]], Dict[str, Any]]] = None,
+) -> Union[List[Tuple[str, Any]], Dict[str, Any]]:
+    """Flattens a Python tree to a list of key, value tuples.
+
+    The keys are using the dot notation to define trees of arbitrary depth and
+    complexity.
+
+    .. code-block:: python
+
+        from mlx.utils import tree_flatten
+
+        print(tree_flatten([[[0]]]))
+        # [("0.0.0", 0)]
+
+        print(tree_flatten([[[0]]], prefix=".hello"))
+        # [("hello.0.0.0", 0)]
+
+        tree_flatten({"a": {"b": 1}}, destination={})
+        {"a.b": 1}
+
+    .. note::
+       Dictionaries should have keys that are valid Python identifiers.
+
+    Args:
+        tree (Any): The Python tree to be flattened.
+        prefix (str): A prefix to use for the keys. The first character is
+            always discarded.
+        is_leaf (callable): An optional callable that returns True if the
+            passed object is considered a leaf or False otherwise.
+        destination (list or dict, optional): A list or dictionary to store the
+            flattened tree. If None an empty list will be used. Default: ``None``.
+
+    Returns:
+        Union[List[Tuple[str, Any]], Dict[str, Any]]: The flat representation of
+            the Python tree.
+    """
+    if destination is None:
+        destination = []
+
+    # Create the function to update the destination. We are taking advantage of
+    # the fact that list.extend and dict.update have the same API to simplify
+    # the code a bit.
+    if isinstance(destination, list):
+        _add_to_destination = destination.extend
+    elif isinstance(destination, dict):
+        _add_to_destination = destination.update
+    else:
+        raise ValueError("Destination should be either a list or a dictionary or None")
+
+    # Leaf identified by is_leaf so add it and return
+    if is_leaf is not None and is_leaf(tree):
+        _add_to_destination([(prefix[1:], tree)])
+        return destination
+
+    # List or tuple so recursively add each subtree
+    if isinstance(tree, (list, tuple)):
+        for i, item in enumerate(tree):
+            tree_flatten(item, f"{prefix}.{i}", is_leaf, destination)
+        return destination
+
+    # Dictionary so recursively add each subtree
+    if isinstance(tree, dict):
+        for key, value in tree.items():
+            tree_flatten(value, f"{prefix}.{key}", is_leaf, destination)
+        return destination
+
+    # Leaf so add it and return
+    _add_to_destination([(prefix[1:], tree)])
+
+    return destination
+
+
+def tree_unflatten(tree: Union[List[Tuple[str, Any]], Dict[str, Any]]) -> Any:
+    """Recreate a Python tree from its flat representation.
+
+    .. code-block:: python
+
+        from mlx.utils import tree_unflatten
+
+        d = tree_unflatten([("hello.world", 42)])
+        print(d)
+        # {"hello": {"world": 42}}
+
+        d = tree_unflatten({"hello.world": 42})
+        print(d)
+        # {"hello": {"world": 42}}
+
+    Args:
+        tree (list[tuple[str, Any]] or dict[str, Any]): The flat representation of a Python tree.
+           For instance as returned by :meth:`tree_flatten`.
+
+    Returns:
+        A Python tree.
+    """
+    items = tree.items() if isinstance(tree, dict) else tree
+
+    # Special case when we have just one element in the tree ie not a tree
+    if len(items) == 1:
+        key, value = next(iter(items))
+        if key == "":
+            return value
+
+    # collect children
+    children = defaultdict(list)
+    for key, value in items:
+        current_idx, *next_idx = key.split(".", maxsplit=1)
+        next_idx = "" if not next_idx else next_idx[0]
+        children[current_idx].append((next_idx, value))
+
+    # Assume they are a list and fail to dict if the keys are not all integers
+    try:
+        keys = sorted((int(idx), idx) for idx in children.keys())
+        l = []
+        for i, k in keys:
+            # if i <= len(l), no {} will be appended.
+            l.extend([{} for _ in range(i - len(l))])
+            l.append(tree_unflatten(children[k]))
+        return l
+    except ValueError:
+        return {k: tree_unflatten(v) for k, v in children.items()}
+
+
+def tree_reduce(fn, tree, initializer=None, is_leaf=None):
+    """Applies a reduction to the leaves of a Python tree.
+
+    This function reduces Python trees into an accumulated result by applying
+    the provided function ``fn`` to the leaves of the tree.
+
+    Example:
+        >>> from mlx.utils import tree_reduce
+        >>> tree = {"a": [1, 2, 3], "b": [4, 5]}
+        >>> tree_reduce(lambda acc, x: acc + x, tree, 0)
+        15
+
+    Args:
+        fn (callable): The reducer function that takes two arguments (accumulator,
+            current value) and returns the updated accumulator.
+        tree (Any): The Python tree to reduce. It can be any nested combination of
+            lists, tuples, or dictionaries.
+        initializer (Any, optional): The initial value to start the reduction. If
+            not provided, the first leaf value is used.
+        is_leaf (callable, optional): A function to determine if an object is a
+            leaf, returning ``True`` for leaf nodes and ``False`` otherwise.
+
+    Returns:
+        Any: The accumulated value.
+    """
+    if is_leaf is not None and is_leaf(tree):
+        return tree if initializer is None else fn(initializer, tree)
+
+    accumulator = initializer
+
+    if isinstance(tree, (list, tuple)):
+        for item in tree:
+            accumulator = tree_reduce(fn, item, accumulator, is_leaf)
+    elif isinstance(tree, dict):
+        for item in tree.values():
+            accumulator = tree_reduce(fn, item, accumulator, is_leaf)
+    else:
+        return tree if accumulator is None else fn(accumulator, tree)
+
+    return accumulator
+
+
+def tree_merge(tree_a, tree_b, merge_fn=None):
+    """Merge two Python trees in one containing the values of both. It can be
+    thought of as a deep dict.update method.
+
+    Args:
+        tree_a (Any): The first Python tree.
+        tree_b (Any): The second Python tree.
+        merge_fn (callable, optional): A function to merge leaves.
+
+    Returns:
+        The Python tree containing the values of both ``tree_a`` and
+        ``tree_b``.
+    """
+    if isinstance(tree_a, (dict, list, tuple)) and len(tree_a) == 0:
+        tree_a = None
+    if isinstance(tree_b, (dict, list, tuple)) and len(tree_b) == 0:
+        tree_b = None
+    if tree_a is None and tree_b is not None:
+        return tree_b
+    if tree_a is not None and tree_b is None:
+        return tree_a
+
+    if isinstance(tree_a, (list, tuple)) and isinstance(tree_b, (list, tuple)):
+        TreeType = type(tree_a)
+        return TreeType(
+            tree_merge(a, b, merge_fn) for a, b in zip_longest(tree_a, tree_b)
+        )
+    elif isinstance(tree_a, dict) and isinstance(tree_b, dict):
+        return {
+            k: tree_merge(tree_a.get(k, None), tree_b.get(k, None), merge_fn)
+            for k in set(tree_a.keys()) | set(tree_b.keys())
+        }
+    else:
+        if merge_fn is None:
+            raise ValueError(
+                (
+                    "Trees contain elements at the same locations but no merge "
+                    "function was provided"
+                )
+            )
+        return merge_fn(tree_a, tree_b)