Merge branch 'ml-explore:main' into adding-Muon-optimizer

python/mlx/extension.py

@@ -6,7 +6,7 @@ import subprocess
 import sys
 from pathlib import Path
 
-from setuptools import Extension, find_namespace_packages, setup
+from setuptools import Extension
 from setuptools.command.build_ext import build_ext
 
 import mlx

python/mlx/utils.py

@@ -62,7 +62,7 @@ def tree_map_with_path(
     tree: Any,
     *rest: Any,
     is_leaf: Optional[Callable] = None,
-    path: Any = 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.
@@ -74,8 +74,9 @@ def tree_map_with_path(
         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``
+        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``.

python/src/random.cpp

@@ -253,7 +253,7 @@ void init_random(nb::module_& parent_module) {
 
         The values are sampled with equal probability from the integers in
         half-open interval ``[low, high)``. The lower and upper bound can be
-        scalars or arrays and must be roadcastable to ``shape``.
+        scalars or arrays and must be broadcastable to ``shape``.
 
         Args:
             low (scalar or array): Lower bound of the interval.