ce304fafe19161978ad512b385c65426bad519e5a0b8fb3f0659eace3d2ea3cc

lib/python3.11/site-packages/mlx-0.0.7.dist-info/INSTALLER

@@ -0,0 +1 @@
+pip

lib/python3.11/site-packages/mlx-0.0.7.dist-info/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright © 2023 Apple Inc.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

lib/python3.11/site-packages/mlx-0.0.7.dist-info/METADATA

@@ -0,0 +1,122 @@
+Metadata-Version: 2.1
+Name: mlx
+Version: 0.0.7
+Summary: A framework for machine learning on Apple silicon.
+Author: MLX Contributors
+Author-email: mlx@group.apple.com
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Provides-Extra: dev
+Requires-Dist: pre-commit ; extra == 'dev'
+Requires-Dist: pybind11-stubgen ; extra == 'dev'
+Provides-Extra: testing
+Requires-Dist: numpy ; extra == 'testing'
+Requires-Dist: torch ; extra == 'testing'
+
+# MLX
+
+[**Quickstart**](#quickstart) | [**Installation**](#installation) |
+[**Documentation**](https://ml-explore.github.io/mlx/build/html/index.html) |
+[**Examples**](#examples) 
+
+[![CircleCI](https://circleci.com/gh/ml-explore/mlx.svg?style=svg)](https://circleci.com/gh/ml-explore/mlx)
+
+MLX is an array framework for machine learning on Apple silicon, brought to you
+by Apple machine learning research.
+
+Some key features of MLX include:
+
+ - **Familiar APIs**: MLX has a Python API that closely follows NumPy.
+   MLX also has a fully featured C++ API, which closely mirrors the Python API. 
+   MLX has higher-level packages like `mlx.nn` and `mlx.optimizers` with APIs
+   that closely follow PyTorch to simplify building more complex models.
+
+ - **Composable function transformations**: MLX supports composable function
+   transformations for automatic differentiation, automatic vectorization,
+   and computation graph optimization.
+
+ - **Lazy computation**: Computations in MLX are lazy. Arrays are only
+   materialized when needed.
+
+ - **Dynamic graph construction**: Computation graphs in MLX are constructed
+   dynamically. Changing the shapes of function arguments does not trigger
+   slow compilations, and debugging is simple and intuitive.
+
+ - **Multi-device**: Operations can run on any of the supported devices
+   (currently the CPU and the GPU).
+
+ - **Unified memory**: A notable difference from MLX and other frameworks
+   is the *unified memory model*. Arrays in MLX live in shared memory.
+   Operations on MLX arrays can be performed on any of the supported
+   device types without transferring data.
+
+MLX is designed by machine learning researchers for machine learning
+researchers. The framework is intended to be user-friendly, but still efficient
+to train and deploy models. The design of the framework itself is also
+conceptually simple. We intend to make it easy for researchers to extend and
+improve MLX with the goal of quickly exploring new ideas. 
+
+The design of MLX is inspired by frameworks like
+[NumPy](https://numpy.org/doc/stable/index.html),
+[PyTorch](https://pytorch.org/), [Jax](https://github.com/google/jax), and
+[ArrayFire](https://arrayfire.org/).
+
+## Examples
+
+The [MLX examples repo](https://github.com/ml-explore/mlx-examples) has a
+variety of examples, including:
+
+- [Transformer language model](https://github.com/ml-explore/mlx-examples/tree/main/transformer_lm) training.
+- Large-scale text generation with
+  [LLaMA](https://github.com/ml-explore/mlx-examples/tree/main/llms/llama) and
+  finetuning with [LoRA](https://github.com/ml-explore/mlx-examples/tree/main/lora).
+- Generating images with [Stable Diffusion](https://github.com/ml-explore/mlx-examples/tree/main/stable_diffusion).
+- Speech recognition with [OpenAI's Whisper](https://github.com/ml-explore/mlx-examples/tree/main/whisper).
+
+## Quickstart
+
+See the [quick start
+guide](https://ml-explore.github.io/mlx/build/html/quick_start.html)
+in the documentation.
+
+## Installation
+
+MLX is available on [PyPI](https://pypi.org/project/mlx/). To install the Python API, run:
+
+```
+pip install mlx
+```
+
+Checkout the
+[documentation](https://ml-explore.github.io/mlx/build/html/install.html#)
+for more information on building the C++ and Python APIs from source.
+
+## Contributing 
+
+Check out the [contribution guidelines](CONTRIBUTING.md) for more information
+on contributing to MLX. See the
+[docs](https://ml-explore.github.io/mlx/build/html/install.html) for more
+information on building from source, and running tests.
+
+We are grateful for all of [our
+contributors](ACKNOWLEDGMENTS.md#Individual-Contributors). If you contribute
+to MLX and wish to be acknowledged, please add your name to the list in your
+pull request.
+
+## Citing MLX
+
+The MLX software suite was initially developed with equal contribution by Awni
+Hannun, Jagrit Digani, Angelos Katharopoulos, and Ronan Collobert. If you find
+MLX useful in your research and wish to cite it, please use the following
+BibTex entry:
+
+```
+@software{mlx2023,
+  author = {Awni Hannun and Jagrit Digani and Angelos Katharopoulos and Ronan Collobert},
+  title = {{MLX}: Efficient and flexible machine learning on Apple silicon},
+  url = {https://github.com/ml-explore},
+  version = {0.0},
+  year = {2023},
+}
+```

lib/python3.11/site-packages/mlx-0.0.7.dist-info/RECORD

@@ -0,0 +1,199 @@
+mlx-0.0.7.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4
+mlx-0.0.7.dist-info/LICENSE,sha256=zPq3zLLqMG9xUxyMp3u1VQdgbNkHaLHjK4tSq1tIzwE,1066
+mlx-0.0.7.dist-info/METADATA,sha256=CK9nqz-nrSgrFZpRYpc86IYysL_W7zObqiiczNHD5OA,4731
+mlx-0.0.7.dist-info/RECORD,,
+mlx-0.0.7.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+mlx-0.0.7.dist-info/WHEEL,sha256=9CmlYCCxe9UEGaJDl4r_tPKbyfF_OdzBb5_jMO3h1CA,110
+mlx-0.0.7.dist-info/top_level.txt,sha256=NCVZzyms-obPlDq2Sa--m3oyfnnmLma38KCnE6gmcFE,4
+mlx/__pycache__/_reprlib_fix.cpython-311.pyc,,
+mlx/__pycache__/extension.cpython-311.pyc,,
+mlx/__pycache__/optimizers.cpython-311.pyc,,
+mlx/__pycache__/utils.cpython-311.pyc,,
+mlx/_reprlib_fix.py,sha256=YoRQib4PSzl1FIMvs4BwkvipAwOrfufqtV89vNF4Yhc,518
+mlx/core.cpython-311-darwin.so,sha256=4GMiWnLS2sZGsrIwB9mAryue5Lztf-SttAn_AzadmCc,1098840
+mlx/extension.py,sha256=tqKrGeP1ElK9UfBMHUeOq56qWnMMxZrt7DVVM8QZ7yU,3769
+mlx/include/metal_cpp/Foundation/Foundation.hpp,sha256=pYbEkzTrb2O0jJPzI0m9cBNvVSDTh_t39Dv7h0ZDma0,1808
+mlx/include/metal_cpp/Foundation/NSArray.hpp,sha256=8nhS0gF060n7taK9aIynkuMPxlbKgqgijsj61yvhrPQ,4818
+mlx/include/metal_cpp/Foundation/NSAutoreleasePool.hpp,sha256=K8_qNgYgyeP49WUSsz90Lo41qTb-x5Xjs23T9UNSkkU,3302
+mlx/include/metal_cpp/Foundation/NSBundle.hpp,sha256=y0Le2v43bCAdIdrSClvLpRtCL8c4qdB3WK59Gz_AYYs,16845
+mlx/include/metal_cpp/Foundation/NSData.hpp,sha256=rAPEZQn1W05O_sZ6LYexLhvv1CmllqyFWddGIrxGJ8s,2198
+mlx/include/metal_cpp/Foundation/NSDate.hpp,sha256=353UgslEWyV-D1_T7_Sr1fCP4nVjixwJK2QGnRZQKgg,2085
+mlx/include/metal_cpp/Foundation/NSDefines.hpp,sha256=c8BfIDwW9yCKiofyatyeLJ9YvnHO87PsBqSQtciN3tY,2203
+mlx/include/metal_cpp/Foundation/NSDictionary.hpp,sha256=Lt7EBR6ppPFhhnoZq9UqMETHf_L1k9A0gs637ZTEQSU,5738
+mlx/include/metal_cpp/Foundation/NSEnumerator.hpp,sha256=P4mkVmpWN0kdXnqqTMkyIovckr4_IZAoFzyg2IZDRs4,3141
+mlx/include/metal_cpp/Foundation/NSError.hpp,sha256=rFqCSIXVsWOsH7Ikegv2irMF3aoMn7fM6teB9VipMgc,7802
+mlx/include/metal_cpp/Foundation/NSLock.hpp,sha256=OjnoMx18m4E89eDPliPMw3cG8LtAqNsoBSxSO0l9DoY,4350
+mlx/include/metal_cpp/Foundation/NSNotification.hpp,sha256=rpAgI4upmIG2inWZ9ACYy-Q8za7muZ2WmUtA5l8WGHw,4696
+mlx/include/metal_cpp/Foundation/NSNumber.hpp,sha256=4a2fEU1Fm_wrWOs4zULhYQne0WymnRGsDeSGq1bDbuQ,22098
+mlx/include/metal_cpp/Foundation/NSObjCRuntime.hpp,sha256=SHYqTLWQ2F8T4dfXJ98OlATOexXQ3NnXzmFyVpb4wFM,1665
+mlx/include/metal_cpp/Foundation/NSObject.hpp,sha256=T_64vByLI4n3HIE05MVJZteUIAdN-y0oogsczYT-JTE,11105
+mlx/include/metal_cpp/Foundation/NSPrivate.hpp,sha256=JYewPlM0BKHLnmQiItPlYuc1HNAO-akJjhXnEKUZO00,20217
+mlx/include/metal_cpp/Foundation/NSProcessInfo.hpp,sha256=1NLI7-3mR8cseE_jE6c-D2_pRJKXA9fRl4V2uRZtiDI,16642
+mlx/include/metal_cpp/Foundation/NSRange.hpp,sha256=7VFC565QzvzfF1CI91osTozl-d1mM9UToCcGyXmf5cc,3147
+mlx/include/metal_cpp/Foundation/NSSet.hpp,sha256=0rORF6rpK2-R1AQCgT4CMNnGr18xZ8P5x82LEQLQ2sc,3368
+mlx/include/metal_cpp/Foundation/NSSharedPtr.hpp,sha256=208e_0hJddLllsT16pJa4rQAyFqA1uYuez5gWyRS6EU,8496
+mlx/include/metal_cpp/Foundation/NSString.hpp,sha256=AqRvSkbHPncbC36TJuR-4beEVT7ZPwnrVjmyoassaL8,11001
+mlx/include/metal_cpp/Foundation/NSTypes.hpp,sha256=9IaLleaqYbdjlMvT_mqFlYDlzg4vF_8MQQTO2Sdp6vo,1893
+mlx/include/metal_cpp/Foundation/NSURL.hpp,sha256=b1oGmKqL5F1WrWSaCK3Iqs0vxL8utVbOgtYW4RZUe_I,3668
+mlx/include/metal_cpp/LICENSE.txt,sha256=mgnj5ca95epXBMNZliB1O16xkRb5_WScN7W8RJeiXOo,11344
+mlx/include/metal_cpp/Metal/MTLAccelerationStructure.hpp,sha256=jqK65WOzpPUNE6H-XYb1PC3myk6vBIQr4fRrClCDZYs,85372
+mlx/include/metal_cpp/Metal/MTLAccelerationStructureCommandEncoder.hpp,sha256=EC1flOCenKN-CNLIbHcpli7kP1PQGxkGUU8ZdMYQ4Qs,16289
+mlx/include/metal_cpp/Metal/MTLAccelerationStructureTypes.hpp,sha256=I29afo7zt4fiCY2ip3vPN72lhb6Whgjpx0lcHabfib4,5800
+mlx/include/metal_cpp/Metal/MTLArgument.hpp,sha256=_gQiklV3kRM2ukcYg0JCz3rUsK-zUOiZMYrp_thHjXU,23515
+mlx/include/metal_cpp/Metal/MTLArgumentEncoder.hpp,sha256=PaV9gzYeC1JLqDdcAJK9H6DaPVV2rclNWvNPK70KUWs,10739
+mlx/include/metal_cpp/Metal/MTLBinaryArchive.hpp,sha256=I8a_7I2WI-ftil5s4bT8N5rnUeRUt2HtsM2O9LW_t2M,5132
+mlx/include/metal_cpp/Metal/MTLBlitCommandEncoder.hpp,sha256=6QF2GR7c4XQXfA5COwCb2Q6xRXVMkYnO1gmUlrNw05o,16482
+mlx/include/metal_cpp/Metal/MTLBlitPass.hpp,sha256=zPt4836TDUuDPyRCyF7tJBEtEMpoWziJLlSRgWleURs,6982
+mlx/include/metal_cpp/Metal/MTLBuffer.hpp,sha256=IDVEI2qXwq6rU_tQsPkOpma-IV7othYmJIT2QCOBsic,3542
+mlx/include/metal_cpp/Metal/MTLCaptureManager.hpp,sha256=ANj6r5G9Gc1l9WikJQmZZlOWvfViJclzyPwKF59UvUg,7524
+mlx/include/metal_cpp/Metal/MTLCaptureScope.hpp,sha256=DU5_HrnsaHBttsqHfXCk-oPr8ifhMdP-DiYpS1G-qC0,3733
+mlx/include/metal_cpp/Metal/MTLCommandBuffer.hpp,sha256=29LVLOZlzB8G0bVfnHm867TXc0FPqbM7eXbNrZ2zRzE,17642
+mlx/include/metal_cpp/Metal/MTLCommandEncoder.hpp,sha256=IQGAhZMikxwU1IINtvw3MSTLwlbQKRWuTNZnZAw_0m8,2946
+mlx/include/metal_cpp/Metal/MTLCommandQueue.hpp,sha256=Uz74PTJAhqMi8XgLrj1RZOio-FYxLnTO9Msdy1dE7RA,2960
+mlx/include/metal_cpp/Metal/MTLComputeCommandEncoder.hpp,sha256=Z2E_Ny8OJj4gvC_YZtku5sni2t4o9Mbm7xvJk2tpgHc,17157
+mlx/include/metal_cpp/Metal/MTLComputePass.hpp,sha256=VOyFJlwaB89gmrtLdmnAz7gl97a5HbTYo-M6icWbN5M,7792
+mlx/include/metal_cpp/Metal/MTLComputePipeline.hpp,sha256=WTBUdVq1sxvHocezHO5wpVmaa94Xnthq27kSUCedcMY,15017
+mlx/include/metal_cpp/Metal/MTLCounters.hpp,sha256=RmlZCNFZIVPC4ue6FoIgYVq1Weq2Nu1GQD-G0yIwimg,9056
+mlx/include/metal_cpp/Metal/MTLDefines.hpp,sha256=1MG6A6zwLk99Yknu7PAvBvIMefVhumr1kFoKxQ-6xKk,1900
+mlx/include/metal_cpp/Metal/MTLDepthStencil.hpp,sha256=nP5OtW6QGUK7tzrXtJWJlch9bXkDYPYgp5BjMAptH6E,9631
+mlx/include/metal_cpp/Metal/MTLDevice.hpp,sha256=HyHg83T53RuLU1y_ZthBvEIzpOi1Lu0nXptaJfdrhwc,63880
+mlx/include/metal_cpp/Metal/MTLDrawable.hpp,sha256=07X3GSg7uJEWoOMWxI77cIWdr2NaBfJEXOtPJZr4M44,3193
+mlx/include/metal_cpp/Metal/MTLDynamicLibrary.hpp,sha256=yBUYfobLLHd00UOGALdsEF0FBiF1hUDr7G3FdKy6uMY,2606
+mlx/include/metal_cpp/Metal/MTLEvent.hpp,sha256=W8QBQWIUn-DcuBYzg0NB0L1W_Zs6G6QuQCeKUKM7tIo,4976
+mlx/include/metal_cpp/Metal/MTLFence.hpp,sha256=DuxPaIkQN9fWi1E6RHg8YrKNeB2nxGOZMuHXd8TiE9c,1723
+mlx/include/metal_cpp/Metal/MTLFunctionConstantValues.hpp,sha256=dz9s7_TYqjTJbthwEJTDoqdXsZ34JE2k4BERCYzWqWQ,3088
+mlx/include/metal_cpp/Metal/MTLFunctionDescriptor.hpp,sha256=JkzBrxhlg4_ezt2NLomw-MsGuBqdV9embOekJow-mks,5410
+mlx/include/metal_cpp/Metal/MTLFunctionHandle.hpp,sha256=HykCHezEwao-pmH6eL57qOyNt_G5RYsrsofxFDyHjiQ,1842
+mlx/include/metal_cpp/Metal/MTLFunctionLog.hpp,sha256=ZiIPh0Ef5oeiYIWsFZJGnOXZj2bejYFDe_3NVoRT2nY,3262
+mlx/include/metal_cpp/Metal/MTLFunctionStitching.hpp,sha256=fEipcQvZ2d8GeyPyn0AiyMup4doS9Yf-uzZRFpjkfyI,11407
+mlx/include/metal_cpp/Metal/MTLHeaderBridge.hpp,sha256=qYNlYMWNAx4z4iNnMyrqdQr_NJ5md2n1HMSxmP5ykRM,105217
+mlx/include/metal_cpp/Metal/MTLHeap.hpp,sha256=O9zMAFEKDGKAt710IDSwdsW3mYSyM7Mt2IVvJNJnAq0,11624
+mlx/include/metal_cpp/Metal/MTLIOCommandBuffer.hpp,sha256=OweZPLAmFoJFypxGIWisfb6Bzlvgv5qIYB2MVu6jBLw,7417
+mlx/include/metal_cpp/Metal/MTLIOCommandQueue.hpp,sha256=Al9BpGDKaSu5zQblMphnr9EXvrVee2ZaZWMba6L7QSs,7465
+mlx/include/metal_cpp/Metal/MTLIOCompressor.hpp,sha256=s7UjSrJLyfQt69RzWBY692G1audMRzjIVR-7GjkZFOQ,3043
+mlx/include/metal_cpp/Metal/MTLIndirectCommandBuffer.hpp,sha256=J3H4HH0p8qq_DYU8pQC_w4qDmYtNqj6FzoOo_AVGfxA,9715
+mlx/include/metal_cpp/Metal/MTLIndirectCommandEncoder.hpp,sha256=AW1Qzyggw9k9EqS71CvlFUHu4lag6nvvzK-LnV4oG0k,11239
+mlx/include/metal_cpp/Metal/MTLIntersectionFunctionTable.hpp,sha256=NUP7rHPWulH7M2M1wvzlyjjF7S0BJvG-_XFn4nWehcg,8181
+mlx/include/metal_cpp/Metal/MTLLibrary.hpp,sha256=le2P__N8XA4RCN62L2ZB-F9uo2p8BjQ-A6j0kEOqzws,23904
+mlx/include/metal_cpp/Metal/MTLLinkedFunctions.hpp,sha256=FmUoNjOGSegryBa0PDa1qcxEF0RUESixbd7OPyU-qDM,3867
+mlx/include/metal_cpp/Metal/MTLParallelRenderCommandEncoder.hpp,sha256=uHyC4bf8BjFUz6xEkFg79Lqzz9CV-0NNnNoQ-U3xfsA,3899
+mlx/include/metal_cpp/Metal/MTLPipeline.hpp,sha256=1EaeWaT8YNfvPaX0gqktNMOLtojBF6cVMlZiFT7vjyg,3775
+mlx/include/metal_cpp/Metal/MTLPixelFormat.hpp,sha256=DYqej_xuztbD0lhsioWTYWO2Sg7jRShyJ7izoVlex5Y,5910
+mlx/include/metal_cpp/Metal/MTLPrivate.hpp,sha256=vPTld95A83b1nIVT4qwsdcrttngMK5L5Ow-0HTr-zm4,6415
+mlx/include/metal_cpp/Metal/MTLRasterizationRate.hpp,sha256=LtcYTnjLSTaxTw98WdUg_9R9BaAozvXOTHcsnmSqbTo,15447
+mlx/include/metal_cpp/Metal/MTLRenderCommandEncoder.hpp,sha256=JNisXL9uSMMm_5B2vkVfuI8dgrMMMIrC8-E7P0pjdeo,60813
+mlx/include/metal_cpp/Metal/MTLRenderPass.hpp,sha256=YQCYdLOnOWsUgIbQimBDFo4l_FmbZfeHIJFADL3NnRQ,32752
+mlx/include/metal_cpp/Metal/MTLRenderPipeline.hpp,sha256=TuBYsiWg4hpumgaKccAn8TVAeLw26l_FdVGa6JsaMLA,72929
+mlx/include/metal_cpp/Metal/MTLResource.hpp,sha256=Be28y-6-9GvFuugCUhnykiTUqbMqYgVovDjUdpIrFDA,5244
+mlx/include/metal_cpp/Metal/MTLResourceStateCommandEncoder.hpp,sha256=eOjzuG-vDtQLJ8QOzhiqxXxiMEShXyn9cn-Uhmtigr8,5330
+mlx/include/metal_cpp/Metal/MTLResourceStatePass.hpp,sha256=N5sMbFCjfEdf10_3VZJfQvDl3nLuydqrKnUgiJzynRU,7585
+mlx/include/metal_cpp/Metal/MTLSampler.hpp,sha256=0rifEvFw4V8_kMmusAU4NXiZeZbCSX_L4x0K6QyEkog,10838
+mlx/include/metal_cpp/Metal/MTLStageInputOutputDescriptor.hpp,sha256=LfDoa6Dj5WQFS2Qg_FQRCkVgBpbbfjBdFtn1SIzsCyI,13120
+mlx/include/metal_cpp/Metal/MTLTexture.hpp,sha256=jM-mSxqfGBNXE6xwzx1M7CC8GOiyungRuop-6uTuf-M,24807
+mlx/include/metal_cpp/Metal/MTLTypes.hpp,sha256=9_le5gsR93J2awZRAk_Rep_yAWKShAQOgI6faZKD_tg,4389
+mlx/include/metal_cpp/Metal/MTLVersion.hpp,sha256=iaG4OsKd3JlcsoLEvyAx_5TdCRup1ykH0FUYqbu6Lxo,1509
+mlx/include/metal_cpp/Metal/MTLVertexDescriptor.hpp,sha256=4oRbxIiJV2eafs-SqTXTVgPttIC22KtTZ0t4AXnLlgc,12017
+mlx/include/metal_cpp/Metal/MTLVisibleFunctionTable.hpp,sha256=P5Cv_bIIfJ4ktplCeZ3fAhHtLEXngF5cAlQgAr55c-0,3833
+mlx/include/metal_cpp/Metal/Metal.hpp,sha256=Bm5ldz-FxywsGmAdjZ8U2X5hlYjwfzHUjKCoW38T-w8,3190
+mlx/include/metal_cpp/MetalFX/MTLFXDefines.hpp,sha256=Ms80cxZbVVAFYDogYahvmzyVZqkDezrwXKXB_79m3KE,2104
+mlx/include/metal_cpp/MetalFX/MTLFXPrivate.hpp,sha256=LcubGmMx98T8OduyEawEx8swCEo_n5Z3ZfwuxwvAtz0,15147
+mlx/include/metal_cpp/MetalFX/MTLFXSpatialScaler.hpp,sha256=8QCySJ9V5JfuTNAHujkvgKQJ_hdM-yV8cqdw5v8l1Ns,18385
+mlx/include/metal_cpp/MetalFX/MTLFXTemporalScaler.hpp,sha256=57Qfhq3YqYOd0qwLfOBvz_viu55-Pmor0_IF99UWMAg,32887
+mlx/include/metal_cpp/MetalFX/MetalFX.hpp,sha256=Vm0W_ycCRb0UOOvLIPYm5gIEHW1nikG5vTvoKfTgAtE,1350
+mlx/include/metal_cpp/QuartzCore/CADefines.hpp,sha256=q6k-jUxe5uulNnsoqkM61OuyEVSJO4W8gIAVM6mSbwE,1895
+mlx/include/metal_cpp/QuartzCore/CAMetalDrawable.hpp,sha256=RE4_Cp9FacJMh7AlqZw-IhelLMXvUv_YpXVLQ2aNVrU,2363
+mlx/include/metal_cpp/QuartzCore/CAMetalLayer.hpp,sha256=i4WjUgWUsQVRPj8dxk4aBC06TKaXsKxRdgNNIgNe0pk,5502
+mlx/include/metal_cpp/QuartzCore/CAPrivate.hpp,sha256=aFEVW34Ih_Gs1aGv5Rml0pNaFxDwWLP6qBQWw3N6EzA,5010
+mlx/include/metal_cpp/QuartzCore/QuartzCore.hpp,sha256=-fQqOpndpszueKWJ4U_jovMID_s0QxCOJHpQoUj8JBE,1346
+mlx/include/metal_cpp/README.md,sha256=ZVbXv3dwSy-n_EmPpYlB2cJXkUCZXdE0-jJsY9pHsd0,14584
+mlx/include/metal_cpp/SingleHeader/MakeSingleHeader.py,sha256=JqlbeaangDoU2dRPyOfKgU5jkbmUghucuumqPEbZOWc,8984
+mlx/include/metal_cpp/SingleHeader/__pycache__/MakeSingleHeader.cpython-311.pyc,,
+mlx/include/mlx/3rdparty/pocketfft.h,sha256=Dq6iEwS_MkY5_xECLuY_r0io-rZK_lw8LR7---HX36Y,110508
+mlx/include/mlx/allocator.h,sha256=r0AgIuBbvq67OeO8eUD8cb4a72HDF6-PISAnOMfC-kM,1519
+mlx/include/mlx/array.h,sha256=D_dkEyq5zbBSc8kK5KNgQbAikvihW9Xn_1HlwI9uaps,11348
+mlx/include/mlx/backend/accelerate/utils.h,sha256=ryQ9t4EI6UkMBCE48MVcGUSstbRDlwM9YI__vs4dTzY,752
+mlx/include/mlx/backend/common/arange.h,sha256=J44RZZx_DB7DSELCtDFHBThHfLopu_JWj_5ew_q0ROw,1823
+mlx/include/mlx/backend/common/binary.h,sha256=uvS_NdciYF-6VHhKrDNFL1Sd2Ag8hn_kClpYmcu6E9E,15169
+mlx/include/mlx/backend/common/copy.h,sha256=5-b2CRTRoyCJ4WL5RkPAQXaLKaEn8pz-XojHGs4u1Rc,690
+mlx/include/mlx/backend/common/erf.h,sha256=Wm6jXR40EgAmyh0zbGOvRbZh58LVhe3_P1-hwN9GPuo,274
+mlx/include/mlx/backend/common/reduce.h,sha256=iR0xKDMZSEDrf2W0TLLXuqxVesf-nOioGu0MDGalT-0,11138
+mlx/include/mlx/backend/common/threefry.h,sha256=sXTdUFRAgEXf789u7IhWIj7n-Gtzwo4HiQXDA6BwFyU,572
+mlx/include/mlx/backend/common/unary.h,sha256=5vAiPfmNjXxMohP6RqZiOngCXcXa76RtrLDEyTeQobs,3356
+mlx/include/mlx/backend/common/utils.h,sha256=4pS1LCC-Tux3x4Y0A3LeTyTLGDSxbjapNkNFWUNjrc8,610
+mlx/include/mlx/backend/metal/allocator.h,sha256=lluCLBSXYbp6WGtl9zKsikpMHrptRjeZJ0Cf70zFkoc,1491
+mlx/include/mlx/backend/metal/copy.h,sha256=jXW0_jDwX97NmnwigeZwovS_MZ1PeZix2iuZf7f9Xts,400
+mlx/include/mlx/backend/metal/device.h,sha256=1Q0xSLn21dTzD5rWanoBBg1eMqOPf2958Z0rJL9GJag,2196
+mlx/include/mlx/backend/metal/kernels/atomic.h,sha256=0Iph_jNZoTyTrR-iEXXDzIwFoRWOkhqndX6IHJxFhKg,9217
+mlx/include/mlx/backend/metal/kernels/bf16.h,sha256=YdmOaznu1O1UNTyGlnpnEkvi1r44EouYNBAZ4SAX2yA,11904
+mlx/include/mlx/backend/metal/kernels/bf16_math.h,sha256=VRM_iZpBcY3Sg3yxkNPzjU30x32zcmd0diIagA1D5l4,26337
+mlx/include/mlx/backend/metal/kernels/complex.h,sha256=XDYKDiBdeM5KGpfpmwXmelZ00LmFCdxiGOIGoUMEeyI,3526
+mlx/include/mlx/backend/metal/kernels/conv_params.h,sha256=pzJn2uZ8TqY1WXRyGRx-MzBYBOfMSyhiDPgu8LMSElk,592
+mlx/include/mlx/backend/metal/kernels/defines.h,sha256=Bkplnhk_KYkqUzpB9TSlYUuXyucJQZYNyW9AR55g8fM,477
+mlx/include/mlx/backend/metal/kernels/erf.h,sha256=k1CfVsb-rEQaXGYJcAbgVR90VwEY5w_88BhJCSsj97k,2736
+mlx/include/mlx/backend/metal/kernels/gemm/conv.h,sha256=ilmFKKdzy8qwn1Zt0cu43f_mKsMLhIJy6DgIdx5ZU5I,14236
+mlx/include/mlx/backend/metal/kernels/gemm/gemm.h,sha256=x80JGVUliIDkFyzoGGd5iC0veHZVPx_BFkeVdGzgVOk,16184
+mlx/include/mlx/backend/metal/kernels/reduce.h,sha256=RgBsV24CtUU_dXCNzDP8sojJdjRz2bdF3PqB8UNJrlA,3554
+mlx/include/mlx/backend/metal/kernels/utils.h,sha256=mRGC9B2cjzCZT2kxY__EM0QmRQQBL8GSHC0rg7_4cJo,7202
+mlx/include/mlx/backend/metal/matmul.h,sha256=n32FxTg1zL-_oIaPRrs9QcI5TywhMemcm2KJ4zF-n0Q,592
+mlx/include/mlx/backend/metal/metal.h,sha256=9kgNkbx0Pqy-W3o7stP0cjvSAcoU0I20dZEsKmsSAS4,552
+mlx/include/mlx/backend/metal/mps/gemm.h,sha256=501lddlggxogceZI3wkWK98yIZurE0qgVxx0ICHyBxs,11307
+mlx/include/mlx/backend/metal/utils.h,sha256=ape52Rl1et96MA4H_a0cbVveLvhASJlszqX6mxl-nis,4205
+mlx/include/mlx/device.h,sha256=KtG7G03J5T0Ad2zn931OyRGfCa1XKn3ZyEUoIBY0No4,563
+mlx/include/mlx/dtype.h,sha256=jrOMl7xEhOZ7Dn0udDfi-Mq8QbT3atbd7CL0SZ4Yxjk,2575
+mlx/include/mlx/fft.h,sha256=MpW4UMjetc81i8pt3NJnTCvtXuaP92CX3w2iatYZHTw,4179
+mlx/include/mlx/graph_utils.h,sha256=23pv6xyLYTl-l2o9Z3DqlcvGVGZEA-uIV29jamcK-1Q,593
+mlx/include/mlx/io/load.h,sha256=vl7Z0pSJTjgG7KkSLXbe5oLjPhnVkWR9XdvC-ZCf_fk,2508
+mlx/include/mlx/io/safetensor.h,sha256=ro-grE_iXiGbCgygGFSeqsvxD_DDiwKckCsyb424e-g,624
+mlx/include/mlx/linalg.h,sha256=-CJ8NhHmycVsT29ObUmVZbxd2gwCpj8qj2RvOv-Glto,1771
+mlx/include/mlx/mlx.h,sha256=raZfU4bjGTSRAZa8P1XkoxHHeiooTZVI9C63ydGwm8c,296
+mlx/include/mlx/ops.h,sha256=322gdgz9l7pz-M0Pgpo0T6LbAe_w0RMV-A1-jRN_Le0,32224
+mlx/include/mlx/primitives.h,sha256=uJhmcwsXyseQlfeJBDjLCNrhPM91vursiw8_R79yUGg,44227
+mlx/include/mlx/random.h,sha256=aNoyWf-zAg9WKgqfe1yBpOvxedygIUZMQBgsd3kNnzQ,4887
+mlx/include/mlx/scheduler.h,sha256=sZX0BU1J15e_IOzBsv95KwAQmU4wrayd3b6ClvxGA4M,3864
+mlx/include/mlx/stream.h,sha256=7BvNtylObXbyIHZ0gCrbMb1bN3DUq7PfzEa-bIvqEXE,686
+mlx/include/mlx/transforms.h,sha256=Qt8pjvxAPYMoa8Dz2WUcJBdeKEFKNdczV0O3UTE_XmY,6527
+mlx/include/mlx/transforms_impl.h,sha256=RttUce4KHg3PLUFc9aRVjxenCIztlBSTlQq5LnVxKVg,542
+mlx/include/mlx/types/bf16.h,sha256=fLnTB0fbL2vIs_CCvWhrr7SfDnFUEpDPmg7RNmrRnzE,6753
+mlx/include/mlx/types/complex.h,sha256=oaTl_rCth3TO33sT2Db6SX--n9lIfE_dw7rjE06Qqmc,2878
+mlx/include/mlx/types/fp16.h,sha256=j0E7-PKkVkCT-bOMRdBEYdvxUdo9rf6w80dNSNCYNaQ,8322
+mlx/include/mlx/types/half_types.h,sha256=-D-ilDAh5shp1xp4lx0J87y9iFfBs43gLNi-XqzmLWA,1367
+mlx/include/mlx/utils.h,sha256=cchUDEk5qbWNIPHf_cRwbsxnczz2_moUr_Dipz-cfhM,1524
+mlx/lib/libmlx.dylib,sha256=ir7-RqHznJKyhGSBTwWnMPqYmbF3V3A0A8bvNi8GrJM,12420704
+mlx/lib/mlx.metallib,sha256=Lu30EADtJwKD2hHYibsQGqTIjG-PDsaP5rBApb5CRQE,59495531
+mlx/nn/__init__.py,sha256=b-1hqkAnzpthKJNKgRGVHsltNEqOU7URFMoY2Z46Buw,126
+mlx/nn/__pycache__/__init__.cpython-311.pyc,,
+mlx/nn/__pycache__/losses.cpython-311.pyc,,
+mlx/nn/__pycache__/utils.cpython-311.pyc,,
+mlx/nn/layers/__init__.py,sha256=2Ge23mO7W5VOMkogmcY2dqTKRRD3lY0g_WhVDP-9Hi0,1249
+mlx/nn/layers/__pycache__/__init__.cpython-311.pyc,,
+mlx/nn/layers/__pycache__/activations.cpython-311.pyc,,
+mlx/nn/layers/__pycache__/base.cpython-311.pyc,,
+mlx/nn/layers/__pycache__/containers.cpython-311.pyc,,
+mlx/nn/layers/__pycache__/convolution.cpython-311.pyc,,
+mlx/nn/layers/__pycache__/dropout.cpython-311.pyc,,
+mlx/nn/layers/__pycache__/embedding.cpython-311.pyc,,
+mlx/nn/layers/__pycache__/linear.cpython-311.pyc,,
+mlx/nn/layers/__pycache__/normalization.cpython-311.pyc,,
+mlx/nn/layers/__pycache__/positional_encoding.cpython-311.pyc,,
+mlx/nn/layers/__pycache__/quantized.cpython-311.pyc,,
+mlx/nn/layers/__pycache__/transformer.cpython-311.pyc,,
+mlx/nn/layers/activations.py,sha256=ukopw8ZFpSIQijXscov80Gv_0j0jKh0Y4Z3qrUFzpkk,11826
+mlx/nn/layers/base.py,sha256=iFp3PVGUhT9pAaWNnWtNY-yZ-uCQxvH-O1CYbKoGCug,19485
+mlx/nn/layers/containers.py,sha256=Ke8gPBPZvrtrjK1qrbRNGtSV4rvBkYTKVBW2pQXrNvY,618
+mlx/nn/layers/convolution.py,sha256=ZP43OBZTKF3ui_2Xrtx7VvAlLic_jcn2oL-MPLU7Rys,4060
+mlx/nn/layers/dropout.py,sha256=1GR5uWU82eFd296WuPB7CqNcNj3acEIlS3_4ZAxnx-g,4529
+mlx/nn/layers/embedding.py,sha256=5G7TvVyEaHTlQCASNVfBGea_1ywDVNqDkxHP2U-JEcs,878
+mlx/nn/layers/linear.py,sha256=ekgkElIUlfdWZNRDZkoyQFTotzyjNpD_0v7Ac8pnako,4137
+mlx/nn/layers/normalization.py,sha256=m92TwFfNQ4NbQlxtlTfJsWZlBJtO0_tR0M8H2d9wZEs,12021
+mlx/nn/layers/positional_encoding.py,sha256=_4zbYLMGYI0EVq4b6eitmB2hgBXblcrJJYsq0KCSafs,6611
+mlx/nn/layers/quantized.py,sha256=FA_NzLBheO2i5AClMWCkaFrNXJGScxlwvumqaJTIxAA,4136
+mlx/nn/layers/transformer.py,sha256=IUTNGvVoaKLS7Gwes_yK3SqgX6voPgMqk7nmdE4nyI0,12235
+mlx/nn/losses.py,sha256=1yWCxt-QSv4lGV9U6KJyeTzybQWW4LPZhipkLFx7hWk,12158
+mlx/nn/utils.py,sha256=tQJ62bu2eAAhzYqV7QriQtNGrYSz85XGxHolEPkGiGs,999
+mlx/optimizers.py,sha256=krPDwdaMy1JTlaVBgRHlg7PNqyxiVRkfNEbEf_wV-5g,17090
+mlx/share/cmake/MLX/MLXConfig.cmake,sha256=02CzwjTAHLTyLzFMtE7LLzyX_CSHreASjFyiYEpfIp4,1884
+mlx/share/cmake/MLX/MLXConfigVersion.cmake,sha256=p6TWi4EVJlFiJx_BxHZqpPBhkKaze00W3S-JRcRHD78,2762
+mlx/share/cmake/MLX/MLXTargets-release.cmake,sha256=bUjlsca3AGGUHl1X9lfm3CiQiSFeYTkm2c7GLzITeJs,788
+mlx/share/cmake/MLX/MLXTargets.cmake,sha256=vS5sI59URF-Pz8SPKMDVbyTPq1ENZ3zoR_VPzWToXkw,4705
+mlx/share/cmake/MLX/extension.cmake,sha256=q09-X3ckTrFCZOc5UNyoqjUC94OQQSXXWW7Bmg7sxok,1637
+mlx/utils.py,sha256=F-bmWn6a4VFRUvnSeZPvIFwT8Z6IHEtgEAfcLiGBrz8,4667

lib/python3.11/site-packages/mlx-0.0.7.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: bdist_wheel (0.41.2)
+Root-Is-Purelib: false
+Tag: cp311-cp311-macosx_14_0_arm64
+

lib/python3.11/site-packages/mlx-0.0.7.dist-info/top_level.txt

@@ -0,0 +1 @@
+mlx

lib/python3.11/site-packages/mlx/nn/layers/base.py

@@ -0,0 +1,532 @@
+# Copyright © 2023 Apple Inc.
+
+import textwrap
+from typing import Any, Callable, List, Optional, Tuple, Union
+
+import mlx.core as mx
+from mlx.utils import tree_flatten, tree_unflatten
+
+
+class Module(dict):
+    """Base class for building neural networks with MLX.
+
+    All the layers provided in :mod:`mlx.nn.layers` subclass this class and
+    your models should do the same.
+
+    A ``Module`` can contain other ``Module`` instances or :class:`mlx.core.array`
+    instances in arbitrary nesting of python lists or dicts. The ``Module``
+    then allows recursively extracting all the :class:`mlx.core.array` instances
+    using :meth:`mlx.nn.Module.parameters`.
+
+    In addition, the ``Module`` has the concept of trainable and non trainable
+    parameters (called "frozen"). When using :func:`mlx.nn.value_and_grad`
+    the gradients are returned only with respect to the trainable parameters.
+    All arrays in a module are trainable unless they are added in the "frozen"
+    set by calling :meth:`freeze`.
+
+    .. code-block:: python
+
+        import mlx.core as mx
+        import mlx.nn as nn
+
+        class MyMLP(nn.Module):
+            def __init__(self, in_dims: int, out_dims: int, hidden_dims: int = 16):
+                super().__init__()
+
+                self.in_proj = nn.Linear(in_dims, hidden_dims)
+                self.out_proj = nn.Linear(hidden_dims, out_dims)
+
+            def __call__(self, x):
+                x = self.in_proj(x)
+                x = mx.maximum(x, 0)
+                return self.out_proj(x)
+
+        model = MyMLP(2, 1)
+
+        # All the model parameters are created but since MLX is lazy by
+        # default, they are not evaluated yet. Calling `mx.eval` actually
+        # allocates memory and initializes the parameters.
+        mx.eval(model.parameters())
+
+        # Setting a parameter to a new value is as simply as accessing that
+        # parameter and assigning a new array to it.
+        model.in_proj.weight = model.in_proj.weight * 2
+        mx.eval(model.parameters())
+    """
+
+    def __init__(self):
+        """Should be called by the subclasses of ``Module``."""
+        self._no_grad = set()
+        self._training = True
+
+    @property
+    def training(self):
+        """Boolean indicating if the model is in training mode."""
+        return self._training
+
+    def _extra_repr(self):
+        return ""
+
+    def __repr__(self):
+        children = tree_flatten(self.children(), is_leaf=self.is_module)
+        value = f"{type(self).__name__}({self._extra_repr()}"
+        for k, v in children:
+            value += "\n"
+            value += textwrap.indent(f"({k}): {repr(v)}", prefix="  ")
+        if children:
+            value += "\n"
+        value += ")"
+
+        return value
+
+    def __getattr__(self, key: str):
+        if key in self:
+            return self[key]
+        else:
+            raise AttributeError(f"{type(self)!r} has no attribute {key!r}")
+
+    def __setattr__(self, key: str, val: Any):
+        self[key] = val
+
+    def load_weights(
+        self,
+        file_or_weights: Union[str, List[Tuple[str, mx.array]]],
+        strict: bool = True,
+    ):
+        """
+        Update the model's weights from a ``.npz`` or a list.
+
+        Args:
+            file_or_weights (str or list(tuple(str, mx.array))): The path to
+                the weights ``.npz`` file or a list of pairs of parameter names
+                and arrays.
+            strict (bool, optional): If ``True`` then checks that the provided
+              weights exactly match the parameters of the model. Otherwise,
+              only the weights actually contained in the model are loaded and
+              shapes are not checked. Default: ``True``.
+
+        Example:
+
+            .. code-block:: python
+
+                import mlx.core as mx
+                import mlx.nn as nn
+                model = nn.Linear(10, 10)
+
+                # Load from file
+                model.load_weights("weights.npz")
+
+                # Load from list
+                weights = [
+                    ("weight", mx.random.uniform(shape=(10, 10))),
+                    ("bias",  mx.zeros((10,))),
+                ]
+                model.load_weights(weights)
+
+                # Missing weight
+                weights = [
+                    ("weight", mx.random.uniform(shape=(10, 10))),
+                ]
+
+                # Raises a ValueError exception
+                model.load_weights(weights)
+
+                # Ok, only updates the weight but not the bias
+                model.load_weights(weights, strict=False)
+        """
+        weights = file_or_weights
+        if isinstance(weights, str):
+            weights = list(mx.load(weights).items())
+
+        if strict:
+            new_weights = dict(weights)
+            curr_weights = dict(tree_flatten(self.parameters()))
+            if extras := (new_weights.keys() - curr_weights.keys()):
+                extras = " ".join(extras)
+                raise ValueError(f"Received parameters not in model: {extras}.")
+            if missing := (curr_weights.keys() - new_weights.keys()):
+                missing = " ".join(missing)
+                raise ValueError(f"Missing parameters: {missing}.")
+            for k, v in curr_weights.items():
+                v_new = new_weights[k]
+                if not isinstance(v_new, mx.array):
+                    raise ValueError(
+                        "Expected mx.array but received "
+                        f"{type(v_new)} for parameter {k}"
+                    )
+                if v_new.shape != v.shape:
+                    raise ValueError(
+                        f"Expected shape {v.shape} but received "
+                        f" shape {v_new.shape} for parameter {k}"
+                    )
+
+        self.update(tree_unflatten(weights))
+
+    def save_weights(self, file: str):
+        """
+        Save the model's weights to a ``.npz`` file.
+        """
+        mx.savez(file, **dict(tree_flatten(self.parameters())))
+
+    @staticmethod
+    def is_module(value):
+        return isinstance(value, Module)
+
+    @staticmethod
+    def valid_child_filter(module, key, value):
+        return isinstance(value, (dict, list))
+
+    @staticmethod
+    def valid_parameter_filter(module, key, value):
+        return isinstance(value, (dict, list, mx.array)) and not key.startswith("_")
+
+    @staticmethod
+    def trainable_parameter_filter(module, key, value):
+        return (
+            Module.valid_parameter_filter(module, key, value)
+            and key not in module._no_grad
+        )
+
+    def filter_and_map(
+        self,
+        filter_fn: Callable[["mlx.nn.Module", str, Any], bool],
+        map_fn: Optional[Callable] = None,
+        is_leaf_fn: Optional[Callable[["mlx.nn.Module", str, Any], bool]] = None,
+    ):
+        """Recursively filter the contents of the module using ``filter_fn``,
+        namely only select keys and values where ``filter_fn`` returns true.
+
+        This is used to implement :meth:`parameters` and :meth:`trainable_parameters`
+        but it can also be used to extract any subset of the module's parameters.
+
+        Args:
+            filter_fn (Callable): Given a value, the key in which it is found
+                and the containing module, decide whether to keep the value or
+                drop it.
+            map_fn (Callable, optional): Optionally transform the value before
+                returning it.
+            is_leaf_fn (Callable, optional): Given a value, the key in which it
+                is found and the containing module decide if it is a leaf.
+
+        Returns:
+            A dictionary containing the contents of the module recursively filtered
+        """
+
+        map_fn = map_fn or (lambda x: x)
+        is_leaf_fn = is_leaf_fn or (
+            lambda m, k, v: not isinstance(v, (Module, dict, list))
+        )
+
+        def unwrap(vk, v):
+            if is_leaf_fn(self, vk, v):
+                return map_fn(v)
+
+            if isinstance(v, Module):
+                return v.filter_and_map(filter_fn, map_fn, is_leaf_fn)
+
+            if isinstance(v, dict):
+                nd = {}
+                for k, v in v.items():
+                    tk = f"{vk}.{k}"
+                    nd[k] = unwrap(tk, v) if filter_fn(self, tk, v) else {}
+                return nd
+
+            if isinstance(v, list):
+                nl = []
+                for i, vi in enumerate(v):
+                    tk = f"{vk}.{i}"
+                    nl.append(unwrap(tk, vi) if filter_fn(self, tk, vi) else {})
+                return nl
+
+            raise RuntimeError("Unexpected leaf found while traversing the module")
+
+        return {k: unwrap(k, v) for k, v in self.items() if filter_fn(self, k, v)}
+
+    def parameters(self):
+        """Recursively return all the :class:`mlx.core.array` members of this Module
+        as a dict of dicts and lists."""
+        return self.filter_and_map(self.valid_parameter_filter)
+
+    def trainable_parameters(self):
+        """Recursively return all the non frozen :class:`mlx.core.array` members of
+        this Module as a dict of dicts and lists."""
+        return self.filter_and_map(self.trainable_parameter_filter)
+
+    def children(self):
+        """Return the direct descendants of this Module instance."""
+        return self.filter_and_map(
+            self.valid_child_filter, is_leaf_fn=lambda m, k, v: isinstance(v, Module)
+        )
+
+    def leaf_modules(self):
+        """Return the submodules that do not contain other modules."""
+
+        def _is_leaf_module(m, k, v):
+            return isinstance(v, Module) and len(tree_flatten(v.children())) == 0
+
+        return self.filter_and_map(self.valid_child_filter, is_leaf_fn=_is_leaf_module)
+
+    def update(self, parameters: dict):
+        """Replace the parameters of this Module with the provided ones in the
+        dict of dicts and lists.
+
+        Commonly used by the optimizer to change the model to the updated
+        (optimized) parameters. Also used by the :meth:`mlx.nn.value_and_grad` to set the
+        tracers in the model in order to compute gradients.
+
+        The passed in parameters dictionary need not be a full dictionary
+        similar to :meth:`parameters`. Only the provided locations will be
+        updated.
+
+        Args:
+            parameters (dict): A complete or partial dictionary of the modules
+                               parameters.
+        """
+
+        def apply(dst, parameters):
+            if isinstance(parameters, dict):
+                for k in parameters:
+                    if k in dst:
+                        current_value = dst[k]
+                        new_value = parameters[k]
+                        if isinstance(current_value, mx.array):
+                            dst[k] = new_value
+                        elif isinstance(current_value, Module):
+                            current_value.update(new_value)
+                        elif isinstance(current_value, (dict, list)):
+                            apply(current_value, new_value)
+            elif isinstance(parameters, list):
+                for i in range(len(dst)):
+                    current_value = dst[i]
+                    new_value = parameters[i]
+                    if isinstance(current_value, mx.array):
+                        dst[i] = new_value
+                    elif isinstance(current_value, Module):
+                        current_value.update(new_value)
+                    elif isinstance(current_value, (dict, list)):
+                        apply(current_value, new_value)
+
+        apply(self, parameters)
+
+    def apply(
+        self,
+        map_fn: Callable[[mx.array], mx.array],
+        filter_fn: Optional[Callable[["mlx.nn.Module", str, Any], bool]] = None,
+    ):
+        """Map all the parameters using the provided ``map_fn`` and immediately
+        update the module with the mapped parameters.
+
+        For instance running ``model.apply(lambda x: x.astype(mx.float16))``
+        casts all parameters to 16 bit floats.
+
+        Args:
+            map_fn (Callable): Maps an array to another array
+            filter_fn (Callable, optional): Filter to select which arrays to
+                map (default: :meth:`Module.valid_parameter_filter`).
+        """
+        filter_fn = filter_fn or Module.valid_parameter_filter
+        self.update(self.filter_and_map(filter_fn, map_fn))
+
+    def update_modules(self, modules: dict):
+        """Replace the child modules of this :class:`Module` instance with the
+        provided ones in the dict of dicts and lists.
+
+        It is the equivalent of :meth:`Module.update` but for modules instead
+        of parameters and allows us to flexibly edit complex architectures by
+        programmatically swapping layers.
+
+        The passed in parameters dictionary need not be a full dictionary
+        similar to :meth:`parameters`. Only the provided locations will be
+        updated.
+
+        Args:
+            modules (dict): A complete or partial dictionary of the modules
+                submodules.
+        """
+
+        def apply(dst, modules):
+            if isinstance(modules, dict):
+                for k in modules:
+                    if k in dst:
+                        current_value = dst[k]
+                        new_value = modules[k]
+                        if self.is_module(current_value) and self.is_module(new_value):
+                            dst[k] = new_value
+                        elif isinstance(current_value, (dict, list)):
+                            apply(current_value, new_value)
+            elif isinstance(modules, list):
+                for i in range(len(dst)):
+                    current_value = dst[i]
+                    new_value = modules[i]
+                    if self.is_module(current_value) and self.is_module(new_value):
+                        dst[i] = new_value
+                    elif isinstance(current_value, (dict, list)):
+                        apply(current_value, new_value)
+
+        apply(self, modules)
+
+    def apply_to_modules(self, apply_fn: Callable[[str, "mlx.nn.Module"], Any]):
+        """Apply a function to all the modules in this instance (including this
+        instance).
+
+        Args:
+            apply_fn (Callable): The function to apply to the modules.
+        """
+        module_stack = [("", self)]
+        while module_stack:
+            prefix, mod = module_stack.pop()
+            apply_fn(prefix, mod)
+            prefix = "." + prefix if prefix else ""
+            module_stack.extend(
+                tree_flatten(mod.children(), prefix=prefix, is_leaf=self.is_module)
+            )
+
+    def modules(self):
+        """Return a list with all the modules in this instance.
+
+        Returns:
+            A list of :class:`mlx.nn.Module` instances.
+        """
+        modulelist = []
+        self.apply_to_modules(lambda k, m: modulelist.append(m))
+        return modulelist
+
+    def named_modules(self):
+        """Return a list with all the modules in this instance and their name
+        with dot notation.
+
+        Returns:
+            A list of tuples (str, :class:`mlx.nn.Module`).
+        """
+        modulelist = []
+        self.apply_to_modules(lambda k, m: modulelist.append((k, m)))
+        return modulelist
+
+    def _validate_keys(self, keys, strict):
+        keys = keys if isinstance(keys, list) else [keys]
+        if strict:
+            for k in keys:
+                if k not in self:
+                    raise KeyError(f"Module doesn't contain member {k}.")
+        return keys
+
+    def freeze(
+        self,
+        *,
+        recurse: bool = True,
+        keys: Optional[Union[str, List[str]]] = None,
+        strict: bool = False,
+    ):
+        """Freeze the Module's parameters or some of them. Freezing a parameter means not
+        computing gradients for it.
+
+        This function is idempotent i.e. freezing a frozen model is a no-op.
+
+        Example:
+            For instance to only train the attention parameters from a Transformer:
+
+            .. code-block:: python
+
+                model = nn.Transformer()
+                model.freeze()
+                model.apply_to_modules(lambda k, v: v.unfreeze() if k.endswith("attention") else None)
+
+        Args:
+            recurse (bool, optional): If True then freeze the parameters of the
+                submodules as well. Default: ``True``.
+            keys (str or list[str], optional): If provided then only these
+                parameters will be frozen otherwise all the parameters of a
+                module. For instance freeze all biases by calling
+                ``module.freeze(keys="bias")``.
+            strict (bool, optional): If set to ``True`` validate that the passed keys exist.
+                Default: ``False``.
+        """
+
+        def _freeze_impl(_, m):
+            local_keys = keys
+            if local_keys is None:
+                local_keys = tree_flatten(
+                    m.filter_and_map(
+                        lambda m, k, v: (not isinstance(v, Module))
+                        and m.valid_parameter_filter(m, k, v)
+                    )
+                )
+                local_keys = [k for (k, v) in local_keys]
+
+            local_keys = m._validate_keys(local_keys, strict)
+            m._no_grad.update(local_keys)
+
+        if recurse:
+            self.apply_to_modules(_freeze_impl)
+        else:
+            _freeze_impl("", self)
+
+    def unfreeze(
+        self,
+        *,
+        recurse: bool = True,
+        keys: Optional[Union[str, List[str]]] = None,
+        strict: bool = False,
+    ):
+        """Unfreeze the Module's parameters or some of them.
+
+        This function is idempotent ie unfreezing a model that is not frozen is
+        a noop.
+
+        Example:
+
+            For instance to only train the biases of a Transformer one can do:
+
+            .. code-block:: python
+
+                model = nn.Transformer()
+                model.freeze()
+                model.unfreeze(keys="bias")
+
+        Args:
+            recurse (bool, optional): If True then unfreeze the parameters of the
+                submodules as well. Default: ``True``.
+            keys (str or list[str], optional): If provided then only these
+                parameters will be unfrozen otherwise all the parameters of a
+                module. For instance unfreeze all biases by calling
+                ``module.unfreeze(keys="bias")``.
+            strict (bool, optional): If set to ``True`` validate that the passed keys exist.
+                Default: ``False``.
+        """
+
+        def _unfreeze_impl(_, m):
+            if keys is None:
+                m._no_grad.clear()
+
+            else:
+                local_keys = m._validate_keys(keys, strict)
+                m._no_grad.difference_update(local_keys)
+
+        if recurse:
+            self.apply_to_modules(_unfreeze_impl)
+        else:
+            _unfreeze_impl("", self)
+
+    def train(self, mode: bool = True):
+        """Set the model in or out of training mode.
+
+        Training mode only applies to certain layers. For example
+        :obj:`Dropout` applies a random mask in training mode, but is the
+        identity in evaluation mode.
+
+        Args:
+            mode (bool): Indicate if the model should be in training or
+                evaluation mode. Default: ``True``.
+        """
+
+        def _set_train(_, m):
+            m._training = mode
+
+        self.apply_to_modules(_set_train)
+
+    def eval(self):
+        """Set the model to evaluation mode.
+
+        See :func:`train`.
+        """
+        self.train(False)

lib/python3.11/site-packages/mlx/nn/layers/containers.py

@@ -0,0 +1,24 @@
+# Copyright © 2023 Apple Inc.
+
+from mlx.nn.layers.base import Module
+
+
+class Sequential(Module):
+    """A layer that calls the passed callables in order.
+
+    We can pass either modules or plain callables to the Sequential module. If
+    our functions have learnable parameters they should be implemented as
+    ``nn.Module`` instances.
+
+    Args:
+        modules (tuple of Callables): The modules to call in order
+    """
+
+    def __init__(self, *modules):
+        super().__init__()
+        self.layers = list(modules)
+
+    def __call__(self, x):
+        for m in self.layers:
+            x = m(x)
+        return x

lib/python3.11/site-packages/mlx/nn/layers/convolution.py

@@ -0,0 +1,126 @@
+# Copyright © 2023 Apple Inc.
+
+import math
+from typing import Union
+
+import mlx.core as mx
+from mlx.nn.layers.base import Module
+
+
+class Conv1d(Module):
+    """Applies a 1-dimensional convolution over the multi-channel input sequence.
+
+    The channels are expected to be last i.e. the input shape should be ``NLC`` where:
+        - ``N`` is the batch dimension
+        - ``L`` is the sequence length
+        - ``C`` is the number of input channels
+
+    Args:
+        in_channels (int): The number of input channels
+        out_channels (int): The number of output channels
+        kernel_size (int): The size of the convolution filters
+        stride (int, optional): The stride when applying the filter.
+            Default: 1.
+        padding (int, optional): How many positions to 0-pad the input with.
+            Default: 0.
+        bias (bool, optional): If ``True`` add a learnable bias to the output.
+            Default: ``True``
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        padding: int = 0,
+        bias: bool = True,
+    ):
+        super().__init__()
+
+        scale = math.sqrt(1 / (in_channels * kernel_size))
+        self.weight = mx.random.uniform(
+            low=-scale,
+            high=scale,
+            shape=(out_channels, kernel_size, in_channels),
+        )
+        if bias:
+            self.bias = mx.zeros((out_channels,))
+
+        self.padding = padding
+        self.stride = stride
+
+    def _extra_repr(self):
+        return (
+            f"{self.weight.shape[-1]}, {self.weight.shape[0]}, "
+            f"kernel_size={self.weight.shape[1]}, stride={self.stride}, "
+            f"padding={self.padding}, bias={'bias' in self}"
+        )
+
+    def __call__(self, x):
+        y = mx.conv1d(x, self.weight, self.stride, self.padding)
+        if "bias" in self:
+            y = y + self.bias
+        return y
+
+
+class Conv2d(Module):
+    """Applies a 2-dimensional convolution over the multi-channel input image.
+
+    The channels are expected to be last i.e. the input shape should be ``NHWC`` where:
+        - ``N`` is the batch dimension
+        - ``H`` is the input image height
+        - ``W`` is the input image width
+        - ``C`` is the number of input channels
+
+    Args:
+        in_channels (int): The number of input channels.
+        out_channels (int): The number of output channels.
+        kernel_size (int or tuple): The size of the convolution filters.
+        stride (int or tuple, optional): The size of the stride when
+            applying the filter. Default: 1.
+        padding (int or tuple, optional): How many positions to 0-pad
+            the input with. Default: 0.
+        bias (bool, optional): If ``True`` add a learnable bias to the
+            output. Default: ``True``
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: Union[int, tuple],
+        stride: Union[int, tuple] = 1,
+        padding: Union[int, tuple] = 0,
+        bias: bool = True,
+    ):
+        super().__init__()
+
+        kernel_size, stride, padding = map(
+            lambda x: (x, x) if isinstance(x, int) else x,
+            (kernel_size, stride, padding),
+        )
+        scale = math.sqrt(1 / (in_channels * kernel_size[0] * kernel_size[1]))
+        self.weight = mx.random.uniform(
+            low=-scale,
+            high=scale,
+            shape=(out_channels, *kernel_size, in_channels),
+        )
+        if bias:
+            self.bias = mx.zeros((out_channels,))
+
+        self.padding = padding
+        self.stride = stride
+
+    def _extra_repr(self):
+        return (
+            f"{self.weight.shape[-1]}, {self.weight.shape[0]}, "
+            f"kernel_size={self.weight.shape[1:2]}, stride={self.stride}, "
+            f"padding={self.padding}, bias={'bias' in self}"
+        )
+
+    def __call__(self, x):
+        y = mx.conv2d(x, self.weight, self.stride, self.padding)
+        if "bias" in self:
+            y = y + self.bias
+        return y

lib/python3.11/site-packages/mlx/nn/layers/dropout.py

@@ -0,0 +1,137 @@
+# Copyright © 2023 Apple Inc.
+
+import mlx.core as mx
+from mlx.nn.layers.base import Module
+
+
+class Dropout(Module):
+    r"""Randomly zero a portion of the elements during training.
+
+    The remaining elements are multiplied with :math:`\frac{1}{1-p}` where
+    :math:`p` is the probability of zeroing an element. This is done so the
+    expected value of a given element will remain the same.
+
+    Args:
+        p (float): The probability to zero an element
+    """
+
+    def __init__(self, p: float = 0.5):
+        super().__init__()
+
+        if p < 0 or p >= 1:
+            raise ValueError(f"The dropout probability {p} is not in [0, 1)")
+
+        self._p_1 = 1 - p
+
+    def _extra_repr(self):
+        return f"p={1-self._p_1}"
+
+    def __call__(self, x):
+        if self._p_1 == 1 or not self.training:
+            return x
+
+        mask = mx.random.bernoulli(self._p_1, x.shape)
+
+        return (1 / self._p_1) * mask * x
+
+
+class Dropout2d(Module):
+    r"""Apply 2D channel-wise dropout during training.
+
+    Randomly zero out entire channels independently with probability :math:`p`.
+    This layer expects the channels to be last, i.e. the input shape should be
+    ``NWHC`` or ``WHC`` where:``N`` is the batch dimension,``H`` is the input
+    image height,``W`` is the input image width, and``C`` is the number of
+    input channels
+
+    The remaining channels are scaled by :math:`\frac{1}{1-p}` to
+    maintain the expected value of each element. Unlike traditional dropout,
+    which zeros individual entries, this layer zeros entire channels. This is
+    beneficial for early convolution layers where adjacent pixels are
+    correlated. In such case, traditional dropout may not effectively
+    regularize activations. For more details, see [1].
+
+    [1]: Thompson, J., Goroshin, R., Jain, A., LeCun, Y. and Bregler C., 2015.
+    Efficient Object Localization Using Convolutional Networks. CVPR 2015.
+
+    Args:
+        p (float): Probability of zeroing a channel during training.
+    """
+
+    def __init__(self, p: float = 0.5):
+        super().__init__()
+
+        if p < 0 or p >= 1:
+            raise ValueError(f"The dropout probability {p} is not in [0, 1)")
+
+        self._p_1 = 1 - p
+
+    def _extra_repr(self):
+        return f"p={1-self._p_1}"
+
+    def __call__(self, x):
+        if x.ndim not in (3, 4):
+            raise ValueError(
+                f"Received input with {x.ndim} dimensions. Expected 3 or 4 dimensions."
+            )
+
+        if self._p_1 == 1 or not self.training:
+            return x
+
+        # Dropout is applied on the whole channel
+        # 3D input: (1, 1, C)
+        # 4D input: (B, 1, 1, C)
+        mask_shape = x.shape
+        mask_shape[-2] = mask_shape[-3] = 1
+
+        mask = mx.random.bernoulli(p=self._p_1, shape=mask_shape)
+        return (1 / self._p_1) * mask * x
+
+
+class Dropout3d(Module):
+    r"""Apply 3D channel-wise dropout during training.
+
+    Randomly zero out entire channels independently with probability :math:`p`.
+    This layer expects the channels to be last, i.e., the input shape should be
+    `NDHWC` or `DHWC` where: `N` is the batch dimension, `D` is the depth,
+    `H` is the input image height, `W` is the input image width, and `C` is
+    the number of input channels.
+
+    The remaining channels are scaled by :math:`\frac{1}{1-p}` to
+    maintain the expected value of each element. Unlike traditional dropout,
+    which zeros individual entries, this layer zeros entire channels. This is
+    often beneficial for convolutional layers processing 3D data, like in
+    medical imaging or video processing.
+
+    Args:
+        p (float): Probability of zeroing a channel during training.
+    """
+
+    def __init__(self, p: float = 0.5):
+        super().__init__()
+
+        if p < 0 or p >= 1:
+            raise ValueError(f"The dropout probability {p} is not in [0, 1)")
+
+        self._p_1 = 1 - p
+
+    def _extra_repr(self):
+        return f"p={1-self._p_1}"
+
+    def __call__(self, x):
+        if x.ndim not in (4, 5):
+            raise ValueError(
+                f"Received input with {x.ndim} dimensions. Expected 4 or 5 dimensions."
+            )
+
+        if self._p_1 == 1 or not self.training:
+            return x
+
+        # Dropout is applied on the whole channel
+        # 4D input: (1, 1, 1, C)
+        # 5D input: (B, 1, 1, 1, C)
+        mask_shape = list(x.shape)
+        mask_shape[-2] = mask_shape[-3] = mask_shape[-4] = 1
+
+        mask = mx.random.bernoulli(p=self._p_1, shape=mask_shape)
+        return (1 / self._p_1) * mask * x

lib/python3.11/site-packages/mlx/nn/layers/embedding.py

@@ -0,0 +1,30 @@
+# Copyright © 2023 Apple Inc.
+
+import math
+
+import mlx.core as mx
+from mlx.nn.layers.base import Module
+
+
+class Embedding(Module):
+    """Implements a simple lookup table that maps each input integer to a
+    high-dimensional vector.
+
+    Typically used to embed discrete tokens for processing by neural networks.
+
+    Args:
+        num_embeddings (int): How many possible discrete tokens can we embed.
+                              Usually called the vocabulary size.
+        dims (int): The dimensionality of the embeddings.
+    """
+
+    def __init__(self, num_embeddings: int, dims: int):
+        super().__init__()
+        scale = math.sqrt(1 / dims)
+        self.weight = mx.random.normal((num_embeddings, dims)) * scale
+
+    def _extra_repr(self):
+        return f"{self.weight.shape[0]}, {self.weight.shape[1]}"
+
+    def __call__(self, x):
+        return self.weight[x]

lib/python3.11/site-packages/mlx/nn/layers/linear.py

@@ -0,0 +1,141 @@
+# Copyright © 2023 Apple Inc.
+
+import math
+from typing import Any
+
+import mlx.core as mx
+from mlx.nn.layers.base import Module
+
+
+class Identity(Module):
+    r"""A placeholder identity operator that is argument-insensitive.
+
+    Args:
+        args: any argument (unused)
+        kwargs: any keyword argument (unused)
+    """
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__()
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return x
+
+
+class Linear(Module):
+    r"""Applies an affine transformation to the input.
+
+    Concretely:
+
+    .. math::
+
+        y = x W^\top + b
+
+    where:
+    where :math:`W` has shape ``[output_dims, input_dims]`` and :math:`b` has shape ``[output_dims]``.
+
+    The values are initialized from the uniform distribution :math:`\mathcal{U}(-{k}, {k})`,
+    where :math:`k = \frac{1}{\sqrt{D_i}}` and :math:`D_i` is equal to ``input_dims``.
+
+    Args:
+        input_dims (int): The dimensionality of the input features
+        output_dims (int): The dimensionality of the output features
+        bias (bool, optional): If set to ``False`` then the layer will
+          not use a bias. Default is ``True``.
+    """
+
+    def __init__(self, input_dims: int, output_dims: int, bias: bool = True) -> None:
+        super().__init__()
+        scale = math.sqrt(1.0 / input_dims)
+        self.weight = mx.random.uniform(
+            low=-scale,
+            high=scale,
+            shape=(output_dims, input_dims),
+        )
+        if bias:
+            self.bias = mx.random.uniform(
+                low=-scale,
+                high=scale,
+                shape=(output_dims,),
+            )
+
+    def _extra_repr(self) -> str:
+        return f"input_dims={self.weight.shape[1]}, output_dims={self.weight.shape[0]}, bias={'bias' in self}"
+
+    def __call__(self, x: mx.array) -> mx.array:
+        x = x @ self.weight.T
+        if "bias" in self:
+            x = x + self.bias
+        return x
+
+
+class Bilinear(Module):
+    r"""Applies a bilinear transformation to the inputs.
+
+    Concretely:
+
+    .. math::
+
+        y_i = x_1^\top W_i x_2 + b_i
+
+    where:
+    :math:`W` has shape ``[output_dims, input1_dims, input2_dims]``, :math:`b` has shape ``[output_dims ]``,
+    and :math:`i` indexes the output dimension.
+
+    The values are initialized from the uniform distribution :math:`\mathcal{U}(-{k}, {k})`,
+    where :math:`k = \frac{1}{\sqrt{D_1}}` and :math:`D_1` is ``input1_dims``.
+
+    Args:
+        input1_dims (int): The dimensionality of the input1 features
+        input2_dims (int): The dimensionality of the input2 features
+        output_dims (int): The dimensionality of the output features
+        bias (bool, optional): If set to ``False`` then the layer will
+          not use a bias. Default is ``True``.
+    """
+
+    def __init__(
+        self, input1_dims: int, input2_dims: int, output_dims: int, bias: bool = True
+    ) -> None:
+        super().__init__()
+        scale = math.sqrt(1.0 / input1_dims)
+        self.weight = mx.random.uniform(
+            low=-scale,
+            high=scale,
+            shape=(output_dims, input2_dims, input1_dims),
+        )
+        if bias:
+            self.bias = mx.random.uniform(
+                low=-scale,
+                high=scale,
+                shape=(output_dims,),
+            )
+
+    def _extra_repr(self) -> str:
+        out, in2, in1 = self.weight.shape
+        return (
+            f"input1_dims={in1}, input2_dims={in2}, output_dims={out}, "
+            f"bias={'bias' in self}"
+        )
+
+    def __call__(self, x1: mx.array, x2: mx.array) -> mx.array:
+        # Normalize shapes
+        out, in2, in1 = self.weight.shape
+        xshape = x1.shape[:-1]
+        x1 = x1.reshape(-1, in1)
+        x2 = x2.reshape(-1, 1, in2)
+
+        # Perform the bilinear transformation
+        w = self.weight.reshape(out * in2, in1)
+        y = x1 @ w.T
+        y = y.reshape(-1, out, in2).swapaxes(-2, -1)
+        y = x2 @ y
+        y = y.squeeze(1)
+
+        # Reset the shape
+        y = y.reshape(*xshape, out)
+
+        # Apply the bias
+        if "bias" in self:
+            y = y + self.bias
+
+        return y

lib/python3.11/site-packages/mlx/nn/layers/normalization.py

@@ -0,0 +1,368 @@
+# Copyright © 2023 Apple Inc.
+
+from typing import Tuple
+
+import mlx.core as mx
+from mlx.nn.layers.base import Module
+
+
+class InstanceNorm(Module):
+    r"""Applies instance normalization [1] on the inputs.
+
+    Computes
+
+    .. math::
+
+        y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta,
+
+    where :math:`\gamma` and :math:`\beta` are learned per feature dimension
+    parameters initialized at 1 and 0 respectively. Both are of size :attr:`dims`,
+    if :attr:`affine` is ``True``.
+
+    Args:
+        dims (int): The number of features of the input.
+        eps (float): A value added to the denominator for numerical stability. Default: ``1e-5``.
+        affine (bool): Default: ``False``.
+
+    Shape:
+      - Input: :math:`(..., C)` where :math:`C` is equal to :attr:`dims`.
+      - Output: Same shape as the input.
+
+    Examples:
+        >>> import mlx.core as mx
+        >>> import mlx.nn as nn
+        >>> x = mx.random.normal((8, 4, 4, 16))
+        >>> inorm = nn.InstanceNorm(dims=16)
+        >>> output = inorm(x)
+
+    References:
+        [1]: https://arxiv.org/abs/1607.08022
+    """
+
+    def __init__(
+        self,
+        dims: int,
+        eps: float = 1e-5,
+        affine: bool = False,
+    ):
+        super().__init__()
+        if affine:
+            self.weight = mx.ones((dims,))
+            self.bias = mx.zeros((dims,))
+        self.dims = dims
+        self.eps = eps
+
+    def _extra_repr(self):
+        return f"{self.dims}, eps={self.eps}, affine={'weight' in self}"
+
+    def __call__(self, x: mx.array) -> mx.array:
+        reduction_axes = tuple(range(1, x.ndim - 1))
+        # Compute stats
+        mean = mx.mean(x, axis=reduction_axes, keepdims=True)
+        var = mx.var(x, axis=reduction_axes, keepdims=True)
+        # Normalize
+        x = (x - mean) * mx.rsqrt(var + self.eps)
+        # Scale and shift if necessary
+        return (self.weight * x + self.bias) if "weight" in self else x
+
+
+class LayerNorm(Module):
+    r"""Applies layer normalization [1] on the inputs.
+
+    Computes
+
+    .. math::
+
+        y = \frac{x - E[x]}{\sqrt{Var[x]} + \epsilon} \gamma + \beta,
+
+    where :math:`\gamma` and :math:`\beta` are learned per feature dimension
+    parameters initialized at 1 and 0 respectively.
+
+    [1]: https://arxiv.org/abs/1607.06450
+
+    Args:
+        dims (int): The feature dimension of the input to normalize over
+        eps (float): A small additive constant for numerical stability
+        affine (bool): If True learn an affine transform to apply after the
+            normalization
+    """
+
+    def __init__(self, dims: int, eps: float = 1e-5, affine: bool = True):
+        super().__init__()
+        if affine:
+            self.bias = mx.zeros((dims,))
+            self.weight = mx.ones((dims,))
+        self.eps = eps
+        self.dims = dims
+
+    def _extra_repr(self):
+        return f"{self.dims}, eps={self.eps}, affine={'weight' in self}"
+
+    def __call__(self, x):
+        means = mx.mean(x, axis=-1, keepdims=True)
+        var = mx.var(x, axis=-1, keepdims=True)
+        x = (x - means) * mx.rsqrt(var + self.eps)
+        return (self.weight * x + self.bias) if "weight" in self else x
+
+
+class RMSNorm(Module):
+    r"""Applies Root Mean Square normalization [1] to the inputs.
+
+    Computes
+
+    ..  math::
+
+        y = \frac{x}{\sqrt{E[x^2] + \epsilon}} \gamma
+
+    where :math:`\gamma` is a learned per feature dimension parameter initialized at
+    1.
+
+    [1]: https://arxiv.org/abs/1910.07467
+
+    Args:
+        dims (int): The feature dimension of the input to normalize over
+        eps (float): A small additive constant for numerical stability
+    """
+
+    def __init__(self, dims: int, eps: float = 1e-5):
+        super().__init__()
+        self.weight = mx.ones((dims,))
+        self.eps = eps
+
+    def _extra_repr(self):
+        return f"{self.weight.shape[0]}, eps={self.eps}"
+
+    def __call__(self, x):
+        # S is 1/sqrt(N) where N is the size of the features of x and is used
+        # to compute a numerically more stable RMS of x by multiplying with S
+        # first and summing.
+        #
+        # This way we prefer underflow over overflow which is controlled with
+        # the parameter epsilon anyway.
+        S = 1 / x.shape[-1] ** 0.5
+
+        n = (x * S).square().sum(axis=-1, keepdims=True)
+        n = mx.rsqrt(n + self.eps)
+
+        return self.weight * x * n
+
+
+class GroupNorm(Module):
+    r"""Applies Group Normalization [1] to the inputs.
+
+    Computes the same normalization as layer norm, namely
+
+    .. math::
+
+        y = \frac{x - E[x]}{\sqrt{Var[x]} + \epsilon} \gamma + \beta,
+
+    where :math:`\gamma` and :math:`\beta` are learned per feature dimension
+    parameters initialized at 1 and 0 respectively. However, the mean and
+    variance are computed over the spatial dimensions and each group of
+    features. In particular, the input is split into num_groups across the
+    feature dimension.
+
+    The feature dimension is assumed to be the last dimension and the dimensions
+    that precede it (except the first) are considered the spatial dimensions.
+
+    [1]: https://arxiv.org/abs/1803.08494
+
+    Args:
+        num_groups (int): Number of groups to separate the features into
+        dims (int): The feature dimensions of the input to normalize over
+        eps (float): A small additive constant for numerical stability
+        affine (bool): If True learn an affine transform to apply after the
+            normalization.
+        pytorch_compatible (bool): If True perform the group normalization in
+            the same order/grouping as PyTorch.
+    """
+
+    def __init__(
+        self,
+        num_groups: int,
+        dims: int,
+        eps: float = 1e-5,
+        affine: bool = True,
+        pytorch_compatible: bool = False,
+    ):
+        super().__init__()
+        if affine:
+            self.bias = mx.zeros((dims,))
+            self.weight = mx.ones((dims,))
+        self.num_groups = num_groups
+        self.dims = dims
+        self.eps = eps
+        self.pytorch_compatible = pytorch_compatible
+
+    def _extra_repr(self):
+        return (
+            f"{self.num_groups}, {self.dims}, eps={self.eps}, "
+            f"affine={'weight' in self}, pytorch_compatible={self.pytorch_compatible}"
+        )
+
+    def _pytorch_compatible_group_norm(self, x):
+        num_groups = self.num_groups
+        batch, *rest, dims = x.shape
+
+        # Split into groups
+        x = x.reshape(batch, -1, num_groups, dims // num_groups)
+        x = x.transpose(0, 1, 3, 2).reshape(batch, -1, num_groups)
+
+        # Normalize
+        means = mx.mean(x, axis=1, keepdims=True)
+        var = mx.var(x, axis=1, keepdims=True)
+        x = (x - means) * mx.rsqrt(var + self.eps)
+        x = x.reshape(batch, -1, dims // num_groups, num_groups)
+        x = x.transpose(0, 1, 3, 2).reshape(batch, *rest, dims)
+
+        return x
+
+    def _group_norm(self, x):
+        num_groups = self.num_groups
+        batch, *rest, dims = x.shape
+
+        # Split into groups
+        x = x.reshape(batch, -1, num_groups)
+
+        # Normalize
+        means = mx.mean(x, axis=1, keepdims=True)
+        var = mx.var(x, axis=1, keepdims=True)
+        x = (x - means) * mx.rsqrt(var + self.eps)
+        x = x.reshape(batch, *rest, dims)
+
+        return x
+
+    def __call__(self, x):
+        group_norm = (
+            self._pytorch_compatible_group_norm
+            if self.pytorch_compatible
+            else self._group_norm
+        )
+        x = group_norm(x)
+        return (self.weight * x + self.bias) if "weight" in self else x
+
+
+class BatchNorm(Module):
+    r"""Applies Batch Normalization over a 2D or 3D input.
+
+    Computes
+
+    .. math::
+
+        y = \frac{x - E[x]}{\sqrt{Var[x]} + \epsilon} \gamma + \beta,
+
+    where :math:`\gamma` and :math:`\beta` are learned per feature dimension
+    parameters initialized at 1 and 0 respectively.
+
+    The input shape is specified as ``NC`` or ``NLC``, where ``N`` is the
+    batch, ``C`` is the number of features or channels, and ``L`` is the
+    sequence length. The output has the same shape as the input. For
+    four-dimensional arrays, the shape is ``NHWC``, where ``H`` and ``W`` are
+    the height and width respectively.
+
+    For more information on Batch Normalization, see the original paper `Batch
+    Normalization: Accelerating Deep Network Training by Reducing Internal
+    Covariate Shift <https://arxiv.org/abs/1502.03167>`_.
+
+    Args:
+        num_features (int): The feature dimension to normalize over.
+        eps (float, optional): A small additive constant for numerical
+            stability. Default: ``1e-5``.
+        momentum (float, optional): The momentum for updating the running
+            mean and variance. Default: ``0.1``.
+        affine (bool, optional): If ``True``, apply a learned affine
+            transformation after the normalization. Default: ``True``.
+        track_running_stats (bool, optional): If ``True``, track the
+            running mean and variance. Default: ``True``.
+
+    Examples:
+        >>> import mlx.core as mx
+        >>> import mlx.nn as nn
+        >>> x = mx.random.normal((5, 4))
+        >>> bn = nn.BatchNorm(num_features=4, affine=True)
+        >>> output = bn(x)
+    """
+
+    def __init__(
+        self,
+        num_features: int,
+        eps: float = 1e-5,
+        momentum: float = 0.1,
+        affine: bool = True,
+        track_running_stats: bool = True,
+    ):
+        super().__init__()
+
+        self.num_features = num_features
+        self.eps = eps
+        self.momentum = momentum
+        self.track_running_stats = track_running_stats
+
+        if affine:
+            self.weight = mx.ones((num_features,))
+            self.bias = mx.zeros((num_features,))
+
+        if self.track_running_stats:
+            self.running_mean = mx.zeros((num_features,))
+            self.running_var = mx.ones((num_features,))
+            self.freeze(keys=["running_mean", "running_var"], recurse=False)
+
+    def unfreeze(self, *args, **kwargs):
+        """Wrap unfreeze to make sure that running_mean and var are always
+        frozen parameters."""
+        super().unfreeze(*args, **kwargs)
+        self.freeze(keys=["running_mean", "running_var"], recurse=False)
+
+    def _extra_repr(self):
+        return (
+            f"{self.num_features}, eps={self.eps}, "
+            f"momentum={self.momentum}, affine={'weight' in self}, "
+            f"track_running_stats={self.track_running_stats}"
+        )
+
+    def _calc_stats(self, x: mx.array) -> Tuple[mx.array, mx.array]:
+        """
+        Calculate the mean and variance of the input tensor across the batch
+        and spatial dimensions.
+
+        Args:
+            x (array): Input tensor.
+
+        Returns:
+            tuple: Tuple containing mean and variance.
+        """
+        reduction_axes = tuple(range(0, x.ndim - 1))
+
+        mean = mx.mean(x, axis=reduction_axes, keepdims=True)
+        var = mx.var(x, axis=reduction_axes, keepdims=True)
+
+        return mean, var
+
+    def __call__(self, x: mx.array) -> mx.array:
+        """
+        Forward pass of BatchNorm.
+
+        Args:
+            x (array): Input tensor.
+
+        Returns:
+            array: Normalized output tensor.
+        """
+        if x.ndim < 2 or x.ndim > 4:
+            raise ValueError(
+                f"Expected input tensor to have 2, 3 or 4 dimensions, but got {x.ndim}"
+            )
+
+        # Calculate the mean and variance used to normalize the input x. If we
+        # are in training mode update the running stats if needed.
+        mean, var = self._calc_stats(x)
+        if self.training and self.track_running_stats:
+            mu = self.momentum
+            self.running_mean = (1 - mu) * self.running_mean + mu * mean
+            self.running_var = (1 - mu) * self.running_var + mu * var
+        elif self.track_running_stats:
+            mean = self.running_mean
+            var = self.running_var
+
+        x = (x - mean) * mx.rsqrt(var + self.eps)
+        return (self.weight * x + self.bias) if "weight" in self else x

lib/python3.11/site-packages/mlx/nn/layers/positional_encoding.py

@@ -0,0 +1,199 @@
+# Copyright © 2023 Apple Inc.
+
+import math
+from typing import Optional
+
+import mlx.core as mx
+from mlx.nn.layers.base import Module
+
+
+class RoPE(Module):
+    """Implements the rotary positional encoding.
+
+    The traditional implementation rotates consecutive pairs of elements in the
+    feature dimension while the default implementation rotates pairs with
+    stride half the feature dimensions for efficiency.
+
+    For more details see `RoFormer: Enhanced Transformer with Rotary Position
+    Embedding <https://arxiv.org/abs/2104.09864>`_.
+
+    Args:
+        dims (int): The feature dimensions to be rotated. If the input feature
+            is larger than dims then the rest is left unchanged.
+        traditional (bool, optional): If set to True choose the traditional
+            implementation which is slightly less efficient. Default: ``False``.
+        base (float, optional): The base used to compute angular frequency for
+            each dimension in the positional encodings. Default: ``10000``.
+        scale (float, optional): The scale used to scale the positions. Default: ``1.0``.
+    """
+
+    def __init__(
+        self,
+        dims: int,
+        traditional: bool = False,
+        base: float = 10000,
+        scale: float = 1.0,
+    ):
+        super().__init__()
+        self.dims = dims
+        self.traditional = traditional
+        self.base = base
+        self.scale = scale
+
+    def _extra_repr(self):
+        return f"{self.dims}, traditional={self.traditional}"
+
+    def _compute_rope(self, costheta, sintheta, x):
+        x1 = x[..., : self.dims // 2]
+        x2 = x[..., self.dims // 2 : self.dims]
+        rx1 = x1 * costheta - x2 * sintheta
+        rx2 = x1 * sintheta + x2 * costheta
+
+        if self.dims < x.shape[-1]:
+            rx = mx.concatenate([rx1, rx2, x[..., self.dims :]], axis=-1)
+        else:
+            rx = mx.concatenate([rx1, rx2], axis=-1)
+
+        return rx
+
+    def _compute_traditional_rope(self, costheta, sintheta, x):
+        x1 = x[..., ::2]
+        x2 = x[..., 1::2]
+        rx1 = x1 * costheta - x2 * sintheta
+        rx2 = x1 * sintheta + x2 * costheta
+
+        if self.dims < x.shape[-1]:
+            raise NotImplementedError(
+                "RoPE doesn't implement partial traditional application"
+            )
+
+        rx = mx.concatenate([rx1[..., None], rx2[..., None]], axis=-1)
+
+        return rx
+
+    def __call__(self, x, offset: int = 0):
+        shape = x.shape
+        x = mx.reshape(x, (-1, shape[-2], shape[-1]))
+        N = x.shape[1] + offset
+        costheta, sintheta = RoPE.create_cos_sin_theta(
+            N, self.dims, offset=offset, base=self.base, scale=self.scale, dtype=x.dtype
+        )
+
+        rope = (
+            self._compute_traditional_rope if self.traditional else self._compute_rope
+        )
+        rx = rope(costheta, sintheta, x)
+
+        return mx.reshape(rx, shape)
+
+    @staticmethod
+    def create_cos_sin_theta(
+        N: int,
+        D: int,
+        offset: int = 0,
+        base: float = 10000,
+        scale: float = 1.0,
+        dtype=mx.float32,
+    ):
+        D = D // 2
+        positions = mx.arange(offset, N, dtype=dtype) * scale
+        freqs = mx.exp(-mx.arange(0.0, D, dtype=dtype) * (math.log(base) / D))
+        theta = mx.reshape(positions, (-1, 1)) * mx.reshape(freqs, (1, -1))
+        return mx.cos(theta), mx.sin(theta)
+
+
+class SinusoidalPositionalEncoding(Module):
+    r"""Implements sinusoidal positional encoding.
+
+    For more details see the paper `Attention Is All You Need
+    <https://arxiv.org/abs/1706.03762>`_.
+
+    Args:
+        dims (int): The dimensionality of the resulting positional embeddings.
+        min_freq (float, optional): The minimum frequency expected. Default:
+            ``0.0001``.
+        max_freq (float, optional): The maximum frequency expected. Default:
+            ``1``.
+        scale (float, optional): A multiplicative scale for the embeddings.
+            Default: ``sqrt(dims//2)``.
+        cos_first (bool, optional): If ``True`` embed using ``[cos(x); sin(x)]``
+            instead of the reverse. Default: ``False``.
+        full_turns (bool, optional): If ``True`` multiply the frequencies with
+            :math:`2\pi`. Default: ``False``.
+    """
+
+    def __init__(
+        self,
+        dims: int,
+        min_freq: float = 0.0001,
+        max_freq: float = 1,
+        scale: Optional[float] = None,
+        cos_first: bool = False,
+        full_turns: bool = False,
+    ):
+        super().__init__()
+
+        one_zero = 1 - mx.arange(0, dims // 2) / (dims // 2 - 1)
+        min_freq = math.log(min_freq)
+        max_freq = math.log(max_freq)
+
+        # Start with underscore so it is not included in the parameters
+        self._sigmas = mx.exp(one_zero * (max_freq - min_freq) + min_freq)
+        if full_turns:
+            self._sigmas = self._sigmas * (2 * math.pi)
+
+        # Save some constants that define the implementation
+        self.scale = scale or (2 / dims) ** 0.5
+        self.cos_first = cos_first
+
+    def __call__(self, x):
+        y = x[..., None] * self._sigmas
+        cosy = mx.cos(y)
+        siny = mx.sin(y)
+
+        if self.cos_first:
+            y = mx.concatenate([cosy, siny], axis=-1)
+        else:
+            y = mx.concatenate([siny, cosy], axis=-1)
+
+        if self.scale != 1:
+            y = y * self.scale
+
+        return y
+
+
+class ALiBi(Module):
+    @staticmethod
+    def create_alibi_matrix(
+        q_sequence_length: int,
+        k_sequence_length: int,
+        num_heads: int,
+        offset: int,
+        dtype=mx.float32,
+    ):
+        x1 = mx.arange(offset, q_sequence_length)
+        x2 = mx.arange(0, k_sequence_length)
+        distance_matrix = -mx.abs(
+            mx.expand_dims(x1[:, None] - x2[None, :], axis=(0, 1))
+        )
+        alibi_slope = ALiBi.create_alibi_slope(num_heads=num_heads)
+        alibi_mask = (distance_matrix * alibi_slope).astype(dtype)
+        return alibi_mask
+
+    @staticmethod
+    def create_alibi_slope(num_heads):
+        x = (2**8) ** (1 / num_heads)
+        out = mx.power(x, -mx.arange(1, num_heads + 1))
+        return mx.expand_dims(out, axis=(-1, -2))
+
+    def __call__(self, attention_scores, offset=0, mask=None):
+        alibi_mask = ALiBi.create_alibi_matrix(
+            q_sequence_length=attention_scores.shape[-2] + offset,
+            k_sequence_length=attention_scores.shape[-1],
+            num_heads=attention_scores.shape[1],
+            offset=offset,
+            dtype=attention_scores.dtype,
+        )
+        if mask is not None:
+            alibi_mask = alibi_mask + mask
+        return attention_scores + alibi_mask

lib/python3.11/site-packages/mlx/nn/layers/quantized.py

@@ -0,0 +1,125 @@
+# Copyright © 2023 Apple Inc.
+
+import math
+
+import mlx.core as mx
+from mlx.nn.layers.base import Module
+from mlx.nn.layers.linear import Linear
+from mlx.utils import tree_flatten, tree_map
+
+
+class QuantizedLinear(Module):
+    """Applies an affine transformation to the input using a quantized weight matrix.
+
+    It is the quantized equivalent of :class:`mlx.nn.Linear`. For now its
+    parameters are frozen and will not be included in any gradient computation
+    but this will probably change in the future.
+
+    QuantizedLinear also provides two useful classmethods to convert linear
+    layers to QuantizedLinear layers.
+
+    - :meth:`from_linear` returns a QuantizedLinear layer that applies the same
+      linear transformation up to the quantization error.
+    - :meth:`quantize_module` swaps all the linear layers of the passed module
+      with QuantizedLinear ones.
+
+    Args:
+        input_dims (int): The dimensionality of the input features
+        output_dims (int): The dimensionality of the output features
+        bias (bool, optional): If set to ``False`` then the layer will not use
+            a bias. (default: True).
+        group_size (int, optional): The group size to use for the quantized
+            weight. See :func:`~mlx.core.quantize`. (default: 64)
+        bits (int, optional): The bit width to use for the quantized weight.
+            See :func:`~mlx.core.quantize`. (default: 4)
+    """
+
+    def __init__(
+        self,
+        input_dims: int,
+        output_dims: int,
+        bias: bool = True,
+        group_size: int = 64,
+        bits: int = 4,
+    ):
+        super().__init__()
+
+        # Quantization config
+        self.group_size = group_size
+        self.bits = bits
+
+        # Initialize the quantized weight
+        scale = math.sqrt(1 / input_dims)
+        weight = mx.random.uniform(
+            low=-scale,
+            high=scale,
+            shape=(output_dims, input_dims),
+        )
+        self.weight, self.scales, self.biases = mx.quantize(weight, group_size, bits)
+
+        # And bias if needed
+        if bias:
+            self.bias = mx.zeros((output_dims,))
+
+        # Freeze this model's parameters
+        self.freeze()
+
+    def unfreeze(self, *args, **kwargs):
+        """Wrap unfreeze so that we unfreeze any layers we might contain but
+        our parameters will remain frozen."""
+        super().unfreeze(*args, **kwargs)
+        self.freeze(recurse=False)
+
+    def _extra_repr(self):
+        out_dims, in_dims = self.weight.shape
+        in_dims *= 32 // self.bits
+        return (
+            f"input_dims={in_dims}, output_dims={out_dims}, bias={'bias' in self},"
+            f"group_size={self.group_size}, bits={self.bits}"
+        )
+
+    def __call__(self, x):
+        x = mx.quantized_matmul(
+            x,
+            self.weight,
+            scales=self.scales,
+            biases=self.biases,
+            transpose=True,
+            group_size=self.group_size,
+            bits=self.bits,
+        )
+        if "bias" in self:
+            x = x + self.bias
+        return x
+
+    @classmethod
+    def from_linear(cls, linear_layer: Module, group_size: int = 64, bits: int = 4):
+        """Create a QuantizedLinear layer from the parameters of a provided
+        linear layer."""
+        output_dims, input_dims = linear_layer.weight.shape
+        ql = cls(input_dims, output_dims, False, group_size, bits)
+        ql.weight, ql.scales, ql.biases = mx.quantize(
+            linear_layer.weight, group_size, bits
+        )
+        if "bias" in linear_layer:
+            ql.bias = linear_layer.bias
+
+        return ql
+
+    @classmethod
+    def quantize_module(
+        cls,
+        model: Module,
+        group_size: int = 64,
+        bits: int = 4,
+        linear_class_predicate=lambda m: isinstance(m, Linear),
+    ):
+        def _quantize_if_linear(m):
+            if linear_class_predicate(m):
+                return cls.from_linear(m, group_size, bits)
+            else:
+                return m
+
+        leaves = model.leaf_modules()
+        leaves = tree_map(_quantize_if_linear, leaves, is_leaf=Module.is_module)
+        model.update_modules(leaves)

lib/python3.11/site-packages/mlx/nn/layers/transformer.py

@@ -0,0 +1,354 @@
+# Copyright © 2023 Apple Inc.
+
+import math
+from typing import Any, Callable, Optional
+
+import mlx.core as mx
+from mlx.nn.layers.activations import relu
+from mlx.nn.layers.base import Module
+from mlx.nn.layers.dropout import Dropout
+from mlx.nn.layers.linear import Linear
+from mlx.nn.layers.normalization import LayerNorm
+
+
+class MultiHeadAttention(Module):
+    """Implements the scaled dot product attention with multiple heads.
+
+    Given inputs for queries, keys and values the ``MultiHeadAttention``
+    produces new values by aggregating information from the input values
+    according to the similarities of the input queries and keys.
+
+    All inputs as well as the output are linearly projected without biases by
+    default.
+
+    ``MultiHeadAttention`` also takes an optional additive attention mask that
+    should be broadcastable with ``(batch, num_heads, # queries, # keys)``. The
+    mask should have ``-inf`` or very large negative numbers at the positions
+    that should *not* be attended to.
+
+    Args:
+        dims (int): The model dimensions. This is also the default
+            value for the queries, keys, values, and the output.
+        num_heads (int): The number of attention heads to use.
+        query_input_dims (int, optional): The input dimensions of the queries.
+            Default: ``dims``.
+        key_input_dims (int, optional): The input dimensions of the keys.
+            Default: ``dims``.
+        value_input_dims (int, optional): The input dimensions of the values.
+            Default: ``key_input_dims``.
+        value_dims (int, optional): The dimensions of the values after the
+            projection. Default: ``dims``.
+        value_output_dims (int, optional): The dimensions the new values will
+            be projected to. Default: ``dims``.
+        bias (bool, optional): Whether or not to use a bias in the projections.
+            Default: ``False``.
+    """
+
+    def __init__(
+        self,
+        dims: int,
+        num_heads: int,
+        query_input_dims: Optional[int] = None,
+        key_input_dims: Optional[int] = None,
+        value_input_dims: Optional[int] = None,
+        value_dims: Optional[int] = None,
+        value_output_dims: Optional[int] = None,
+        bias: bool = False,
+    ):
+        super().__init__()
+
+        if (dims % num_heads) != 0:
+            raise ValueError(
+                "The input feature dimensions should be divisible by the "
+                f"number of heads ({dims} % {num_heads}) != 0"
+            )
+
+        query_input_dims = query_input_dims or dims
+        key_input_dims = key_input_dims or dims
+        value_input_dims = value_input_dims or key_input_dims
+        value_dims = value_dims or dims
+        value_output_dims = value_output_dims or dims
+
+        self.num_heads = num_heads
+        self.query_proj = Linear(query_input_dims, dims, bias=bias)
+        self.key_proj = Linear(key_input_dims, dims, bias=bias)
+        self.value_proj = Linear(value_input_dims, value_dims, bias=bias)
+        self.out_proj = Linear(value_dims, value_output_dims, bias=bias)
+
+    def __call__(self, queries, keys, values, mask=None):
+        queries = self.query_proj(queries)
+        keys = self.key_proj(keys)
+        values = self.value_proj(values)
+
+        num_heads = self.num_heads
+        B, L, D = queries.shape
+        _, S, _ = keys.shape
+        queries = queries.reshape(B, L, num_heads, -1).transpose(0, 2, 1, 3)
+        keys = keys.reshape(B, S, num_heads, -1).transpose(0, 2, 3, 1)
+        values = values.reshape(B, S, num_heads, -1).transpose(0, 2, 1, 3)
+
+        # Dimensions are [batch x num heads x sequence x hidden dim]
+        scale = math.sqrt(1 / queries.shape[-1])
+        scores = (queries * scale) @ keys
+        if mask is not None:
+            scores = scores + mask.astype(scores.dtype)
+        scores = mx.softmax(scores, axis=-1)
+        values_hat = (scores @ values).transpose(0, 2, 1, 3).reshape(B, L, -1)
+
+        return self.out_proj(values_hat)
+
+    @staticmethod
+    def create_additive_causal_mask(N: int, dtype: mx.Dtype = mx.float32):
+        indices = mx.arange(N)
+        mask = indices[:, None] < indices[None]
+        # usually inf but 1e9 is as good and softmax(full(1e9)) != nan
+        # TODO: Should replace this with finfo(dtype).min
+        mask = mask.astype(dtype) * -1e9
+        return mask
+
+
+class TransformerEncoderLayer(Module):
+    def __init__(
+        self,
+        dims: int,
+        num_heads: int,
+        mlp_dims: Optional[int] = None,
+        dropout: float = 0.0,
+        activation: Callable[[Any], Any] = relu,
+        norm_first: bool = False,
+    ):
+        super().__init__()
+        mlp_dims = mlp_dims or dims * 4
+        self.attention = MultiHeadAttention(dims, num_heads)
+        self.ln1 = LayerNorm(dims)
+        self.ln2 = LayerNorm(dims)
+        self.linear1 = Linear(dims, mlp_dims)
+        self.linear2 = Linear(mlp_dims, dims)
+        self.dropout1 = Dropout(dropout)
+        self.dropout2 = Dropout(dropout)
+        self.activation = activation
+        self.norm_first = norm_first
+
+    def __call__(self, x, mask):
+        if self.norm_first:
+            y = self.ln1(x)
+            y = self.attention(y, y, y, mask)
+            y = self.dropout1(y)
+            x = x + y
+
+            y = self.ln2(x)
+            y = self.linear1(y)
+            y = self.activation(y)
+            y = self.dropout2(y)
+            y = self.linear2(y)
+            y = x + y
+
+        else:
+            y = self.attention(x, x, x, mask)
+            y = self.dropout1(y)
+            y = self.ln1(x + y)
+
+            y = self.linear1(y)
+            y = self.activation(y)
+            y = self.dropout2(y)
+            y = self.linear2(y)
+            y = self.ln2(x + y)
+
+        return y
+
+
+class TransformerEncoder(Module):
+    def __init__(
+        self,
+        num_layers: int,
+        dims: int,
+        num_heads: int,
+        mlp_dims: Optional[int] = None,
+        dropout: float = 0.0,
+        activation=relu,
+        norm_first: bool = False,
+    ):
+        super().__init__()
+        self.layers = [
+            TransformerEncoderLayer(
+                dims, num_heads, mlp_dims, dropout, activation, norm_first
+            )
+            for i in range(num_layers)
+        ]
+        self.ln = LayerNorm(dims)
+
+    def __call__(self, x, mask):
+        for l in self.layers:
+            x = l(x, mask)
+        return self.ln(x)
+
+
+class TransformerDecoderLayer(Module):
+    def __init__(
+        self,
+        dims: int,
+        num_heads: int,
+        mlp_dims: Optional[int] = None,
+        dropout: float = 0.0,
+        activation: Callable[[Any], Any] = relu,
+        norm_first: bool = False,
+    ):
+        super().__init__()
+        mlp_dims = mlp_dims or dims * 4
+        self.self_attention = MultiHeadAttention(dims, num_heads)
+        self.cross_attention = MultiHeadAttention(dims, num_heads)
+        self.ln1 = LayerNorm(dims)
+        self.ln2 = LayerNorm(dims)
+        self.ln3 = LayerNorm(dims)
+        self.linear1 = Linear(dims, mlp_dims)
+        self.linear2 = Linear(mlp_dims, dims)
+        self.dropout1 = Dropout(dropout)
+        self.dropout2 = Dropout(dropout)
+        self.dropout3 = Dropout(dropout)
+        self.activation = activation
+        self.norm_first = norm_first
+
+    def __call__(self, x, memory, x_mask, memory_mask):
+        if self.norm_first:
+            y = self.ln1(x)
+            y = self.self_attention(y, y, y, x_mask)
+            y = self.dropout1(y)
+            x = x + y
+
+            y = self.ln2(x)
+            y = self.cross_attention(y, memory, memory, memory_mask)
+            y = self.dropout2(y)
+            x = x + y
+
+            y = self.ln3(x)
+            y = self.linear1(y)
+            y = self.activation(y)
+            y = self.dropout3(y)
+            y = self.linear2(y)
+            y = x + y
+
+        else:
+            y = self.self_attention(x, x, x, x_mask)
+            y = self.dropout1(y)
+            x = self.ln1(x + y)
+
+            y = self.cross_attention(y, memory, memory, memory_mask)
+            y = self.dropout2(y)
+            x = self.ln1(x + y)
+
+            y = self.linear1(x)
+            y = self.activation(y)
+            y = self.dropout3(y)
+            y = self.linear2(y)
+            y = self.ln3(x + y)
+
+        return y
+
+
+class TransformerDecoder(Module):
+    def __init__(
+        self,
+        num_layers: int,
+        dims: int,
+        num_heads: int,
+        mlp_dims: Optional[int] = None,
+        dropout: float = 0.0,
+        activation=relu,
+        norm_first: bool = False,
+    ):
+        super().__init__()
+        self.layers = [
+            TransformerDecoderLayer(
+                dims, num_heads, mlp_dims, dropout, activation, norm_first
+            )
+            for i in range(num_layers)
+        ]
+        self.ln = LayerNorm(dims)
+
+    def __call__(self, x, memory, x_mask, memory_mask):
+        for l in self.layers:
+            x = l(x, memory, x_mask, memory_mask)
+        return self.ln(x)
+
+
+class Transformer(Module):
+    """
+    Implements a standard Transformer model.
+
+    The implementation is based on `Attention Is All You Need
+    <https://arxiv.org/abs/1706.03762>`_.
+
+    The Transformer model contains an encoder and a decoder. The encoder
+    processes the input sequence and the decoder generates the output sequence.
+    The interaction between encoder and decoder happens through the attention
+    mechanism.
+
+    Args:
+        dims (int, optional): The number of expected features in the
+            encoder/decoder inputs. Default: ``512``.
+        num_heads (int, optional): The number of attention heads. Default:
+            ``8``.
+        num_encoder_layers (int, optional): The number of encoder layers in the
+            Transformer encoder. Default: ``6``.
+        num_decoder_layers (int, optional): The number of decoder layers in the
+            Transformer decoder. Default: ``6``.
+        mlp_dims (int, optional): The hidden dimension of the MLP block in each
+            Transformer layer. Defaults to ``4*dims`` if not provided. Default:
+            ``None``.
+        dropout (float, optional): The dropout value for the Transformer
+            encoder and decoder. Dropout is used after each attention layer and
+            the activation in the MLP layer. Default: ``0.0``.
+        activation (function, optional): the activation function for the MLP
+            hidden layer. Default: :func:`mlx.nn.relu`.
+        custom_encoder (nn.Module, optional): A custom encoder to replace the
+            standard Transformer encoder. Default: ``None``.
+        custom_decoder (nn.Module, optional): A custom decoder to replace the
+            standard Transformer decoder. Default: ``None``.
+        norm_first (bool, optional): if ``True``, encoder and decoder layers
+            will perform layer normalization before attention and MLP
+            operations, otherwise after. Default: ``False``.
+    """
+
+    def __init__(
+        self,
+        dims: int = 512,
+        num_heads: int = 8,
+        num_encoder_layers: int = 6,
+        num_decoder_layers: int = 6,
+        mlp_dims: Optional[int] = None,
+        dropout: float = 0.0,
+        activation: Callable[[Any], Any] = relu,
+        custom_encoder: Optional[Any] = None,
+        custom_decoder: Optional[Any] = None,
+        norm_first: bool = False,
+    ):
+        super().__init__()
+        if custom_encoder is not None:
+            self.encoder = custom_encoder
+        else:
+            self.encoder = TransformerEncoder(
+                num_encoder_layers,
+                dims,
+                num_heads,
+                mlp_dims,
+                dropout,
+                activation,
+                norm_first,
+            )
+
+        if custom_decoder is not None:
+            self.decoder = custom_decoder
+        else:
+            self.decoder = TransformerDecoder(
+                num_decoder_layers,
+                dims,
+                num_heads,
+                mlp_dims,
+                dropout,
+                activation,
+                norm_first,
+            )
+
+    def __call__(self, src, tgt, src_mask, tgt_mask, memory_mask):
+        memory = self.encoder(src, src_mask)
+        return self.decoder(tgt, memory, tgt_mask, memory_mask)

lib/python3.11/site-packages/mlx/nn/losses.py

@@ -0,0 +1,374 @@
+# Copyright © 2023 Apple Inc.
+
+import math
+
+import mlx.core as mx
+from mlx.nn.layers.base import Module
+
+
+def cross_entropy(
+    logits: mx.array,
+    targets: mx.array,
+    weights: mx.array = None,
+    axis: int = -1,
+    label_smoothing: float = 0.0,
+    reduction: str = "none",
+) -> mx.array:
+    """
+    Computes the cross entropy loss.
+
+    Args:
+        logits (array): The unnormalized predicted logits.
+        targets (array): The target values, as class indices.
+        weights (array, optional): Weights for each target. Default: ``None``.
+        axis (int, optional): The axis over which to compute softmax. Default: ``-1``.
+        label_smoothing (float, optional): Label smoothing factor. Default: ``0``.
+        reduction (str, optional): Specifies the reduction to apply to the output:
+            ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``.
+
+    Returns:
+        array: The computed cross entropy loss.
+    """
+    if label_smoothing < 0 or label_smoothing >= 1:
+        raise ValueError(f"Label smoothing must in [0, 1), got {label_smoothing}.")
+
+    score = mx.take_along_axis(logits, targets[..., None], axis).squeeze(-1)
+    logsumexp_logits = mx.logsumexp(logits, axis=axis)
+    if label_smoothing > 0:
+        # Adjust the true class score with label smoothing
+        adjusted_score = (1 - label_smoothing) * score
+
+        # Calculate the mean logit across the classes for smoothed loss
+        mean_logits = logits.mean(axis=axis)
+        smoothed_loss = -mean_logits * label_smoothing
+
+        # Combine the adjusted score and smoothed loss with the logsumexp logits
+        loss = logsumexp_logits - adjusted_score + smoothed_loss
+    else:
+        loss = logsumexp_logits - score
+
+    # Apply weights if provided
+    if weights is not None:
+        if weights.shape != targets.shape:
+            raise ValueError(
+                f"Weights with shape {weights.shape} is not the same as "
+                f"targets with shape {targets.shape}."
+            )
+        loss *= weights
+
+    # Apply reduction
+    return _reduce(loss, reduction)
+
+
+def binary_cross_entropy(
+    logits: mx.array, targets: mx.array, reduction: str = "none"
+) -> mx.array:
+    """
+    Computes the binary cross entropy loss.
+
+    Args:
+        logits (array): The unnormalized (pre-sigmoid) predicted logits.
+        targets (array): The binary target values in {0, 1}.
+        reduction (str, optional): Specifies the reduction to apply to the output:
+          ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``.
+
+    Returns:
+        array: The computed binary cross entropy loss.
+    Examples:
+        >>> import mlx.core as mx
+        >>> import mlx.nn as nn
+        >>> inputs = mx.array([0.105361, 0.223144, 1.20397, 0.916291])
+        >>> targets = mx.array([0, 0, 1, 1])
+        >>> loss = nn.losses.binary_cross_entropy(inputs, targets, "mean")
+        >>> loss
+        array([0.612192], dtype=float32)
+    """
+    loss = mx.logaddexp(0.0, logits) - targets * logits
+    return _reduce(loss, reduction)
+
+
+def l1_loss(
+    predictions: mx.array, targets: mx.array, reduction: str = "mean"
+) -> mx.array:
+    """
+    Computes the L1 loss.
+
+    Args:
+        predictions (array): The predicted values.
+        targets (array): The target values.
+        reduction (str, optional): Specifies the reduction to apply to the output:
+          ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'mean'``.
+
+    Returns:
+        array: The computed L1 loss.
+    """
+    if predictions.shape != targets.shape:
+        raise ValueError(
+            f"Predictions shape {predictions.shape} does not match "
+            f"targets shape {targets.shape}."
+        )
+    loss = mx.abs(predictions - targets)
+
+    return _reduce(loss, reduction)
+
+
+def mse_loss(
+    predictions: mx.array, targets: mx.array, reduction: str = "mean"
+) -> mx.array:
+    """
+    Computes the mean squared error loss.
+
+    Args:
+        predictions (array): The predicted values.
+        targets (array): The target values.
+        reduction (str, optional): Specifies the reduction to apply to the output:
+          ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'mean'``.
+
+    Returns:
+        array: The computed mean squared error loss.
+    """
+    if predictions.shape != targets.shape:
+        raise ValueError(
+            f"Predictions shape {predictions.shape} does not match "
+            f"targets shape {targets.shape}."
+        )
+
+    loss = mx.square(predictions - targets)
+    return _reduce(loss, reduction)
+
+
+def nll_loss(
+    inputs: mx.array, targets: mx.array, axis: int = -1, reduction: str = "none"
+) -> mx.array:
+    """
+    Computes the negative log likelihood loss.
+
+    Args:
+        inputs (array): The predicted distribution in log space.
+        targets (array): The target values.
+        axis (int, optional): The distribution axis. Default: ``-1``.
+        reduction (str, optional): Specifies the reduction to apply to the output:
+          ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``.
+
+    Returns:
+        array: The computed NLL loss.
+    """
+    loss = -mx.take_along_axis(inputs, targets[..., None], axis).squeeze(-1)
+
+    return _reduce(loss, reduction)
+
+
+def kl_div_loss(
+    inputs: mx.array, targets: mx.array, axis: int = -1, reduction: str = "none"
+) -> mx.array:
+    """
+    Computes the Kullback-Leibler divergence loss.
+
+    Computes the following when ``reduction == 'none'``:
+
+    .. code-block:: python
+
+        mx.exp(targets) * (targets - inputs).sum(axis)
+
+    Args:
+        inputs (array): Log probabilities for the predicted distribution.
+        targets (array): Log probabilities for the target distribution.
+        axis (int, optional): The distribution axis. Default: ``-1``.
+        reduction (str, optional): Specifies the reduction to apply to the output:
+          ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``.
+
+    Returns:
+        array: The computed Kullback-Leibler divergence loss.
+    """
+    loss = mx.sum(mx.exp(targets) * (targets - inputs), axis)
+
+    return _reduce(loss, reduction)
+
+
+def smooth_l1_loss(
+    predictions: mx.array, targets: mx.array, beta: float = 1.0, reduction: str = "mean"
+) -> mx.array:
+    r"""
+    Computes the smooth L1 loss.
+
+    The smooth L1 loss is a variant of the L1 loss which replaces the absolute
+    difference with a squared difference when the absolute difference is less
+    than ``beta``.
+
+    The formula for the smooth L1 Loss is:
+
+    .. math::
+
+       l =
+          \begin{cases}
+            0.5 (x - y)^2, & \text{ if } & (x - y) < \beta \\
+            |x - y| - 0.5 \beta, &  & \text{otherwise}
+          \end{cases}
+
+    Args:
+        predictions (array): Predicted values.
+        targets (array): Ground truth values.
+        beta (float, optional): The threshold after which the loss changes
+          from the squared to the absolute difference. Default: ``1.0``.
+        reduction (str, optional): Specifies the reduction to apply to the output:
+          ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'mean'``.
+
+    Returns:
+        array: The computed smooth L1 loss.
+    """
+    if predictions.shape != targets.shape:
+        raise ValueError(
+            f"Predictions shape {predictions.shape} does not match "
+            f"targets shape {targets.shape}."
+        )
+
+    diff = predictions - targets
+    loss = mx.where(
+        diff < beta, 0.5 * mx.square(diff) / beta, mx.abs(diff) - 0.5 * beta
+    )
+
+    return _reduce(loss, reduction)
+
+
+def triplet_loss(
+    anchors: mx.array,
+    positives: mx.array,
+    negatives: mx.array,
+    axis: int = -1,
+    p: int = 2,
+    margin: float = 1.0,
+    eps: float = 1e-6,
+    reduction: str = "none",
+) -> mx.array:
+    r"""
+    Computes the triplet loss for a set of anchor, positive, and negative samples.
+    Margin is represented with alpha in the math section.
+
+    .. math::
+
+       L_{\text{triplet}} = \max\left(\|A - P\|_p - \|A - N\|_p + \alpha, 0\right)
+
+    Args:
+        anchors (array): The anchor samples.
+        positives (array): The positive samples.
+        negatives (array): The negative samples.
+        axis (int, optional): The distribution axis. Default: ``-1``.
+        p (int, optional): The norm degree for pairwise distance. Default: ``2``.
+        margin (float, optional): Margin for the triplet loss. Defaults to ``1.0``.
+        eps (float, optional): Small positive constant to prevent numerical instability. Defaults to ``1e-6``.
+        reduction (str, optional): Specifies the reduction to apply to the output:
+          ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``.
+
+    Returns:
+        array: Computed triplet loss. If reduction is "none", returns a tensor of the same shape as input;
+                  if reduction is "mean" or "sum", returns a scalar tensor.
+    """
+    loss = mx.maximum(
+        mx.sqrt(mx.power(anchors - positives, p).sum(axis) + eps)
+        - mx.sqrt(mx.power(anchors - negatives, p).sum(axis) + eps)
+        + margin,
+        0,
+    )
+    return _reduce(loss, reduction)
+
+
+def _reduce(loss: mx.array, reduction: str = "none"):
+    if reduction == "mean":
+        return mx.mean(loss)
+    elif reduction == "sum":
+        return mx.sum(loss)
+    elif reduction == "none":
+        return loss
+    else:
+        raise ValueError("Invalid reduction. Must be 'none', 'mean', or 'sum'.")
+
+
+def hinge_loss(
+    inputs: mx.array, targets: mx.array, reduction: str = "none"
+) -> mx.array:
+    r"""
+    Computes the hinge loss between inputs and targets.
+
+    .. math::
+
+       \text{hinge}(y, y_{\text{pred}}) = \max(0, 1 - y \cdot y_{\text{pred}})
+
+
+    Args:
+        inputs (array): The predicted values.
+        targets (array): The target values. They should be -1 or 1.
+        reduction (str, optional): Specifies the reduction to apply to the output:
+          ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``.
+
+    Returns:
+        array: The computed hinge loss.
+    """
+    loss = mx.maximum(1 - inputs * targets, 0)
+
+    return _reduce(loss, reduction)
+
+
+def huber_loss(
+    inputs: mx.array, targets: mx.array, delta: float = 1.0, reduction: str = "none"
+) -> mx.array:
+    r"""
+    Computes the Huber loss between inputs and targets.
+
+    .. math::
+
+        L_{\delta}(a) =
+        \left\{ \begin{array}{ll}
+            \frac{1}{2} a^2 & \text{for } |a| \leq \delta, \\
+            \delta \left( |a| - \frac{1}{2} \delta \right) & \text{otherwise.}
+        \end{array} \right.
+
+    Args:
+        inputs (array): The predicted values.
+        targets (array): The target values.
+        delta (float, optional): The threshold at which to change between L1 and L2 loss.
+          Default: ``1.0``.
+        reduction (str, optional): Specifies the reduction to apply to the output:
+          ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``.
+
+    Returns:
+        array: The computed Huber loss.
+    """
+    errors = inputs - targets
+    abs_errors = mx.abs(errors)
+    quadratic = mx.minimum(abs_errors, delta)
+    linear = abs_errors - quadratic
+    loss = 0.5 * quadratic**2 + delta * linear
+
+    return _reduce(loss, reduction)
+
+
+def log_cosh_loss(
+    inputs: mx.array, targets: mx.array, reduction: str = "none"
+) -> mx.array:
+    r"""
+    Computes the log cosh loss between inputs and targets.
+
+    Logcosh acts like L2 loss for small errors, ensuring stable gradients,
+    and like the L1 loss for large errors, reducing sensitivity to outliers. This
+    dual behavior offers a balanced, robust approach for regression tasks.
+
+    .. math::
+
+       \text{logcosh}(y_{\text{true}}, y_{\text{pred}}) =
+            \frac{1}{n} \sum_{i=1}^{n}
+            \log(\cosh(y_{\text{pred}}^{(i)} - y_{\text{true}}^{(i)}))
+
+
+    Args:
+        inputs (array): The predicted values.
+        targets (array): The target values.
+        reduction (str, optional): Specifies the reduction to apply to the output:
+          ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``.
+
+    Returns:
+        array: The computed log cosh loss.
+    """
+    errors = inputs - targets
+    loss = mx.logaddexp(errors, -errors) - math.log(2)
+
+    return _reduce(loss, reduction)

lib/python3.11/site-packages/mlx/nn/utils.py

@@ -0,0 +1,33 @@
+# Copyright © 2023 Apple Inc.
+
+from typing import Callable
+
+import mlx.core as mx
+
+
+def value_and_grad(model: "mlx.nn.Module", fn: Callable):
+    """Transform the passed function ``fn`` to a function that computes the
+    gradients of ``fn`` wrt the model's trainable parameters and also its
+    value.
+
+    Args:
+        model (mlx.nn.Module): The model whose trainable parameters to compute
+                               gradients for
+        fn (Callable): The scalar function to compute gradients for
+
+    Returns:
+        A callable that returns the value of ``fn`` and the gradients wrt the
+        trainable parameters of ``model``
+    """
+
+    def inner_fn(params, *args, **kwargs):
+        model.update(params)
+        return fn(*args, **kwargs)
+
+    value_grad_fn = mx.value_and_grad(inner_fn)
+
+    def wrapped_value_grad_fn(*args, **kwargs):
+        value, grad = value_grad_fn(model.trainable_parameters(), *args, **kwargs)
+        return value, grad
+
+    return wrapped_value_grad_fn

lib/python3.11/site-packages/mlx/optimizers.py

@@ -0,0 +1,500 @@
+# Copyright © 2023 Apple Inc.
+
+import math
+from typing import List
+
+import mlx.core as mx
+from mlx.utils import tree_map
+
+
+class OptimizerState(dict):
+    """The optimizer state implements a recursively defined
+    :class:`collections.defaultdict`, namely a missing key in an optimizer
+    state is an :class:`OptimizerState`.
+
+    .. note::
+       :meth:`OptimizerState.get` in contrast to a normal dictionary also sets
+       the key to the ``default`` value if the ``key`` was not present in the
+       dictionary.
+    """
+
+    def __getitem__(self, key):
+        if key not in self:
+            self[key] = OptimizerState()
+        return super().__getitem__(key)
+
+    def get(self, key, default):
+        """If ``key`` doesn't exist set its value to ``default`` and then return it."""
+        if key not in self:
+            self[key] = default
+        return super().__getitem__(key)
+
+
+class Optimizer:
+    """The base class for all optimizers. It allows us to implement an
+    optimizer on a per-parameter basis and apply it to a parameter tree.
+
+    Attributes:
+        state (OptimizerState): It holds the optimizer's state dictionary.
+    """
+
+    def __init__(self):
+        self.state = OptimizerState()
+
+    def update(self, model: "mlx.nn.Module", gradients: dict):
+        """Apply the gradients to the parameters of the model and update the
+        model with the new parameters.
+
+        Args:
+            model (mlx.nn.Module): An mlx module to be updated.
+            gradients (dict): A Python tree of gradients, most likely computed
+                              via :func:`mlx.nn.value_and_grad`.
+        """
+        model.update(self.apply_gradients(gradients, model))
+
+    def apply_gradients(self, gradients: dict, model: dict):
+        """Apply the gradients to the parameters and return the updated parameters.
+
+        Can be used to update a model via
+        ``model.update(opt.apply_gradients(grads, model))`` which is precisely
+        how :meth:`Optimizer.update` is implemented.
+
+        Args:
+            gradients (dict): A Python tree of gradients.
+            model (dict): A Python tree of parameters. It can be a superset of
+                          the gradients. In that case the returned python tree
+                          will be of the same structure as the gradients.
+        """
+        return tree_map(self.apply_single, gradients, model, self.state)
+
+    def apply_single(
+        self, gradient: mx.array, parameter: mx.array, state: OptimizerState
+    ):
+        """To be extended by the children classes to implement each optimizer's
+        update."""
+        raise NotImplementedError()
+
+
+class SGD(Optimizer):
+    r"""Stochastic gradient descent optimizer.
+
+    Updates a parameter :math:`w` with a gradient :math:`g` as follows
+
+    .. math::
+
+        v_{t+1} &= \mu v_t + (1 - \tau) g_t \\
+        w_{t+1} &= w_t - \lambda v_{t+1}
+
+    Args:
+        learning_rate (float): The learning rate :math:`\lambda`.
+        momentum (float, optional): The momentum strength :math:`\mu`. Default: ``0``
+        weight_decay (float, optional): The weight decay (L2 penalty). Default: ``0``
+        dampening (float, optional): Dampening for momentum :math:`\tau`. Default: ``0``
+        nesterov (bool, optional): Enables Nesterov momentum. Default: ``False``
+    """
+
+    def __init__(
+        self,
+        learning_rate: float,
+        momentum: float = 0.0,
+        weight_decay: float = 0.0,
+        dampening: float = 0.0,
+        nesterov: bool = False,
+    ):
+        if nesterov and (momentum <= 0 or dampening != 0):
+            raise ValueError(
+                "Nesterov momentum requires a momentum and zero dampening."
+            )
+        super().__init__()
+
+        self.learning_rate = learning_rate
+        self.momentum = momentum
+        self.weight_decay = weight_decay
+        self.dampening = dampening
+        self.nesterov = nesterov
+
+    def apply_single(
+        self, gradient: mx.array, parameter: mx.array, state: OptimizerState
+    ):
+        """Performs the SGD parameter update and stores :math:`v` in the
+        optimizer state."""
+        if self.momentum <= 0:
+            return parameter - self.learning_rate * gradient
+
+        v = state.get("v", mx.zeros_like(gradient))
+
+        if self.weight_decay != 0:
+            gradient += self.weight_decay * parameter
+
+        v = self.momentum * v
+        if self.dampening > 0:
+            v += (1 - self.dampening) * gradient
+        else:
+            v += gradient
+
+        if self.nesterov:
+            update = gradient + self.momentum * v
+        else:
+            update = v
+        state["v"] = v
+        return parameter - self.learning_rate * update
+
+
+class RMSprop(Optimizer):
+    r"""Implementation of the RMSprop optimizer [1].
+
+    [1]: Tieleman, T. and Hinton, G. 2012. Lecture 6.5-rmsprop, coursera: Neural networks for machine learning
+
+    .. math::
+
+        v_{t+1} &= \alpha v_t + (1 - \alpha) g_t^2 \\
+        w_{t+1} &= w_t - \lambda \frac{g_t}{\sqrt{v_{t+1}} + \epsilon}
+
+    Args:
+        learning_rate (float): The learning rate :math:`\lambda`.
+        alpha (float, optional): The smoothing constant :math:`\alpha`.
+          Default: ``0.99``
+        eps (float, optional): The term :math:`\epsilon` added to the denominator
+          to improve numerical stability. Default: ``1e-8``
+    """
+
+    def __init__(self, learning_rate: float, alpha: float = 0.99, eps: float = 1e-8):
+        super().__init__()
+
+        self.learning_rate = learning_rate
+        self.alpha = alpha
+        self.eps = eps
+
+        if self.alpha < 0.0:
+            raise ValueError(
+                f"RMSprop alpha should be >=0, {self.alpha} was provided instead"
+            )
+        if self.eps < 0.0:
+            raise ValueError(
+                f"RMSprop epsilon should be >0, {self.eps} was provided instead"
+            )
+
+    def apply_single(
+        self, gradient: mx.array, parameter: mx.array, state: OptimizerState
+    ):
+        """Performs the RMSprop parameter update and stores :math:`v` in the optimizer state."""
+        lr = self.learning_rate
+        alpha = self.alpha
+        eps = self.eps
+
+        v = state.get("v", mx.zeros_like(gradient))
+        v = alpha * v + (1 - alpha) * mx.square(gradient)
+        state["v"] = v
+
+        return parameter - lr * gradient / (mx.sqrt(v) + eps)
+
+
+class Adagrad(Optimizer):
+    r"""Implementation of the Adagrad optimizer [1].
+
+    Our Adagrad implementation follows the original paper. In detail,
+
+    [1]: Duchi, J., Hazan, E. and Singer, Y., 2011. Adaptive subgradient methods
+    for online learning and stochastic optimization. JMLR 2011.
+
+    .. math::
+
+        v_{t+1} &= v_t + g_t^2 \\
+        w_{t+1} &= w_t - \lambda \frac{g_t}{\sqrt{v_{t+1}} + \epsilon}
+
+    Args:
+        learning_rate (float): The learning rate :math:`\lambda`.
+        eps (float, optional): The term :math:`\epsilon` added to the
+          denominator to improve numerical stability. Default: ``1e-8``
+    """
+
+    def __init__(self, learning_rate: float, eps: float = 1e-8):
+        super().__init__()
+
+        self.learning_rate = learning_rate
+        self.eps = eps
+
+        if self.eps < 0.0:
+            raise ValueError(
+                f"Adagrad epsilon should be >0, {self.eps} was provided instead"
+            )
+
+    def apply_single(
+        self, gradient: mx.array, parameter: mx.array, state: OptimizerState
+    ):
+        """Performs the Adagrad parameter update and stores :math:`v` in the
+        optimizer state."""
+        lr = self.learning_rate
+        eps = self.eps
+
+        v = state.get("v", mx.zeros_like(gradient))
+        v = v + mx.square(gradient)
+        state["v"] = v
+
+        return parameter - lr * gradient / (mx.sqrt(v) + eps)
+
+
+class AdaDelta(Optimizer):
+    r"""Implementation of the AdaDelta optimizer with learning rate[1].
+
+    Our AdaDelta implementation follows the original paper. In detail,
+
+    [1]: Zeiler, M.D., 2012. ADADELTA: an adaptive learning rate method. arXiv preprint arXiv:1212.5701.
+
+    .. math::
+
+        v_{t+1} &= \rho v_t + (1 - \rho) g_t^2 \\
+        \Delta w_{t+1} &= \frac{\sqrt{u_t + \epsilon}}{\sqrt{v_{t+1} + \epsilon}} g_t \\
+        u_{t+1} &= \rho u_t + (1 - \rho) \Delta w_{t+1}^2 \\
+        w_{t+1} &= w_t - \lambda \Delta w_{t+1}
+
+    Args:
+        learning_rate (float): The learning rate :math:`\lambda`.
+        rho (float, optional): The coefficient :math:`\rho` used for computing a
+            running average of squared gradients. Default: ``0.9``
+        eps (float, optional): The term :math:`\epsilon` added to the denominator to improve
+          numerical stability. Default: `1e-8`
+    """
+
+    def __init__(self, learning_rate: float, rho: float = 0.9, eps: float = 1e-6):
+        super().__init__()
+
+        self.learning_rate = learning_rate
+        self.rho = rho
+        self.eps = eps
+        if self.rho < 0.0:
+            raise ValueError(
+                f"AdaDelta rho should be >=0, {self.rho} was provided instead"
+            )
+        if self.eps < 0.0:
+            raise ValueError(
+                f"AdaDelta epsilon should be >0, {self.eps} was provided instead"
+            )
+
+    def apply_single(
+        self, gradient: mx.array, parameter: mx.array, state: OptimizerState
+    ):
+        """Performs the AdaDelta parameter update and stores :math:`v` and
+        :math:`u` in the optimizer state."""
+        lr = self.learning_rate
+        rho = self.rho
+        eps = self.eps
+
+        v = state.get("v", mx.zeros_like(gradient))
+        u = state.get("s", mx.zeros_like(gradient))
+
+        v = rho * v + (1 - rho) * mx.square(gradient)
+        d = mx.sqrt(u + eps) / mx.sqrt(v + eps) * gradient
+        u = rho * u + (1 - rho) * mx.square(d)
+
+        state["v"] = v
+        state["u"] = u
+
+        return parameter - lr * d
+
+
+class Adam(Optimizer):
+    r"""Implementation of the Adam optimizer [1].
+
+    Our Adam implementation follows the original paper and omits the bias
+    correction in the first and second moment estimates. In detail,
+
+    [1]: Kingma, D.P. and Ba, J., 2015. Adam: A method for stochastic
+    optimization. ICLR 2015.
+
+    .. math::
+
+        m_{t+1} &= \beta_1 m_t + (1 - \beta_1) g_t \\
+        v_{t+1} &= \beta_2 v_t + (1 - \beta_2) g_t^2 \\
+        w_{t+1} &= w_t - \lambda \frac{m_{t+1}}{\sqrt{v_{t+1} + \epsilon}}
+
+    Args:
+        learning_rate (float): The learning rate :math:`\lambda`.
+        betas (Tuple[float, float], optional): The coefficients
+          :math:`(\beta_1, \beta_2)` used for computing running averages of the
+          gradient and its square. Default: ``(0.9, 0.999)``
+        eps (float, optional): The term :math:`\epsilon` added to the
+          denominator to improve numerical stability. Default: ``1e-8``
+    """
+
+    def __init__(
+        self, learning_rate: float, betas: List[float] = [0.9, 0.999], eps: float = 1e-8
+    ):
+        super().__init__()
+
+        self.learning_rate = learning_rate
+        self.betas = betas
+        self.eps = eps
+
+    def apply_single(
+        self, gradient: mx.array, parameter: mx.array, state: OptimizerState
+    ):
+        """Performs the Adam parameter update and stores :math:`v` and
+        :math:`m` in the optimizer state."""
+        lr = self.learning_rate
+        b1, b2 = self.betas
+        eps = self.eps
+
+        m = state.get("m", gradient)
+        v = state.get("v", mx.square(gradient))
+        m = b1 * m + (1 - b1) * gradient
+        v = b2 * v + (1 - b2) * mx.square(gradient)
+        state["m"] = m
+        state["v"] = v
+
+        return parameter - lr * m / (mx.sqrt(v) + eps)
+
+
+class AdamW(Adam):
+    r"""Implementation of the AdamW optimizer [1].
+
+    Following the above convention, in contrast with [1], we do not use bias
+    correction in the first and second moments for AdamW. We update the weights
+    with a weight_decay (:math:`\lambda`) value:
+
+    [1]: Loshchilov, I. and Hutter, F., 2019. Decoupled weight decay
+    regularization. ICLR 2019.
+
+    .. math::
+
+        m_{t+1} &= \beta_1 m_t + (1 - \beta_1) g_t \\
+        v_{t+1} &= \beta_2 v_t + (1 - \beta_2) g_t^2 \\
+        w_{t+1} &= w_t - \alpha (\frac{m_{t+1}}{\sqrt{v_{t+1} + \epsilon}} + \lambda w_t)
+
+    Args:
+        learning_rate (float): The learning rate :math:`\alpha`.
+        betas (Tuple[float, float], optional): The coefficients
+          :math:`(\beta_1, \beta_2)` used for computing running averages of the
+          gradient and its square. Default: ``(0.9, 0.999)``
+        eps (float, optional): The term :math:`\epsilon` added to the
+          denominator to improve numerical stability. Default: ``1e-8``
+        weight_decay (float, optional): The weight decay :math:`\lambda`.
+          Default: ``0``.
+    """
+
+    def __init__(
+        self,
+        learning_rate: float,
+        betas: List[float] = [0.9, 0.999],
+        eps: float = 1e-8,
+        weight_decay: float = 0.01,
+    ):
+        super().__init__(learning_rate=learning_rate, betas=betas, eps=eps)
+        self.weight_decay = weight_decay
+
+    def apply_single(
+        self, gradient: mx.array, parameter: mx.array, state: OptimizerState
+    ):
+        """Performs the AdamW parameter update by modifying the parameters
+        passed into Adam.
+        """
+
+        return super().apply_single(
+            gradient, parameter * (1 - self.learning_rate * self.weight_decay), state
+        )
+
+
+class Adamax(Adam):
+    r"""Implementation of the Adamax optimizer. It is a variant of Adam based
+    on the infinity norm [1].
+
+    Our Adam implementation follows the original paper and omits the bias
+    correction in the first and second moment estimates. In detail,
+
+    [1]: Kingma, D.P. and Ba, J., 2015. Adam: A method for stochastic
+    optimization. ICLR 2015.
+
+    .. math::
+
+        m_{t+1} &= \beta_1 m_t + (1 - \beta_1) g_t \\
+        v_{t+1} &= \max(\beta_2 v_t, |g_t|) \\
+        w_{t+1} &= w_t - \lambda \frac{m_{t+1}}{v_{t+1} + \epsilon}
+
+    Args:
+        learning_rate (float): The learning rate :math:`\lambda`.
+        betas (Tuple[float, float], optional): The coefficients
+          :math:`(\beta_1, \beta_2)` used for computing running averages of the
+          gradient and its square. Default: ``(0.9, 0.999)``
+        eps (float, optional): The term :math:`\epsilon` added to the
+          denominator to improve numerical stability. Default: ``1e-8``
+    """
+
+    def __init__(
+        self, learning_rate: float, betas: List[float] = [0.9, 0.999], eps: float = 1e-8
+    ):
+        super().__init__(learning_rate, betas, eps)
+
+    def apply_single(
+        self, gradient: mx.array, parameter: mx.array, state: OptimizerState
+    ):
+        """Performs the Adamax parameter update and stores :math:`v` and
+        :math:`m` in the optimizer state."""
+        lr = self.learning_rate
+        b1, b2 = self.betas
+        eps = self.eps
+
+        m = state.get("m", mx.zeros_like(gradient))
+        v = state.get("v", mx.zeros_like(gradient))
+
+        m = b1 * m + (1 - b1) * gradient
+        v = mx.maximum(b2 * v, mx.abs(gradient))
+        state["m"] = m
+        state["v"] = v
+
+        return parameter - lr * m / (v + eps)
+
+
+class Lion(Optimizer):
+    r"""Implementation of the Lion optimizer [1].
+
+    Since updates are computed through the sign operation, they tend to
+    have larger norm than for other optimizers such as SGD and Adam.
+    We recommend a learning rate that is 3-10x smaller than AdamW and a
+    weight decay 3-10x larger than AdamW to maintain the strength
+    (lr * wd). Our Lion implementation follows the original paper. In
+    detail,
+
+    [1]: Chen, X. Symbolic Discovery of Optimization Algorithms. arXiv
+    preprint arXiv:2302.06675.
+
+    .. math::
+
+        c_{t + 1} &= \beta_1 m_t + (1 - \beta_1) g_t
+        m_{t + 1} &= \beta_2 m_t + (1 - \beta_2) g_t
+        w_{t + 1} &= w_t - \eta (\text{sign}(c_t) + \lambda w_t)
+
+    Args:
+        learning_rate (float): The learning rate :math:`\eta`.
+        betas (Tuple[float, float], optional): The coefficients
+          :math:`(\beta_1, \beta_2)` used for computing the gradient
+          momentum and update direction. Default: ``(0.9, 0.99)``
+        weight_decay (float, optional): The weight decay :math:`\lambda`. Default: ``0.0``
+    """
+
+    def __init__(
+        self,
+        learning_rate: float,
+        betas: List[float] = [0.9, 0.99],
+        weight_decay: float = 0.0,
+    ):
+        super().__init__()
+
+        self.learning_rate = learning_rate
+        self.betas = betas
+        self.weight_decay = weight_decay
+
+    def apply_single(
+        self, gradient: mx.array, parameter: mx.array, state: OptimizerState
+    ):
+        """Performs the Lion parameter update and stores :math:`m`
+        in the optimizer state."""
+        lr = self.learning_rate
+        b1, b2 = self.betas
+        weight_decay = self.weight_decay
+
+        m = state.get("m", gradient)
+        c = b1 * m + (1 - b1) * gradient
+        state["m"] = b2 * m + (1 - b2) * gradient
+        if weight_decay > 0:
+            parameter = (1 - lr * weight_decay) * parameter
+        return parameter - lr * mx.sign(c)

lib/python3.11/site-packages/mlx/share/cmake/MLX/MLXConfig.cmake

@@ -0,0 +1,57 @@
+# Find MLX
+#
+# Defines the following variables:
+#
+#   MLX_FOUND            : True if MLX is found
+#   MLX_INCLUDE_DIRS     : Include directory
+#   MLX_LIBRARIES        : Libraries to link against
+#   MLX_CXX_FLAGS        : Additional compiler flags
+#   MLX_BUILD_ACCELERATE : True if MLX was built with accelerate 
+#   MLX_BUILD_METAL      : True if MLX was built with metal 
+
+
+####### Expanded from @PACKAGE_INIT@ by configure_package_config_file() #######
+####### Any changes to this file will be overwritten by the next CMake run ####
+####### The input file was mlx.pc.in                            ########
+
+get_filename_component(PACKAGE_PREFIX_DIR "${CMAKE_CURRENT_LIST_DIR}/../../../" ABSOLUTE)
+
+macro(set_and_check _var _file)
+  set(${_var} "${_file}")
+  if(NOT EXISTS "${_file}")
+    message(FATAL_ERROR "File or directory ${_file} referenced by variable ${_var} does not exist !")
+  endif()
+endmacro()
+
+####################################################################################
+
+include(${PACKAGE_PREFIX_DIR}/share/cmake/MLX/MLXTargets.cmake)
+include(${PACKAGE_PREFIX_DIR}/share/cmake/MLX/extension.cmake)
+
+set_and_check(MLX_LIBRARY_DIRS ${PACKAGE_PREFIX_DIR}/lib)
+set_and_check(MLX_INCLUDE_DIRS ${PACKAGE_PREFIX_DIR}/include)
+set(MLX_LIBRARIES mlx)
+
+find_library(MLX_LIBRARY mlx PATHS ${MLX_LIBRARY_DIRS})
+
+if (ON)
+    set(MLX_BUILD_ACCELERATE ON)
+    set(MLX_CXX_FLAGS ${MLX_CXX_FLAGS} -DACCELERATE_NEW_LAPACK)
+endif()
+
+if (ON)
+    set(MLX_BUILD_METAL ON)
+    set(MLX_CXX_FLAGS ${MLX_CXX_FLAGS} -D_METAL_)
+    set_and_check(MLX_INCLUDE_DIRS 
+        ${MLX_INCLUDE_DIRS} 
+        ${PACKAGE_PREFIX_DIR}/include/metal_cpp
+    )
+endif()
+
+set_target_properties(mlx PROPERTIES
+    CXX_STANDARD 17
+    INTERFACE_COMPILE_OPTIONS "${MLX_CXX_FLAGS}"
+)
+
+include(FindPackageHandleStandardArgs)
+find_package_handle_standard_args(MLX DEFAULT_MSG MLX_LIBRARY MLX_INCLUDE_DIRS)

lib/python3.11/site-packages/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.0.7")
+
+if(PACKAGE_VERSION VERSION_LESS PACKAGE_FIND_VERSION)
+  set(PACKAGE_VERSION_COMPATIBLE FALSE)
+else()
+
+  if("0.0.7" 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.0.7")
+  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()

lib/python3.11/site-packages/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.dylib"
+  IMPORTED_SONAME_RELEASE "@rpath/libmlx.dylib"
+  )
+
+list(APPEND _cmake_import_check_targets mlx )
+list(APPEND _cmake_import_check_files_for_mlx "${_IMPORT_PREFIX}/lib/libmlx.dylib" )
+
+# Commands beyond this point should not need to know the version.
+set(CMAKE_IMPORT_FILE_VERSION)

lib/python3.11/site-packages/mlx/share/cmake/MLX/MLXTargets.cmake

@@ -0,0 +1,107 @@
+# Generated by CMake
+
+if("${CMAKE_MAJOR_VERSION}.${CMAKE_MINOR_VERSION}" LESS 2.8)
+   message(FATAL_ERROR "CMake >= 2.8.0 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...3.24)
+#----------------------------------------------------------------
+# 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/metal_cpp;${_IMPORT_PREFIX}/include/json;${_IMPORT_PREFIX}/include;${_IMPORT_PREFIX}/include"
+  INTERFACE_LINK_LIBRARIES "/Applications/Xcode-beta.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX14.0.sdk/System/Library/Frameworks/Metal.framework;/Applications/Xcode-beta.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX14.0.sdk/System/Library/Frameworks/Foundation.framework;/Applications/Xcode-beta.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX14.0.sdk/System/Library/Frameworks/QuartzCore.framework;/Applications/Xcode-beta.app/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX14.0.sdk/System/Library/Frameworks/Accelerate.framework"
+)
+
+if(CMAKE_VERSION VERSION_LESS 2.8.12)
+  message(FATAL_ERROR "This file relies on consumers using CMake 2.8.12 or greater.")
+endif()
+
+# 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)
+  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()
+  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)

lib/python3.11/site-packages/mlx/share/cmake/MLX/extension.cmake

@@ -0,0 +1,56 @@
+include(CMakeParseArguments)
+
+###############################################################################
+# 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)
+#
+macro(mlx_build_metallib)
+  # Parse args
+  set(oneValueArgs TARGET TITLE OUTPUT_DIRECTORY)
+  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)
+
+  # 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)

lib/python3.11/site-packages/mlx/utils.py

@@ -0,0 +1,145 @@
+# Copyright © 2023 Apple Inc.
+
+from collections import defaultdict
+
+
+def tree_map(fn, tree, *rest, is_leaf=None):
+    """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 (Optional[Callable]): 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)
+        return TreeType(
+            tree_map(fn, child, *(r[i] for r in rest), is_leaf=is_leaf)
+            for i, child in enumerate(tree)
+        )
+    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_flatten(tree, prefix="", is_leaf=None):
+    """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]]], ".hello"))
+        # [("hello.0.0.0", 0)]
+
+    .. 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.
+
+    Returns:
+        List[Tuple[str, Any]]: The flat representation of the python tree.
+    """
+    flat_tree = []
+
+    if is_leaf is None or not is_leaf(tree):
+        if isinstance(tree, (list, tuple)):
+            for i, t in enumerate(tree):
+                flat_tree.extend(tree_flatten(t, f"{prefix}.{i}", is_leaf))
+            return flat_tree
+        if isinstance(tree, dict):
+            for k, t in tree.items():
+                flat_tree.extend(tree_flatten(t, f"{prefix}.{k}", is_leaf))
+            return flat_tree
+
+    return [(prefix[1:], tree)]
+
+
+def tree_unflatten(tree):
+    """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}}
+
+    Args:
+        tree (List[Tuple[str, Any]]): The flat representation of a python tree.
+                                      For instance as returned by :meth:`tree_flatten`.
+
+    Returns:
+        A python tree.
+    """
+    if len(tree) == 1 and tree[0][0] == "":
+        return tree[0][1]
+
+    try:
+        int(tree[0][0].split(".", maxsplit=1)[0])
+        is_list = True
+    except ValueError:
+        is_list = False
+
+    # collect children
+    children = defaultdict(list)
+    for key, value in tree:
+        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))
+
+    # recursively map them to the original container
+    if is_list:
+        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
+    else:
+        return {k: tree_unflatten(v) for k, v in children.items()}

lib/python3.11/site-packages/more_itertools-10.1.0.dist-info/INSTALLER

@@ -0,0 +1 @@
+pip

lib/python3.11/site-packages/more_itertools-10.1.0.dist-info/LICENSE

@@ -0,0 +1,19 @@
+Copyright (c) 2012 Erik Rose
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of
+this software and associated documentation files (the "Software"), to deal in
+the Software without restriction, including without limitation the rights to
+use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
+of the Software, and to permit persons to whom the Software is furnished to do
+so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

lib/python3.11/site-packages/more_itertools-10.1.0.dist-info/METADATA

@@ -0,0 +1,253 @@
+Metadata-Version: 2.1
+Name: more-itertools
+Version: 10.1.0
+Summary: More routines for operating on iterables, beyond itertools
+Keywords: itertools,iterator,iteration,filter,peek,peekable,chunk,chunked
+Author-email: Erik Rose <erikrose@grinchcentral.com>
+Requires-Python: >=3.8
+Description-Content-Type: text/x-rst
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: Intended Audience :: Developers
+Classifier: Natural Language :: English
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: Implementation :: CPython
+Classifier: Programming Language :: Python :: Implementation :: PyPy
+Classifier: Topic :: Software Development :: Libraries
+Project-URL: Homepage, https://github.com/more-itertools/more-itertools
+
+==============
+More Itertools
+==============
+
+.. image:: https://readthedocs.org/projects/more-itertools/badge/?version=latest
+  :target: https://more-itertools.readthedocs.io/en/stable/
+
+Python's ``itertools`` library is a gem - you can compose elegant solutions
+for a variety of problems with the functions it provides. In ``more-itertools``
+we collect additional building blocks, recipes, and routines for working with
+Python iterables.
+
++------------------------+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+| Grouping               | `chunked <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.chunked>`_,                                                                               |
+|                        | `ichunked <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.ichunked>`_,                                                                             |
+|                        | `chunked_even <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.chunked_even>`_,                                                                     |
+|                        | `sliced <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.sliced>`_,                                                                                 |
+|                        | `constrained_batches <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.constrained_batches>`_,                                                       |
+|                        | `distribute <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.distribute>`_,                                                                         |
+|                        | `divide <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.divide>`_,                                                                                 |
+|                        | `split_at <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.split_at>`_,                                                                             |
+|                        | `split_before <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.split_before>`_,                                                                     |
+|                        | `split_after <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.split_after>`_,                                                                       |
+|                        | `split_into <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.split_into>`_,                                                                         |
+|                        | `split_when <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.split_when>`_,                                                                         |
+|                        | `bucket <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.bucket>`_,                                                                                 |
+|                        | `unzip <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.unzip>`_,                                                                                   |
+|                        | `batched <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.batched>`_,                                                                               |
+|                        | `grouper <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.grouper>`_,                                                                               |
+|                        | `partition <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.partition>`_,                                                                           |
+|                        | `transpose <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.transpose>`_                                                                            |
++------------------------+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+| Lookahead and lookback | `spy <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.spy>`_,                                                                                       |
+|                        | `peekable <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.peekable>`_,                                                                             |
+|                        | `seekable <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.seekable>`_                                                                              |
++------------------------+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+| Windowing              | `windowed <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.windowed>`_,                                                                             |
+|                        | `substrings <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.substrings>`_,                                                                         |
+|                        | `substrings_indexes <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.substrings_indexes>`_,                                                         |
+|                        | `stagger <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.stagger>`_,                                                                               |
+|                        | `windowed_complete <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.windowed_complete>`_,                                                           |
+|                        | `pairwise <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.pairwise>`_,                                                                             |
+|                        | `triplewise <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.triplewise>`_,                                                                         |
+|                        | `sliding_window <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.sliding_window>`_,                                                                 |
+|                        | `subslices <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.subslices>`_                                                                            |
++------------------------+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+| Augmenting             | `count_cycle <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.count_cycle>`_,                                                                       |
+|                        | `intersperse <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.intersperse>`_,                                                                       |
+|                        | `padded <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.padded>`_,                                                                                 |
+|                        | `repeat_each <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.repeat_each>`_,                                                                       |
+|                        | `mark_ends <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.mark_ends>`_,                                                                           |
+|                        | `repeat_last <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.repeat_last>`_,                                                                       |
+|                        | `adjacent <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.adjacent>`_,                                                                             |
+|                        | `groupby_transform <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.groupby_transform>`_,                                                           |
+|                        | `pad_none <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.pad_none>`_,                                                                             |
+|                        | `ncycles <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.ncycles>`_                                                                                |
++------------------------+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+| Combining              | `collapse <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.collapse>`_,                                                                             |
+|                        | `sort_together <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.sort_together>`_,                                                                   |
+|                        | `interleave <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.interleave>`_,                                                                         |
+|                        | `interleave_longest <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.interleave_longest>`_,                                                         |
+|                        | `interleave_evenly <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.interleave_evenly>`_,                                                           |
+|                        | `zip_offset <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.zip_offset>`_,                                                                         |
+|                        | `zip_equal <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.zip_equal>`_,                                                                           |
+|                        | `zip_broadcast <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.zip_broadcast>`_,                                                                   |
+|                        | `dotproduct <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.dotproduct>`_,                                                                         |
+|                        | `convolve <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.convolve>`_,                                                                             |
+|                        | `flatten <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.flatten>`_,                                                                               |
+|                        | `roundrobin <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.roundrobin>`_,                                                                         |
+|                        | `prepend <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.prepend>`_,                                                                               |
+|                        | `value_chain <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.value_chain>`_,                                                                       |
+|                        | `partial_product <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.partial_product>`_                                                                |
++------------------------+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+| Summarizing            | `ilen <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.ilen>`_,                                                                                     |
+|                        | `unique_to_each <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.unique_to_each>`_,                                                                 |
+|                        | `sample <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.sample>`_,                                                                                 |
+|                        | `consecutive_groups <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.consecutive_groups>`_,                                                         |
+|                        | `run_length <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.run_length>`_,                                                                         |
+|                        | `map_reduce <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.map_reduce>`_,                                                                         |
+|                        | `exactly_n <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.exactly_n>`_,                                                                           |
+|                        | `is_sorted <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.is_sorted>`_,                                                                           |
+|                        | `all_equal <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.all_equal>`_,                                                                           |
+|                        | `all_unique <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.all_unique>`_,                                                                         |
+|                        | `minmax <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.minmax>`_,                                                                                 |
+|                        | `first_true <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.first_true>`_,                                                                         |
+|                        | `quantify <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.quantify>`_,                                                                             |
+|                        | `iequals <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.iequals>`_                                                                                |
++------------------------+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+| Selecting              | `islice_extended <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.islice_extended>`_,                                                               |
+|                        | `first <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.first>`_,                                                                                   |
+|                        | `last <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.last>`_,                                                                                     |
+|                        | `one <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.one>`_,                                                                                       |
+|                        | `only <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.only>`_,                                                                                     |
+|                        | `strictly_n <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.strictly_n>`_,                                                                         |
+|                        | `strip <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.strip>`_,                                                                                   |
+|                        | `lstrip <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.lstrip>`_,                                                                                 |
+|                        | `rstrip <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.rstrip>`_,                                                                                 |
+|                        | `filter_except <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.filter_except>`_,                                                                   |
+|                        | `map_except <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.map_except>`_,                                                                         |
+|                        | `nth_or_last <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.nth_or_last>`_,                                                                       |
+|                        | `unique_in_window <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.unique_in_window>`_,                                                             |
+|                        | `before_and_after <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.before_and_after>`_,                                                             |
+|                        | `nth <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.nth>`_,                                                                                       |
+|                        | `take <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.take>`_,                                                                                     |
+|                        | `tail <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.tail>`_,                                                                                     |
+|                        | `unique_everseen <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.unique_everseen>`_,                                                               |
+|                        | `unique_justseen <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.unique_justseen>`_,                                                               |
+|                        | `duplicates_everseen <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.duplicates_everseen>`_,                                                       |
+|                        | `duplicates_justseen <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.duplicates_justseen>`_,                                                       |
+|                        | `longest_common_prefix <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.longest_common_prefix>`_,                                                   |
+|                        | `takewhile_inclusive <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.takewhile_inclusive>`_                                                        |
++------------------------+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+| Combinatorics          | `distinct_permutations <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.distinct_permutations>`_,                                                   |
+|                        | `distinct_combinations <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.distinct_combinations>`_,                                                   |
+|                        | `circular_shifts <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.circular_shifts>`_,                                                               |
+|                        | `partitions <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.partitions>`_,                                                                         |
+|                        | `set_partitions <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.set_partitions>`_,                                                                 |
+|                        | `product_index <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.product_index>`_,                                                                   |
+|                        | `combination_index <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.combination_index>`_,                                                           |
+|                        | `permutation_index <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.permutation_index>`_,                                                           |
+|                        | `combination_with_replacement_index <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.combination_with_replacement_index>`_,                         |
+|                        | `gray_product  <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.gray_product>`_,                                                                    |
+|                        | `outer_product  <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.outer_product>`_,                                                                  |
+|                        | `powerset <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.powerset>`_,                                                                             |
+|                        | `random_product <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.random_product>`_,                                                                 |
+|                        | `random_permutation <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.random_permutation>`_,                                                         |
+|                        | `random_combination <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.random_combination>`_,                                                         |
+|                        | `random_combination_with_replacement <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.random_combination_with_replacement>`_,                       |
+|                        | `nth_product <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.nth_product>`_,                                                                       |
+|                        | `nth_permutation <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.nth_permutation>`_,                                                               |
+|                        | `nth_combination <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.nth_combination>`_,                                                               |
+|                        | `nth_combination_with_replacement <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.nth_combination_with_replacement>`_                              |
++------------------------+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+| Wrapping               | `always_iterable <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.always_iterable>`_,                                                               |
+|                        | `always_reversible <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.always_reversible>`_,                                                           |
+|                        | `countable <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.countable>`_,                                                                           |
+|                        | `consumer <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.consumer>`_,                                                                             |
+|                        | `with_iter <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.with_iter>`_,                                                                           |
+|                        | `iter_except <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.iter_except>`_                                                                        |
++------------------------+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+| Others                 | `locate <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.locate>`_,                                                                                 |
+|                        | `rlocate <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.rlocate>`_,                                                                               |
+|                        | `replace <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.replace>`_,                                                                               |
+|                        | `numeric_range <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.numeric_range>`_,                                                                   |
+|                        | `side_effect <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.side_effect>`_,                                                                       |
+|                        | `iterate <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.iterate>`_,                                                                               |
+|                        | `difference <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.difference>`_,                                                                         |
+|                        | `make_decorator <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.make_decorator>`_,                                                                 |
+|                        | `SequenceView <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.SequenceView>`_,                                                                     |
+|                        | `time_limited <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.time_limited>`_,                                                                     |
+|                        | `map_if <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.map_if>`_,                                                                                 |
+|                        | `iter_index <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.iter_index>`_,                                                                         |
+|                        | `consume <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.consume>`_,                                                                               |
+|                        | `tabulate <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.tabulate>`_,                                                                             |
+|                        | `repeatfunc <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.repeatfunc>`_,                                                                         |
+|                        | `polynomial_from_roots <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.polynomial_from_roots>`_,                                                   |
+|                        | `polynomial_eval <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.polynomial_eval>`_,                                                               |
+|                        | `polynomial_derivative <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.polynomial_derivative>`_,                                                   |
+|                        | `sieve <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.sieve>`_,                                                                                   |
+|                        | `factor <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.factor>`_,                                                                                 |
+|                        | `matmul <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.matmul>`_,                                                                                 |
+|                        | `sum_of_squares <https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.sum_of_squares>`_                                                                  |
++------------------------+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+
+
+Getting started
+===============
+
+To get started, install the library with `pip <https://pip.pypa.io/en/stable/>`_:
+
+.. code-block:: shell
+
+    pip install more-itertools
+
+The recipes from the `itertools docs <https://docs.python.org/3/library/itertools.html#itertools-recipes>`_
+are included in the top-level package:
+
+.. code-block:: python
+
+    >>> from more_itertools import flatten
+    >>> iterable = [(0, 1), (2, 3)]
+    >>> list(flatten(iterable))
+    [0, 1, 2, 3]
+
+Several new recipes are available as well:
+
+.. code-block:: python
+
+    >>> from more_itertools import chunked
+    >>> iterable = [0, 1, 2, 3, 4, 5, 6, 7, 8]
+    >>> list(chunked(iterable, 3))
+    [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
+
+    >>> from more_itertools import spy
+    >>> iterable = (x * x for x in range(1, 6))
+    >>> head, iterable = spy(iterable, n=3)
+    >>> list(head)
+    [1, 4, 9]
+    >>> list(iterable)
+    [1, 4, 9, 16, 25]
+
+
+
+For the full listing of functions, see the `API documentation <https://more-itertools.readthedocs.io/en/stable/api.html>`_.
+
+
+Links elsewhere
+===============
+
+Blog posts about ``more-itertools``:
+
+* `Yo, I heard you like decorators <https://www.bbayles.com/index/decorator_factory>`__
+* `Tour of Python Itertools <https://martinheinz.dev/blog/16>`__ (`Alternate <https://dev.to/martinheinz/tour-of-python-itertools-4122>`__)
+* `Real-World Python More Itertools <https://www.gidware.com/real-world-more-itertools/>`_
+
+
+Development
+===========
+
+``more-itertools`` is maintained by `@erikrose <https://github.com/erikrose>`_
+and `@bbayles <https://github.com/bbayles>`_, with help from `many others <https://github.com/more-itertools/more-itertools/graphs/contributors>`_.
+If you have a problem or suggestion, please file a bug or pull request in this
+repository. Thanks for contributing!
+
+
+Version History
+===============
+
+The version history can be found in `documentation <https://more-itertools.readthedocs.io/en/stable/versions.html>`_.
+

lib/python3.11/site-packages/more_itertools-10.1.0.dist-info/RECORD

@@ -0,0 +1,16 @@
+more_itertools-10.1.0.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4
+more_itertools-10.1.0.dist-info/LICENSE,sha256=CfHIyelBrz5YTVlkHqm4fYPAyw_QB-te85Gn4mQ8GkY,1053
+more_itertools-10.1.0.dist-info/METADATA,sha256=s6T6Rg5Sq9hJE6KhX38MrPxWh8X0d9Fsdld4qdbrQVQ,33830
+more_itertools-10.1.0.dist-info/RECORD,,
+more_itertools-10.1.0.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+more_itertools-10.1.0.dist-info/WHEEL,sha256=rSgq_JpHF9fHR1lx53qwg_1-2LypZE_qmcuXbVUq948,81
+more_itertools/__init__.py,sha256=weQyJxnCVH2EuAaoxC8ZEX-ViZmpHm5kLHO05O8LfRM,149
+more_itertools/__init__.pyi,sha256=5B3eTzON1BBuOLob1vCflyEb2lSd6usXQQ-Cv-hXkeA,43
+more_itertools/__pycache__/__init__.cpython-311.pyc,,
+more_itertools/__pycache__/more.cpython-311.pyc,,
+more_itertools/__pycache__/recipes.cpython-311.pyc,,
+more_itertools/more.py,sha256=6XrPO3vvd3miZb6ygQtY5SoJeLk2PAwfEQOlZHwZgXw,140481
+more_itertools/more.pyi,sha256=nbYRtg-0YR0XMnFOGMzRzbgEOtNC_cfkUaOMOmNVWMA,20726
+more_itertools/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+more_itertools/recipes.py,sha256=GNzJluz0ZYBZhF_H3bC7YkNCm9fZzOoBiQvh1WNNDOU,26354
+more_itertools/recipes.pyi,sha256=77IGo2ZqPtp9IkRu5MzQBgatLBl52vO3MRrFvPDcHzM,4237

lib/python3.11/site-packages/more_itertools-10.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: flit 3.8.0
+Root-Is-Purelib: true
+Tag: py3-none-any

lib/python3.11/site-packages/more_itertools/__init__.py

@@ -0,0 +1,6 @@
+"""More routines for operating on iterables, beyond itertools"""
+
+from .more import *  # noqa
+from .recipes import *  # noqa
+
+__version__ = '10.1.0'

lib/python3.11/site-packages/more_itertools/__init__.pyi

@@ -0,0 +1,2 @@
+from .more import *
+from .recipes import *

lib/python3.11/site-packages/more_itertools/more.pyi

@@ -0,0 +1,684 @@
+"""Stubs for more_itertools.more"""
+from __future__ import annotations
+
+from types import TracebackType
+from typing import (
+    Any,
+    Callable,
+    Container,
+    ContextManager,
+    Generic,
+    Hashable,
+    Iterable,
+    Iterator,
+    overload,
+    Reversible,
+    Sequence,
+    Sized,
+    Type,
+    TypeVar,
+    type_check_only,
+)
+from typing_extensions import Protocol
+
+# Type and type variable definitions
+_T = TypeVar('_T')
+_T1 = TypeVar('_T1')
+_T2 = TypeVar('_T2')
+_U = TypeVar('_U')
+_V = TypeVar('_V')
+_W = TypeVar('_W')
+_T_co = TypeVar('_T_co', covariant=True)
+_GenFn = TypeVar('_GenFn', bound=Callable[..., Iterator[object]])
+_Raisable = BaseException | Type[BaseException]
+
+@type_check_only
+class _SizedIterable(Protocol[_T_co], Sized, Iterable[_T_co]): ...
+
+@type_check_only
+class _SizedReversible(Protocol[_T_co], Sized, Reversible[_T_co]): ...
+
+@type_check_only
+class _SupportsSlicing(Protocol[_T_co]):
+    def __getitem__(self, __k: slice) -> _T_co: ...
+
+def chunked(
+    iterable: Iterable[_T], n: int | None, strict: bool = ...
+) -> Iterator[list[_T]]: ...
+@overload
+def first(iterable: Iterable[_T]) -> _T: ...
+@overload
+def first(iterable: Iterable[_T], default: _U) -> _T | _U: ...
+@overload
+def last(iterable: Iterable[_T]) -> _T: ...
+@overload
+def last(iterable: Iterable[_T], default: _U) -> _T | _U: ...
+@overload
+def nth_or_last(iterable: Iterable[_T], n: int) -> _T: ...
+@overload
+def nth_or_last(iterable: Iterable[_T], n: int, default: _U) -> _T | _U: ...
+
+class peekable(Generic[_T], Iterator[_T]):
+    def __init__(self, iterable: Iterable[_T]) -> None: ...
+    def __iter__(self) -> peekable[_T]: ...
+    def __bool__(self) -> bool: ...
+    @overload
+    def peek(self) -> _T: ...
+    @overload
+    def peek(self, default: _U) -> _T | _U: ...
+    def prepend(self, *items: _T) -> None: ...
+    def __next__(self) -> _T: ...
+    @overload
+    def __getitem__(self, index: int) -> _T: ...
+    @overload
+    def __getitem__(self, index: slice) -> list[_T]: ...
+
+def consumer(func: _GenFn) -> _GenFn: ...
+def ilen(iterable: Iterable[object]) -> int: ...
+def iterate(func: Callable[[_T], _T], start: _T) -> Iterator[_T]: ...
+def with_iter(
+    context_manager: ContextManager[Iterable[_T]],
+) -> Iterator[_T]: ...
+def one(
+    iterable: Iterable[_T],
+    too_short: _Raisable | None = ...,
+    too_long: _Raisable | None = ...,
+) -> _T: ...
+def raise_(exception: _Raisable, *args: Any) -> None: ...
+def strictly_n(
+    iterable: Iterable[_T],
+    n: int,
+    too_short: _GenFn | None = ...,
+    too_long: _GenFn | None = ...,
+) -> list[_T]: ...
+def distinct_permutations(
+    iterable: Iterable[_T], r: int | None = ...
+) -> Iterator[tuple[_T, ...]]: ...
+def intersperse(
+    e: _U, iterable: Iterable[_T], n: int = ...
+) -> Iterator[_T | _U]: ...
+def unique_to_each(*iterables: Iterable[_T]) -> list[list[_T]]: ...
+@overload
+def windowed(
+    seq: Iterable[_T], n: int, *, step: int = ...
+) -> Iterator[tuple[_T | None, ...]]: ...
+@overload
+def windowed(
+    seq: Iterable[_T], n: int, fillvalue: _U, step: int = ...
+) -> Iterator[tuple[_T | _U, ...]]: ...
+def substrings(iterable: Iterable[_T]) -> Iterator[tuple[_T, ...]]: ...
+def substrings_indexes(
+    seq: Sequence[_T], reverse: bool = ...
+) -> Iterator[tuple[Sequence[_T], int, int]]: ...
+
+class bucket(Generic[_T, _U], Container[_U]):
+    def __init__(
+        self,
+        iterable: Iterable[_T],
+        key: Callable[[_T], _U],
+        validator: Callable[[object], object] | None = ...,
+    ) -> None: ...
+    def __contains__(self, value: object) -> bool: ...
+    def __iter__(self) -> Iterator[_U]: ...
+    def __getitem__(self, value: object) -> Iterator[_T]: ...
+
+def spy(
+    iterable: Iterable[_T], n: int = ...
+) -> tuple[list[_T], Iterator[_T]]: ...
+def interleave(*iterables: Iterable[_T]) -> Iterator[_T]: ...
+def interleave_longest(*iterables: Iterable[_T]) -> Iterator[_T]: ...
+def interleave_evenly(
+    iterables: list[Iterable[_T]], lengths: list[int] | None = ...
+) -> Iterator[_T]: ...
+def collapse(
+    iterable: Iterable[Any],
+    base_type: type | None = ...,
+    levels: int | None = ...,
+) -> Iterator[Any]: ...
+@overload
+def side_effect(
+    func: Callable[[_T], object],
+    iterable: Iterable[_T],
+    chunk_size: None = ...,
+    before: Callable[[], object] | None = ...,
+    after: Callable[[], object] | None = ...,
+) -> Iterator[_T]: ...
+@overload
+def side_effect(
+    func: Callable[[list[_T]], object],
+    iterable: Iterable[_T],
+    chunk_size: int,
+    before: Callable[[], object] | None = ...,
+    after: Callable[[], object] | None = ...,
+) -> Iterator[_T]: ...
+def sliced(
+    seq: _SupportsSlicing[_T], n: int, strict: bool = ...
+) -> Iterator[_T]: ...
+def split_at(
+    iterable: Iterable[_T],
+    pred: Callable[[_T], object],
+    maxsplit: int = ...,
+    keep_separator: bool = ...,
+) -> Iterator[list[_T]]: ...
+def split_before(
+    iterable: Iterable[_T], pred: Callable[[_T], object], maxsplit: int = ...
+) -> Iterator[list[_T]]: ...
+def split_after(
+    iterable: Iterable[_T], pred: Callable[[_T], object], maxsplit: int = ...
+) -> Iterator[list[_T]]: ...
+def split_when(
+    iterable: Iterable[_T],
+    pred: Callable[[_T, _T], object],
+    maxsplit: int = ...,
+) -> Iterator[list[_T]]: ...
+def split_into(
+    iterable: Iterable[_T], sizes: Iterable[int | None]
+) -> Iterator[list[_T]]: ...
+@overload
+def padded(
+    iterable: Iterable[_T],
+    *,
+    n: int | None = ...,
+    next_multiple: bool = ...,
+) -> Iterator[_T | None]: ...
+@overload
+def padded(
+    iterable: Iterable[_T],
+    fillvalue: _U,
+    n: int | None = ...,
+    next_multiple: bool = ...,
+) -> Iterator[_T | _U]: ...
+@overload
+def repeat_last(iterable: Iterable[_T]) -> Iterator[_T]: ...
+@overload
+def repeat_last(iterable: Iterable[_T], default: _U) -> Iterator[_T | _U]: ...
+def distribute(n: int, iterable: Iterable[_T]) -> list[Iterator[_T]]: ...
+@overload
+def stagger(
+    iterable: Iterable[_T],
+    offsets: _SizedIterable[int] = ...,
+    longest: bool = ...,
+) -> Iterator[tuple[_T | None, ...]]: ...
+@overload
+def stagger(
+    iterable: Iterable[_T],
+    offsets: _SizedIterable[int] = ...,
+    longest: bool = ...,
+    fillvalue: _U = ...,
+) -> Iterator[tuple[_T | _U, ...]]: ...
+
+class UnequalIterablesError(ValueError):
+    def __init__(self, details: tuple[int, int, int] | None = ...) -> None: ...
+
+@overload
+def zip_equal(__iter1: Iterable[_T1]) -> Iterator[tuple[_T1]]: ...
+@overload
+def zip_equal(
+    __iter1: Iterable[_T1], __iter2: Iterable[_T2]
+) -> Iterator[tuple[_T1, _T2]]: ...
+@overload
+def zip_equal(
+    __iter1: Iterable[_T],
+    __iter2: Iterable[_T],
+    __iter3: Iterable[_T],
+    *iterables: Iterable[_T],
+) -> Iterator[tuple[_T, ...]]: ...
+@overload
+def zip_offset(
+    __iter1: Iterable[_T1],
+    *,
+    offsets: _SizedIterable[int],
+    longest: bool = ...,
+    fillvalue: None = None,
+) -> Iterator[tuple[_T1 | None]]: ...
+@overload
+def zip_offset(
+    __iter1: Iterable[_T1],
+    __iter2: Iterable[_T2],
+    *,
+    offsets: _SizedIterable[int],
+    longest: bool = ...,
+    fillvalue: None = None,
+) -> Iterator[tuple[_T1 | None, _T2 | None]]: ...
+@overload
+def zip_offset(
+    __iter1: Iterable[_T],
+    __iter2: Iterable[_T],
+    __iter3: Iterable[_T],
+    *iterables: Iterable[_T],
+    offsets: _SizedIterable[int],
+    longest: bool = ...,
+    fillvalue: None = None,
+) -> Iterator[tuple[_T | None, ...]]: ...
+@overload
+def zip_offset(
+    __iter1: Iterable[_T1],
+    *,
+    offsets: _SizedIterable[int],
+    longest: bool = ...,
+    fillvalue: _U,
+) -> Iterator[tuple[_T1 | _U]]: ...
+@overload
+def zip_offset(
+    __iter1: Iterable[_T1],
+    __iter2: Iterable[_T2],
+    *,
+    offsets: _SizedIterable[int],
+    longest: bool = ...,
+    fillvalue: _U,
+) -> Iterator[tuple[_T1 | _U, _T2 | _U]]: ...
+@overload
+def zip_offset(
+    __iter1: Iterable[_T],
+    __iter2: Iterable[_T],
+    __iter3: Iterable[_T],
+    *iterables: Iterable[_T],
+    offsets: _SizedIterable[int],
+    longest: bool = ...,
+    fillvalue: _U,
+) -> Iterator[tuple[_T | _U, ...]]: ...
+def sort_together(
+    iterables: Iterable[Iterable[_T]],
+    key_list: Iterable[int] = ...,
+    key: Callable[..., Any] | None = ...,
+    reverse: bool = ...,
+) -> list[tuple[_T, ...]]: ...
+def unzip(iterable: Iterable[Sequence[_T]]) -> tuple[Iterator[_T], ...]: ...
+def divide(n: int, iterable: Iterable[_T]) -> list[Iterator[_T]]: ...
+def always_iterable(
+    obj: object,
+    base_type: type | tuple[type | tuple[Any, ...], ...] | None = ...,
+) -> Iterator[Any]: ...
+def adjacent(
+    predicate: Callable[[_T], bool],
+    iterable: Iterable[_T],
+    distance: int = ...,
+) -> Iterator[tuple[bool, _T]]: ...
+@overload
+def groupby_transform(
+    iterable: Iterable[_T],
+    keyfunc: None = None,
+    valuefunc: None = None,
+    reducefunc: None = None,
+) -> Iterator[tuple[_T, Iterator[_T]]]: ...
+@overload
+def groupby_transform(
+    iterable: Iterable[_T],
+    keyfunc: Callable[[_T], _U],
+    valuefunc: None,
+    reducefunc: None,
+) -> Iterator[tuple[_U, Iterator[_T]]]: ...
+@overload
+def groupby_transform(
+    iterable: Iterable[_T],
+    keyfunc: None,
+    valuefunc: Callable[[_T], _V],
+    reducefunc: None,
+) -> Iterable[tuple[_T, Iterable[_V]]]: ...
+@overload
+def groupby_transform(
+    iterable: Iterable[_T],
+    keyfunc: Callable[[_T], _U],
+    valuefunc: Callable[[_T], _V],
+    reducefunc: None,
+) -> Iterable[tuple[_U, Iterator[_V]]]: ...
+@overload
+def groupby_transform(
+    iterable: Iterable[_T],
+    keyfunc: None,
+    valuefunc: None,
+    reducefunc: Callable[[Iterator[_T]], _W],
+) -> Iterable[tuple[_T, _W]]: ...
+@overload
+def groupby_transform(
+    iterable: Iterable[_T],
+    keyfunc: Callable[[_T], _U],
+    valuefunc: None,
+    reducefunc: Callable[[Iterator[_T]], _W],
+) -> Iterable[tuple[_U, _W]]: ...
+@overload
+def groupby_transform(
+    iterable: Iterable[_T],
+    keyfunc: None,
+    valuefunc: Callable[[_T], _V],
+    reducefunc: Callable[[Iterable[_V]], _W],
+) -> Iterable[tuple[_T, _W]]: ...
+@overload
+def groupby_transform(
+    iterable: Iterable[_T],
+    keyfunc: Callable[[_T], _U],
+    valuefunc: Callable[[_T], _V],
+    reducefunc: Callable[[Iterable[_V]], _W],
+) -> Iterable[tuple[_U, _W]]: ...
+
+class numeric_range(Generic[_T, _U], Sequence[_T], Hashable, Reversible[_T]):
+    @overload
+    def __init__(self, __stop: _T) -> None: ...
+    @overload
+    def __init__(self, __start: _T, __stop: _T) -> None: ...
+    @overload
+    def __init__(self, __start: _T, __stop: _T, __step: _U) -> None: ...
+    def __bool__(self) -> bool: ...
+    def __contains__(self, elem: object) -> bool: ...
+    def __eq__(self, other: object) -> bool: ...
+    @overload
+    def __getitem__(self, key: int) -> _T: ...
+    @overload
+    def __getitem__(self, key: slice) -> numeric_range[_T, _U]: ...
+    def __hash__(self) -> int: ...
+    def __iter__(self) -> Iterator[_T]: ...
+    def __len__(self) -> int: ...
+    def __reduce__(
+        self,
+    ) -> tuple[Type[numeric_range[_T, _U]], tuple[_T, _T, _U]]: ...
+    def __repr__(self) -> str: ...
+    def __reversed__(self) -> Iterator[_T]: ...
+    def count(self, value: _T) -> int: ...
+    def index(self, value: _T) -> int: ...  # type: ignore
+
+def count_cycle(
+    iterable: Iterable[_T], n: int | None = ...
+) -> Iterable[tuple[int, _T]]: ...
+def mark_ends(
+    iterable: Iterable[_T],
+) -> Iterable[tuple[bool, bool, _T]]: ...
+def locate(
+    iterable: Iterable[object],
+    pred: Callable[..., Any] = ...,
+    window_size: int | None = ...,
+) -> Iterator[int]: ...
+def lstrip(
+    iterable: Iterable[_T], pred: Callable[[_T], object]
+) -> Iterator[_T]: ...
+def rstrip(
+    iterable: Iterable[_T], pred: Callable[[_T], object]
+) -> Iterator[_T]: ...
+def strip(
+    iterable: Iterable[_T], pred: Callable[[_T], object]
+) -> Iterator[_T]: ...
+
+class islice_extended(Generic[_T], Iterator[_T]):
+    def __init__(self, iterable: Iterable[_T], *args: int | None) -> None: ...
+    def __iter__(self) -> islice_extended[_T]: ...
+    def __next__(self) -> _T: ...
+    def __getitem__(self, index: slice) -> islice_extended[_T]: ...
+
+def always_reversible(iterable: Iterable[_T]) -> Iterator[_T]: ...
+def consecutive_groups(
+    iterable: Iterable[_T], ordering: Callable[[_T], int] = ...
+) -> Iterator[Iterator[_T]]: ...
+@overload
+def difference(
+    iterable: Iterable[_T],
+    func: Callable[[_T, _T], _U] = ...,
+    *,
+    initial: None = ...,
+) -> Iterator[_T | _U]: ...
+@overload
+def difference(
+    iterable: Iterable[_T], func: Callable[[_T, _T], _U] = ..., *, initial: _U
+) -> Iterator[_U]: ...
+
+class SequenceView(Generic[_T], Sequence[_T]):
+    def __init__(self, target: Sequence[_T]) -> None: ...
+    @overload
+    def __getitem__(self, index: int) -> _T: ...
+    @overload
+    def __getitem__(self, index: slice) -> Sequence[_T]: ...
+    def __len__(self) -> int: ...
+
+class seekable(Generic[_T], Iterator[_T]):
+    def __init__(
+        self, iterable: Iterable[_T], maxlen: int | None = ...
+    ) -> None: ...
+    def __iter__(self) -> seekable[_T]: ...
+    def __next__(self) -> _T: ...
+    def __bool__(self) -> bool: ...
+    @overload
+    def peek(self) -> _T: ...
+    @overload
+    def peek(self, default: _U) -> _T | _U: ...
+    def elements(self) -> SequenceView[_T]: ...
+    def seek(self, index: int) -> None: ...
+    def relative_seek(self, count: int) -> None: ...
+
+class run_length:
+    @staticmethod
+    def encode(iterable: Iterable[_T]) -> Iterator[tuple[_T, int]]: ...
+    @staticmethod
+    def decode(iterable: Iterable[tuple[_T, int]]) -> Iterator[_T]: ...
+
+def exactly_n(
+    iterable: Iterable[_T], n: int, predicate: Callable[[_T], object] = ...
+) -> bool: ...
+def circular_shifts(iterable: Iterable[_T]) -> list[tuple[_T, ...]]: ...
+def make_decorator(
+    wrapping_func: Callable[..., _U], result_index: int = ...
+) -> Callable[..., Callable[[Callable[..., Any]], Callable[..., _U]]]: ...
+@overload
+def map_reduce(
+    iterable: Iterable[_T],
+    keyfunc: Callable[[_T], _U],
+    valuefunc: None = ...,
+    reducefunc: None = ...,
+) -> dict[_U, list[_T]]: ...
+@overload
+def map_reduce(
+    iterable: Iterable[_T],
+    keyfunc: Callable[[_T], _U],
+    valuefunc: Callable[[_T], _V],
+    reducefunc: None = ...,
+) -> dict[_U, list[_V]]: ...
+@overload
+def map_reduce(
+    iterable: Iterable[_T],
+    keyfunc: Callable[[_T], _U],
+    valuefunc: None = ...,
+    reducefunc: Callable[[list[_T]], _W] = ...,
+) -> dict[_U, _W]: ...
+@overload
+def map_reduce(
+    iterable: Iterable[_T],
+    keyfunc: Callable[[_T], _U],
+    valuefunc: Callable[[_T], _V],
+    reducefunc: Callable[[list[_V]], _W],
+) -> dict[_U, _W]: ...
+def rlocate(
+    iterable: Iterable[_T],
+    pred: Callable[..., object] = ...,
+    window_size: int | None = ...,
+) -> Iterator[int]: ...
+def replace(
+    iterable: Iterable[_T],
+    pred: Callable[..., object],
+    substitutes: Iterable[_U],
+    count: int | None = ...,
+    window_size: int = ...,
+) -> Iterator[_T | _U]: ...
+def partitions(iterable: Iterable[_T]) -> Iterator[list[list[_T]]]: ...
+def set_partitions(
+    iterable: Iterable[_T], k: int | None = ...
+) -> Iterator[list[list[_T]]]: ...
+
+class time_limited(Generic[_T], Iterator[_T]):
+    def __init__(
+        self, limit_seconds: float, iterable: Iterable[_T]
+    ) -> None: ...
+    def __iter__(self) -> islice_extended[_T]: ...
+    def __next__(self) -> _T: ...
+
+@overload
+def only(
+    iterable: Iterable[_T], *, too_long: _Raisable | None = ...
+) -> _T | None: ...
+@overload
+def only(
+    iterable: Iterable[_T], default: _U, too_long: _Raisable | None = ...
+) -> _T | _U: ...
+def ichunked(iterable: Iterable[_T], n: int) -> Iterator[Iterator[_T]]: ...
+def distinct_combinations(
+    iterable: Iterable[_T], r: int
+) -> Iterator[tuple[_T, ...]]: ...
+def filter_except(
+    validator: Callable[[Any], object],
+    iterable: Iterable[_T],
+    *exceptions: Type[BaseException],
+) -> Iterator[_T]: ...
+def map_except(
+    function: Callable[[Any], _U],
+    iterable: Iterable[_T],
+    *exceptions: Type[BaseException],
+) -> Iterator[_U]: ...
+def map_if(
+    iterable: Iterable[Any],
+    pred: Callable[[Any], bool],
+    func: Callable[[Any], Any],
+    func_else: Callable[[Any], Any] | None = ...,
+) -> Iterator[Any]: ...
+def sample(
+    iterable: Iterable[_T],
+    k: int,
+    weights: Iterable[float] | None = ...,
+) -> list[_T]: ...
+def is_sorted(
+    iterable: Iterable[_T],
+    key: Callable[[_T], _U] | None = ...,
+    reverse: bool = False,
+    strict: bool = False,
+) -> bool: ...
+
+class AbortThread(BaseException):
+    pass
+
+class callback_iter(Generic[_T], Iterator[_T]):
+    def __init__(
+        self,
+        func: Callable[..., Any],
+        callback_kwd: str = ...,
+        wait_seconds: float = ...,
+    ) -> None: ...
+    def __enter__(self) -> callback_iter[_T]: ...
+    def __exit__(
+        self,
+        exc_type: Type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> bool | None: ...
+    def __iter__(self) -> callback_iter[_T]: ...
+    def __next__(self) -> _T: ...
+    def _reader(self) -> Iterator[_T]: ...
+    @property
+    def done(self) -> bool: ...
+    @property
+    def result(self) -> Any: ...
+
+def windowed_complete(
+    iterable: Iterable[_T], n: int
+) -> Iterator[tuple[_T, ...]]: ...
+def all_unique(
+    iterable: Iterable[_T], key: Callable[[_T], _U] | None = ...
+) -> bool: ...
+def nth_product(index: int, *args: Iterable[_T]) -> tuple[_T, ...]: ...
+def nth_combination_with_replacement(
+    iterable: Iterable[_T], r: int, index: int
+) -> tuple[_T, ...]: ...
+def nth_permutation(
+    iterable: Iterable[_T], r: int, index: int
+) -> tuple[_T, ...]: ...
+def value_chain(*args: _T | Iterable[_T]) -> Iterable[_T]: ...
+def product_index(element: Iterable[_T], *args: Iterable[_T]) -> int: ...
+def combination_index(
+    element: Iterable[_T], iterable: Iterable[_T]
+) -> int: ...
+def combination_with_replacement_index(
+    element: Iterable[_T], iterable: Iterable[_T]
+) -> int: ...
+def permutation_index(
+    element: Iterable[_T], iterable: Iterable[_T]
+) -> int: ...
+def repeat_each(iterable: Iterable[_T], n: int = ...) -> Iterator[_T]: ...
+
+class countable(Generic[_T], Iterator[_T]):
+    def __init__(self, iterable: Iterable[_T]) -> None: ...
+    def __iter__(self) -> countable[_T]: ...
+    def __next__(self) -> _T: ...
+
+def chunked_even(iterable: Iterable[_T], n: int) -> Iterator[list[_T]]: ...
+def zip_broadcast(
+    *objects: _T | Iterable[_T],
+    scalar_types: type | tuple[type | tuple[Any, ...], ...] | None = ...,
+    strict: bool = ...,
+) -> Iterable[tuple[_T, ...]]: ...
+def unique_in_window(
+    iterable: Iterable[_T], n: int, key: Callable[[_T], _U] | None = ...
+) -> Iterator[_T]: ...
+def duplicates_everseen(
+    iterable: Iterable[_T], key: Callable[[_T], _U] | None = ...
+) -> Iterator[_T]: ...
+def duplicates_justseen(
+    iterable: Iterable[_T], key: Callable[[_T], _U] | None = ...
+) -> Iterator[_T]: ...
+
+class _SupportsLessThan(Protocol):
+    def __lt__(self, __other: Any) -> bool: ...
+
+_SupportsLessThanT = TypeVar("_SupportsLessThanT", bound=_SupportsLessThan)
+
+@overload
+def minmax(
+    iterable_or_value: Iterable[_SupportsLessThanT], *, key: None = None
+) -> tuple[_SupportsLessThanT, _SupportsLessThanT]: ...
+@overload
+def minmax(
+    iterable_or_value: Iterable[_T], *, key: Callable[[_T], _SupportsLessThan]
+) -> tuple[_T, _T]: ...
+@overload
+def minmax(
+    iterable_or_value: Iterable[_SupportsLessThanT],
+    *,
+    key: None = None,
+    default: _U,
+) -> _U | tuple[_SupportsLessThanT, _SupportsLessThanT]: ...
+@overload
+def minmax(
+    iterable_or_value: Iterable[_T],
+    *,
+    key: Callable[[_T], _SupportsLessThan],
+    default: _U,
+) -> _U | tuple[_T, _T]: ...
+@overload
+def minmax(
+    iterable_or_value: _SupportsLessThanT,
+    __other: _SupportsLessThanT,
+    *others: _SupportsLessThanT,
+) -> tuple[_SupportsLessThanT, _SupportsLessThanT]: ...
+@overload
+def minmax(
+    iterable_or_value: _T,
+    __other: _T,
+    *others: _T,
+    key: Callable[[_T], _SupportsLessThan],
+) -> tuple[_T, _T]: ...
+def longest_common_prefix(
+    iterables: Iterable[Iterable[_T]],
+) -> Iterator[_T]: ...
+def iequals(*iterables: Iterable[object]) -> bool: ...
+def constrained_batches(
+    iterable: Iterable[object],
+    max_size: int,
+    max_count: int | None = ...,
+    get_len: Callable[[_T], object] = ...,
+    strict: bool = ...,
+) -> Iterator[tuple[_T]]: ...
+def gray_product(*iterables: Iterable[_T]) -> Iterator[tuple[_T, ...]]: ...
+def partial_product(*iterables: Iterable[_T]) -> Iterator[tuple[_T, ...]]: ...
+def takewhile_inclusive(
+    predicate: Callable[[_T], bool], iterable: Iterable[_T]
+) -> Iterator[_T]: ...
+def outer_product(
+    func: Callable[[_T, _U], _V],
+    xs: Iterable[_T],
+    ys: Iterable[_U],
+    *args: Any,
+    **kwargs: Any,
+) -> Iterator[tuple[_V, ...]]: ...

lib/python3.11/site-packages/more_itertools/recipes.py

@@ -0,0 +1,977 @@
+"""Imported from the recipes section of the itertools documentation.
+
+All functions taken from the recipes section of the itertools library docs
+[1]_.
+Some backward-compatible usability improvements have been made.
+
+.. [1] http://docs.python.org/library/itertools.html#recipes
+
+"""
+import math
+import operator
+
+from collections import deque
+from collections.abc import Sized
+from functools import partial, reduce
+from itertools import (
+    chain,
+    combinations,
+    compress,
+    count,
+    cycle,
+    groupby,
+    islice,
+    product,
+    repeat,
+    starmap,
+    tee,
+    zip_longest,
+)
+from random import randrange, sample, choice
+
+__all__ = [
+    'all_equal',
+    'batched',
+    'before_and_after',
+    'consume',
+    'convolve',
+    'dotproduct',
+    'first_true',
+    'factor',
+    'flatten',
+    'grouper',
+    'iter_except',
+    'iter_index',
+    'matmul',
+    'ncycles',
+    'nth',
+    'nth_combination',
+    'padnone',
+    'pad_none',
+    'pairwise',
+    'partition',
+    'polynomial_eval',
+    'polynomial_from_roots',
+    'polynomial_derivative',
+    'powerset',
+    'prepend',
+    'quantify',
+    'random_combination_with_replacement',
+    'random_combination',
+    'random_permutation',
+    'random_product',
+    'repeatfunc',
+    'roundrobin',
+    'sieve',
+    'sliding_window',
+    'subslices',
+    'sum_of_squares',
+    'tabulate',
+    'tail',
+    'take',
+    'transpose',
+    'triplewise',
+    'unique_everseen',
+    'unique_justseen',
+]
+
+_marker = object()
+
+
+# zip with strict is available for Python 3.10+
+try:
+    zip(strict=True)
+except TypeError:
+    _zip_strict = zip
+else:
+    _zip_strict = partial(zip, strict=True)
+
+# math.sumprod is available for Python 3.12+
+_sumprod = getattr(math, 'sumprod', lambda x, y: dotproduct(x, y))
+
+
+def take(n, iterable):
+    """Return first *n* items of the iterable as a list.
+
+        >>> take(3, range(10))
+        [0, 1, 2]
+
+    If there are fewer than *n* items in the iterable, all of them are
+    returned.
+
+        >>> take(10, range(3))
+        [0, 1, 2]
+
+    """
+    return list(islice(iterable, n))
+
+
+def tabulate(function, start=0):
+    """Return an iterator over the results of ``func(start)``,
+    ``func(start + 1)``, ``func(start + 2)``...
+
+    *func* should be a function that accepts one integer argument.
+
+    If *start* is not specified it defaults to 0. It will be incremented each
+    time the iterator is advanced.
+
+        >>> square = lambda x: x ** 2
+        >>> iterator = tabulate(square, -3)
+        >>> take(4, iterator)
+        [9, 4, 1, 0]
+
+    """
+    return map(function, count(start))
+
+
+def tail(n, iterable):
+    """Return an iterator over the last *n* items of *iterable*.
+
+    >>> t = tail(3, 'ABCDEFG')
+    >>> list(t)
+    ['E', 'F', 'G']
+
+    """
+    # If the given iterable has a length, then we can use islice to get its
+    # final elements. Note that if the iterable is not actually Iterable,
+    # either islice or deque will throw a TypeError. This is why we don't
+    # check if it is Iterable.
+    if isinstance(iterable, Sized):
+        yield from islice(iterable, max(0, len(iterable) - n), None)
+    else:
+        yield from iter(deque(iterable, maxlen=n))
+
+
+def consume(iterator, n=None):
+    """Advance *iterable* by *n* steps. If *n* is ``None``, consume it
+    entirely.
+
+    Efficiently exhausts an iterator without returning values. Defaults to
+    consuming the whole iterator, but an optional second argument may be
+    provided to limit consumption.
+
+        >>> i = (x for x in range(10))
+        >>> next(i)
+        0
+        >>> consume(i, 3)
+        >>> next(i)
+        4
+        >>> consume(i)
+        >>> next(i)
+        Traceback (most recent call last):
+          File "<stdin>", line 1, in <module>
+        StopIteration
+
+    If the iterator has fewer items remaining than the provided limit, the
+    whole iterator will be consumed.
+
+        >>> i = (x for x in range(3))
+        >>> consume(i, 5)
+        >>> next(i)
+        Traceback (most recent call last):
+          File "<stdin>", line 1, in <module>
+        StopIteration
+
+    """
+    # Use functions that consume iterators at C speed.
+    if n is None:
+        # feed the entire iterator into a zero-length deque
+        deque(iterator, maxlen=0)
+    else:
+        # advance to the empty slice starting at position n
+        next(islice(iterator, n, n), None)
+
+
+def nth(iterable, n, default=None):
+    """Returns the nth item or a default value.
+
+    >>> l = range(10)
+    >>> nth(l, 3)
+    3
+    >>> nth(l, 20, "zebra")
+    'zebra'
+
+    """
+    return next(islice(iterable, n, None), default)
+
+
+def all_equal(iterable):
+    """
+    Returns ``True`` if all the elements are equal to each other.
+
+        >>> all_equal('aaaa')
+        True
+        >>> all_equal('aaab')
+        False
+
+    """
+    g = groupby(iterable)
+    return next(g, True) and not next(g, False)
+
+
+def quantify(iterable, pred=bool):
+    """Return the how many times the predicate is true.
+
+    >>> quantify([True, False, True])
+    2
+
+    """
+    return sum(map(pred, iterable))
+
+
+def pad_none(iterable):
+    """Returns the sequence of elements and then returns ``None`` indefinitely.
+
+        >>> take(5, pad_none(range(3)))
+        [0, 1, 2, None, None]
+
+    Useful for emulating the behavior of the built-in :func:`map` function.
+
+    See also :func:`padded`.
+
+    """
+    return chain(iterable, repeat(None))
+
+
+padnone = pad_none
+
+
+def ncycles(iterable, n):
+    """Returns the sequence elements *n* times
+
+    >>> list(ncycles(["a", "b"], 3))
+    ['a', 'b', 'a', 'b', 'a', 'b']
+
+    """
+    return chain.from_iterable(repeat(tuple(iterable), n))
+
+
+def dotproduct(vec1, vec2):
+    """Returns the dot product of the two iterables.
+
+    >>> dotproduct([10, 10], [20, 20])
+    400
+
+    """
+    return sum(map(operator.mul, vec1, vec2))
+
+
+def flatten(listOfLists):
+    """Return an iterator flattening one level of nesting in a list of lists.
+
+        >>> list(flatten([[0, 1], [2, 3]]))
+        [0, 1, 2, 3]
+
+    See also :func:`collapse`, which can flatten multiple levels of nesting.
+
+    """
+    return chain.from_iterable(listOfLists)
+
+
+def repeatfunc(func, times=None, *args):
+    """Call *func* with *args* repeatedly, returning an iterable over the
+    results.
+
+    If *times* is specified, the iterable will terminate after that many
+    repetitions:
+
+        >>> from operator import add
+        >>> times = 4
+        >>> args = 3, 5
+        >>> list(repeatfunc(add, times, *args))
+        [8, 8, 8, 8]
+
+    If *times* is ``None`` the iterable will not terminate:
+
+        >>> from random import randrange
+        >>> times = None
+        >>> args = 1, 11
+        >>> take(6, repeatfunc(randrange, times, *args))  # doctest:+SKIP
+        [2, 4, 8, 1, 8, 4]
+
+    """
+    if times is None:
+        return starmap(func, repeat(args))
+    return starmap(func, repeat(args, times))
+
+
+def _pairwise(iterable):
+    """Returns an iterator of paired items, overlapping, from the original
+
+    >>> take(4, pairwise(count()))
+    [(0, 1), (1, 2), (2, 3), (3, 4)]
+
+    On Python 3.10 and above, this is an alias for :func:`itertools.pairwise`.
+
+    """
+    a, b = tee(iterable)
+    next(b, None)
+    return zip(a, b)
+
+
+try:
+    from itertools import pairwise as itertools_pairwise
+except ImportError:
+    pairwise = _pairwise
+else:
+
+    def pairwise(iterable):
+        return itertools_pairwise(iterable)
+
+    pairwise.__doc__ = _pairwise.__doc__
+
+
+class UnequalIterablesError(ValueError):
+    def __init__(self, details=None):
+        msg = 'Iterables have different lengths'
+        if details is not None:
+            msg += (': index 0 has length {}; index {} has length {}').format(
+                *details
+            )
+
+        super().__init__(msg)
+
+
+def _zip_equal_generator(iterables):
+    for combo in zip_longest(*iterables, fillvalue=_marker):
+        for val in combo:
+            if val is _marker:
+                raise UnequalIterablesError()
+        yield combo
+
+
+def _zip_equal(*iterables):
+    # Check whether the iterables are all the same size.
+    try:
+        first_size = len(iterables[0])
+        for i, it in enumerate(iterables[1:], 1):
+            size = len(it)
+            if size != first_size:
+                raise UnequalIterablesError(details=(first_size, i, size))
+        # All sizes are equal, we can use the built-in zip.
+        return zip(*iterables)
+    # If any one of the iterables didn't have a length, start reading
+    # them until one runs out.
+    except TypeError:
+        return _zip_equal_generator(iterables)
+
+
+def grouper(iterable, n, incomplete='fill', fillvalue=None):
+    """Group elements from *iterable* into fixed-length groups of length *n*.
+
+    >>> list(grouper('ABCDEF', 3))
+    [('A', 'B', 'C'), ('D', 'E', 'F')]
+
+    The keyword arguments *incomplete* and *fillvalue* control what happens for
+    iterables whose length is not a multiple of *n*.
+
+    When *incomplete* is `'fill'`, the last group will contain instances of
+    *fillvalue*.
+
+    >>> list(grouper('ABCDEFG', 3, incomplete='fill', fillvalue='x'))
+    [('A', 'B', 'C'), ('D', 'E', 'F'), ('G', 'x', 'x')]
+
+    When *incomplete* is `'ignore'`, the last group will not be emitted.
+
+    >>> list(grouper('ABCDEFG', 3, incomplete='ignore', fillvalue='x'))
+    [('A', 'B', 'C'), ('D', 'E', 'F')]
+
+    When *incomplete* is `'strict'`, a subclass of `ValueError` will be raised.
+
+    >>> it = grouper('ABCDEFG', 3, incomplete='strict')
+    >>> list(it)  # doctest: +IGNORE_EXCEPTION_DETAIL
+    Traceback (most recent call last):
+    ...
+    UnequalIterablesError
+
+    """
+    args = [iter(iterable)] * n
+    if incomplete == 'fill':
+        return zip_longest(*args, fillvalue=fillvalue)
+    if incomplete == 'strict':
+        return _zip_equal(*args)
+    if incomplete == 'ignore':
+        return zip(*args)
+    else:
+        raise ValueError('Expected fill, strict, or ignore')
+
+
+def roundrobin(*iterables):
+    """Yields an item from each iterable, alternating between them.
+
+        >>> list(roundrobin('ABC', 'D', 'EF'))
+        ['A', 'D', 'E', 'B', 'F', 'C']
+
+    This function produces the same output as :func:`interleave_longest`, but
+    may perform better for some inputs (in particular when the number of
+    iterables is small).
+
+    """
+    # Recipe credited to George Sakkis
+    pending = len(iterables)
+    nexts = cycle(iter(it).__next__ for it in iterables)
+    while pending:
+        try:
+            for next in nexts:
+                yield next()
+        except StopIteration:
+            pending -= 1
+            nexts = cycle(islice(nexts, pending))
+
+
+def partition(pred, iterable):
+    """
+    Returns a 2-tuple of iterables derived from the input iterable.
+    The first yields the items that have ``pred(item) == False``.
+    The second yields the items that have ``pred(item) == True``.
+
+        >>> is_odd = lambda x: x % 2 != 0
+        >>> iterable = range(10)
+        >>> even_items, odd_items = partition(is_odd, iterable)
+        >>> list(even_items), list(odd_items)
+        ([0, 2, 4, 6, 8], [1, 3, 5, 7, 9])
+
+    If *pred* is None, :func:`bool` is used.
+
+        >>> iterable = [0, 1, False, True, '', ' ']
+        >>> false_items, true_items = partition(None, iterable)
+        >>> list(false_items), list(true_items)
+        ([0, False, ''], [1, True, ' '])
+
+    """
+    if pred is None:
+        pred = bool
+
+    t1, t2, p = tee(iterable, 3)
+    p1, p2 = tee(map(pred, p))
+    return (compress(t1, map(operator.not_, p1)), compress(t2, p2))
+
+
+def powerset(iterable):
+    """Yields all possible subsets of the iterable.
+
+        >>> list(powerset([1, 2, 3]))
+        [(), (1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
+
+    :func:`powerset` will operate on iterables that aren't :class:`set`
+    instances, so repeated elements in the input will produce repeated elements
+    in the output. Use :func:`unique_everseen` on the input to avoid generating
+    duplicates:
+
+        >>> seq = [1, 1, 0]
+        >>> list(powerset(seq))
+        [(), (1,), (1,), (0,), (1, 1), (1, 0), (1, 0), (1, 1, 0)]
+        >>> from more_itertools import unique_everseen
+        >>> list(powerset(unique_everseen(seq)))
+        [(), (1,), (0,), (1, 0)]
+
+    """
+    s = list(iterable)
+    return chain.from_iterable(combinations(s, r) for r in range(len(s) + 1))
+
+
+def unique_everseen(iterable, key=None):
+    """
+    Yield unique elements, preserving order.
+
+        >>> list(unique_everseen('AAAABBBCCDAABBB'))
+        ['A', 'B', 'C', 'D']
+        >>> list(unique_everseen('ABBCcAD', str.lower))
+        ['A', 'B', 'C', 'D']
+
+    Sequences with a mix of hashable and unhashable items can be used.
+    The function will be slower (i.e., `O(n^2)`) for unhashable items.
+
+    Remember that ``list`` objects are unhashable - you can use the *key*
+    parameter to transform the list to a tuple (which is hashable) to
+    avoid a slowdown.
+
+        >>> iterable = ([1, 2], [2, 3], [1, 2])
+        >>> list(unique_everseen(iterable))  # Slow
+        [[1, 2], [2, 3]]
+        >>> list(unique_everseen(iterable, key=tuple))  # Faster
+        [[1, 2], [2, 3]]
+
+    Similary, you may want to convert unhashable ``set`` objects with
+    ``key=frozenset``. For ``dict`` objects,
+    ``key=lambda x: frozenset(x.items())`` can be used.
+
+    """
+    seenset = set()
+    seenset_add = seenset.add
+    seenlist = []
+    seenlist_add = seenlist.append
+    use_key = key is not None
+
+    for element in iterable:
+        k = key(element) if use_key else element
+        try:
+            if k not in seenset:
+                seenset_add(k)
+                yield element
+        except TypeError:
+            if k not in seenlist:
+                seenlist_add(k)
+                yield element
+
+
+def unique_justseen(iterable, key=None):
+    """Yields elements in order, ignoring serial duplicates
+
+    >>> list(unique_justseen('AAAABBBCCDAABBB'))
+    ['A', 'B', 'C', 'D', 'A', 'B']
+    >>> list(unique_justseen('ABBCcAD', str.lower))
+    ['A', 'B', 'C', 'A', 'D']
+
+    """
+    return map(next, map(operator.itemgetter(1), groupby(iterable, key)))
+
+
+def iter_except(func, exception, first=None):
+    """Yields results from a function repeatedly until an exception is raised.
+
+    Converts a call-until-exception interface to an iterator interface.
+    Like ``iter(func, sentinel)``, but uses an exception instead of a sentinel
+    to end the loop.
+
+        >>> l = [0, 1, 2]
+        >>> list(iter_except(l.pop, IndexError))
+        [2, 1, 0]
+
+    Multiple exceptions can be specified as a stopping condition:
+
+        >>> l = [1, 2, 3, '...', 4, 5, 6]
+        >>> list(iter_except(lambda: 1 + l.pop(), (IndexError, TypeError)))
+        [7, 6, 5]
+        >>> list(iter_except(lambda: 1 + l.pop(), (IndexError, TypeError)))
+        [4, 3, 2]
+        >>> list(iter_except(lambda: 1 + l.pop(), (IndexError, TypeError)))
+        []
+
+    """
+    try:
+        if first is not None:
+            yield first()
+        while 1:
+            yield func()
+    except exception:
+        pass
+
+
+def first_true(iterable, default=None, pred=None):
+    """
+    Returns the first true value in the iterable.
+
+    If no true value is found, returns *default*
+
+    If *pred* is not None, returns the first item for which
+    ``pred(item) == True`` .
+
+        >>> first_true(range(10))
+        1
+        >>> first_true(range(10), pred=lambda x: x > 5)
+        6
+        >>> first_true(range(10), default='missing', pred=lambda x: x > 9)
+        'missing'
+
+    """
+    return next(filter(pred, iterable), default)
+
+
+def random_product(*args, repeat=1):
+    """Draw an item at random from each of the input iterables.
+
+        >>> random_product('abc', range(4), 'XYZ')  # doctest:+SKIP
+        ('c', 3, 'Z')
+
+    If *repeat* is provided as a keyword argument, that many items will be
+    drawn from each iterable.
+
+        >>> random_product('abcd', range(4), repeat=2)  # doctest:+SKIP
+        ('a', 2, 'd', 3)
+
+    This equivalent to taking a random selection from
+    ``itertools.product(*args, **kwarg)``.
+
+    """
+    pools = [tuple(pool) for pool in args] * repeat
+    return tuple(choice(pool) for pool in pools)
+
+
+def random_permutation(iterable, r=None):
+    """Return a random *r* length permutation of the elements in *iterable*.
+
+    If *r* is not specified or is ``None``, then *r* defaults to the length of
+    *iterable*.
+
+        >>> random_permutation(range(5))  # doctest:+SKIP
+        (3, 4, 0, 1, 2)
+
+    This equivalent to taking a random selection from
+    ``itertools.permutations(iterable, r)``.
+
+    """
+    pool = tuple(iterable)
+    r = len(pool) if r is None else r
+    return tuple(sample(pool, r))
+
+
+def random_combination(iterable, r):
+    """Return a random *r* length subsequence of the elements in *iterable*.
+
+        >>> random_combination(range(5), 3)  # doctest:+SKIP
+        (2, 3, 4)
+
+    This equivalent to taking a random selection from
+    ``itertools.combinations(iterable, r)``.
+
+    """
+    pool = tuple(iterable)
+    n = len(pool)
+    indices = sorted(sample(range(n), r))
+    return tuple(pool[i] for i in indices)
+
+
+def random_combination_with_replacement(iterable, r):
+    """Return a random *r* length subsequence of elements in *iterable*,
+    allowing individual elements to be repeated.
+
+        >>> random_combination_with_replacement(range(3), 5) # doctest:+SKIP
+        (0, 0, 1, 2, 2)
+
+    This equivalent to taking a random selection from
+    ``itertools.combinations_with_replacement(iterable, r)``.
+
+    """
+    pool = tuple(iterable)
+    n = len(pool)
+    indices = sorted(randrange(n) for i in range(r))
+    return tuple(pool[i] for i in indices)
+
+
+def nth_combination(iterable, r, index):
+    """Equivalent to ``list(combinations(iterable, r))[index]``.
+
+    The subsequences of *iterable* that are of length *r* can be ordered
+    lexicographically. :func:`nth_combination` computes the subsequence at
+    sort position *index* directly, without computing the previous
+    subsequences.
+
+        >>> nth_combination(range(5), 3, 5)
+        (0, 3, 4)
+
+    ``ValueError`` will be raised If *r* is negative or greater than the length
+    of *iterable*.
+    ``IndexError`` will be raised if the given *index* is invalid.
+    """
+    pool = tuple(iterable)
+    n = len(pool)
+    if (r < 0) or (r > n):
+        raise ValueError
+
+    c = 1
+    k = min(r, n - r)
+    for i in range(1, k + 1):
+        c = c * (n - k + i) // i
+
+    if index < 0:
+        index += c
+
+    if (index < 0) or (index >= c):
+        raise IndexError
+
+    result = []
+    while r:
+        c, n, r = c * r // n, n - 1, r - 1
+        while index >= c:
+            index -= c
+            c, n = c * (n - r) // n, n - 1
+        result.append(pool[-1 - n])
+
+    return tuple(result)
+
+
+def prepend(value, iterator):
+    """Yield *value*, followed by the elements in *iterator*.
+
+        >>> value = '0'
+        >>> iterator = ['1', '2', '3']
+        >>> list(prepend(value, iterator))
+        ['0', '1', '2', '3']
+
+    To prepend multiple values, see :func:`itertools.chain`
+    or :func:`value_chain`.
+
+    """
+    return chain([value], iterator)
+
+
+def convolve(signal, kernel):
+    """Convolve the iterable *signal* with the iterable *kernel*.
+
+        >>> signal = (1, 2, 3, 4, 5)
+        >>> kernel = [3, 2, 1]
+        >>> list(convolve(signal, kernel))
+        [3, 8, 14, 20, 26, 14, 5]
+
+    Note: the input arguments are not interchangeable, as the *kernel*
+    is immediately consumed and stored.
+
+    """
+    # This implementation intentionally doesn't match the one in the itertools
+    # documentation.
+    kernel = tuple(kernel)[::-1]
+    n = len(kernel)
+    window = deque([0], maxlen=n) * n
+    for x in chain(signal, repeat(0, n - 1)):
+        window.append(x)
+        yield _sumprod(kernel, window)
+
+
+def before_and_after(predicate, it):
+    """A variant of :func:`takewhile` that allows complete access to the
+    remainder of the iterator.
+
+         >>> it = iter('ABCdEfGhI')
+         >>> all_upper, remainder = before_and_after(str.isupper, it)
+         >>> ''.join(all_upper)
+         'ABC'
+         >>> ''.join(remainder) # takewhile() would lose the 'd'
+         'dEfGhI'
+
+    Note that the first iterator must be fully consumed before the second
+    iterator can generate valid results.
+    """
+    it = iter(it)
+    transition = []
+
+    def true_iterator():
+        for elem in it:
+            if predicate(elem):
+                yield elem
+            else:
+                transition.append(elem)
+                return
+
+    # Note: this is different from itertools recipes to allow nesting
+    # before_and_after remainders into before_and_after again. See tests
+    # for an example.
+    remainder_iterator = chain(transition, it)
+
+    return true_iterator(), remainder_iterator
+
+
+def triplewise(iterable):
+    """Return overlapping triplets from *iterable*.
+
+    >>> list(triplewise('ABCDE'))
+    [('A', 'B', 'C'), ('B', 'C', 'D'), ('C', 'D', 'E')]
+
+    """
+    for (a, _), (b, c) in pairwise(pairwise(iterable)):
+        yield a, b, c
+
+
+def sliding_window(iterable, n):
+    """Return a sliding window of width *n* over *iterable*.
+
+        >>> list(sliding_window(range(6), 4))
+        [(0, 1, 2, 3), (1, 2, 3, 4), (2, 3, 4, 5)]
+
+    If *iterable* has fewer than *n* items, then nothing is yielded:
+
+        >>> list(sliding_window(range(3), 4))
+        []
+
+    For a variant with more features, see :func:`windowed`.
+    """
+    it = iter(iterable)
+    window = deque(islice(it, n - 1), maxlen=n)
+    for x in it:
+        window.append(x)
+        yield tuple(window)
+
+
+def subslices(iterable):
+    """Return all contiguous non-empty subslices of *iterable*.
+
+        >>> list(subslices('ABC'))
+        [['A'], ['A', 'B'], ['A', 'B', 'C'], ['B'], ['B', 'C'], ['C']]
+
+    This is similar to :func:`substrings`, but emits items in a different
+    order.
+    """
+    seq = list(iterable)
+    slices = starmap(slice, combinations(range(len(seq) + 1), 2))
+    return map(operator.getitem, repeat(seq), slices)
+
+
+def polynomial_from_roots(roots):
+    """Compute a polynomial's coefficients from its roots.
+
+    >>> roots = [5, -4, 3]  # (x - 5) * (x + 4) * (x - 3)
+    >>> polynomial_from_roots(roots)  # x^3 - 4 * x^2 - 17 * x + 60
+    [1, -4, -17, 60]
+    """
+    factors = zip(repeat(1), map(operator.neg, roots))
+    return list(reduce(convolve, factors, [1]))
+
+
+def iter_index(iterable, value, start=0):
+    """Yield the index of each place in *iterable* that *value* occurs,
+    beginning with index *start*.
+
+    See :func:`locate` for a more general means of finding the indexes
+    associated with particular values.
+
+    >>> list(iter_index('AABCADEAF', 'A'))
+    [0, 1, 4, 7]
+    """
+    try:
+        seq_index = iterable.index
+    except AttributeError:
+        # Slow path for general iterables
+        it = islice(iterable, start, None)
+        i = start - 1
+        try:
+            while True:
+                i = i + operator.indexOf(it, value) + 1
+                yield i
+        except ValueError:
+            pass
+    else:
+        # Fast path for sequences
+        i = start - 1
+        try:
+            while True:
+                i = seq_index(value, i + 1)
+                yield i
+        except ValueError:
+            pass
+
+
+def sieve(n):
+    """Yield the primes less than n.
+
+    >>> list(sieve(30))
+    [2, 3, 5, 7, 11, 13, 17, 19, 23, 29]
+    """
+    data = bytearray((0, 1)) * (n // 2)
+    data[:3] = 0, 0, 0
+    limit = math.isqrt(n) + 1
+    for p in compress(range(limit), data):
+        data[p * p : n : p + p] = bytes(len(range(p * p, n, p + p)))
+    data[2] = 1
+    return iter_index(data, 1) if n > 2 else iter([])
+
+
+def _batched(iterable, n):
+    """Batch data into lists of length *n*. The last batch may be shorter.
+
+    >>> list(batched('ABCDEFG', 3))
+    [('A', 'B', 'C'), ('D', 'E', 'F'), ('G',)]
+
+    On Python 3.12 and above, this is an alias for :func:`itertools.batched`.
+    """
+    if n < 1:
+        raise ValueError('n must be at least one')
+    it = iter(iterable)
+    while True:
+        batch = tuple(islice(it, n))
+        if not batch:
+            break
+        yield batch
+
+
+try:
+    from itertools import batched as itertools_batched
+except ImportError:
+    batched = _batched
+else:
+
+    def batched(iterable, n):
+        return itertools_batched(iterable, n)
+
+    batched.__doc__ = _batched.__doc__
+
+
+def transpose(it):
+    """Swap the rows and columns of the input.
+
+    >>> list(transpose([(1, 2, 3), (11, 22, 33)]))
+    [(1, 11), (2, 22), (3, 33)]
+
+    The caller should ensure that the dimensions of the input are compatible.
+    If the input is empty, no output will be produced.
+    """
+    return _zip_strict(*it)
+
+
+def matmul(m1, m2):
+    """Multiply two matrices.
+    >>> list(matmul([(7, 5), (3, 5)], [(2, 5), (7, 9)]))
+    [(49, 80), (41, 60)]
+
+    The caller should ensure that the dimensions of the input matrices are
+    compatible with each other.
+    """
+    n = len(m2[0])
+    return batched(starmap(_sumprod, product(m1, transpose(m2))), n)
+
+
+def factor(n):
+    """Yield the prime factors of n.
+    >>> list(factor(360))
+    [2, 2, 2, 3, 3, 5]
+    """
+    for prime in sieve(math.isqrt(n) + 1):
+        while True:
+            if n % prime:
+                break
+            yield prime
+            n //= prime
+            if n == 1:
+                return
+    if n > 1:
+        yield n
+
+
+def polynomial_eval(coefficients, x):
+    """Evaluate a polynomial at a specific value.
+
+    Example: evaluating x^3 - 4 * x^2 - 17 * x + 60 at x = 2.5:
+
+    >>> coefficients = [1, -4, -17, 60]
+    >>> x = 2.5
+    >>> polynomial_eval(coefficients, x)
+    8.125
+    """
+    n = len(coefficients)
+    if n == 0:
+        return x * 0  # coerce zero to the type of x
+    powers = map(pow, repeat(x), reversed(range(n)))
+    return _sumprod(coefficients, powers)
+
+
+def sum_of_squares(it):
+    """Return the sum of the squares of the input values.
+
+    >>> sum_of_squares([10, 20, 30])
+    1400
+    """
+    return _sumprod(*tee(it))
+
+
+def polynomial_derivative(coefficients):
+    """Compute the first derivative of a polynomial.
+
+    Example: evaluating the derivative of x^3 - 4 * x^2 - 17 * x + 60
+
+    >>> coefficients = [1, -4, -17, 60]
+    >>> derivative_coefficients = polynomial_derivative(coefficients)
+    >>> derivative_coefficients
+    [3, -8, -17]
+    """
+    n = len(coefficients)
+    powers = reversed(range(1, n))
+    return list(map(operator.mul, coefficients, powers))

lib/python3.11/site-packages/more_itertools/recipes.pyi

@@ -0,0 +1,122 @@
+"""Stubs for more_itertools.recipes"""
+from __future__ import annotations
+
+from typing import (
+    Any,
+    Callable,
+    Iterable,
+    Iterator,
+    overload,
+    Sequence,
+    Type,
+    TypeVar,
+)
+
+# Type and type variable definitions
+_T = TypeVar('_T')
+_U = TypeVar('_U')
+
+def take(n: int, iterable: Iterable[_T]) -> list[_T]: ...
+def tabulate(
+    function: Callable[[int], _T], start: int = ...
+) -> Iterator[_T]: ...
+def tail(n: int, iterable: Iterable[_T]) -> Iterator[_T]: ...
+def consume(iterator: Iterable[_T], n: int | None = ...) -> None: ...
+@overload
+def nth(iterable: Iterable[_T], n: int) -> _T | None: ...
+@overload
+def nth(iterable: Iterable[_T], n: int, default: _U) -> _T | _U: ...
+def all_equal(iterable: Iterable[object]) -> bool: ...
+def quantify(
+    iterable: Iterable[_T], pred: Callable[[_T], bool] = ...
+) -> int: ...
+def pad_none(iterable: Iterable[_T]) -> Iterator[_T | None]: ...
+def padnone(iterable: Iterable[_T]) -> Iterator[_T | None]: ...
+def ncycles(iterable: Iterable[_T], n: int) -> Iterator[_T]: ...
+def dotproduct(vec1: Iterable[object], vec2: Iterable[object]) -> object: ...
+def flatten(listOfLists: Iterable[Iterable[_T]]) -> Iterator[_T]: ...
+def repeatfunc(
+    func: Callable[..., _U], times: int | None = ..., *args: Any
+) -> Iterator[_U]: ...
+def pairwise(iterable: Iterable[_T]) -> Iterator[tuple[_T, _T]]: ...
+def grouper(
+    iterable: Iterable[_T],
+    n: int,
+    incomplete: str = ...,
+    fillvalue: _U = ...,
+) -> Iterator[tuple[_T | _U, ...]]: ...
+def roundrobin(*iterables: Iterable[_T]) -> Iterator[_T]: ...
+def partition(
+    pred: Callable[[_T], object] | None, iterable: Iterable[_T]
+) -> tuple[Iterator[_T], Iterator[_T]]: ...
+def powerset(iterable: Iterable[_T]) -> Iterator[tuple[_T, ...]]: ...
+def unique_everseen(
+    iterable: Iterable[_T], key: Callable[[_T], _U] | None = ...
+) -> Iterator[_T]: ...
+def unique_justseen(
+    iterable: Iterable[_T], key: Callable[[_T], object] | None = ...
+) -> Iterator[_T]: ...
+@overload
+def iter_except(
+    func: Callable[[], _T],
+    exception: Type[BaseException] | tuple[Type[BaseException], ...],
+    first: None = ...,
+) -> Iterator[_T]: ...
+@overload
+def iter_except(
+    func: Callable[[], _T],
+    exception: Type[BaseException] | tuple[Type[BaseException], ...],
+    first: Callable[[], _U],
+) -> Iterator[_T | _U]: ...
+@overload
+def first_true(
+    iterable: Iterable[_T], *, pred: Callable[[_T], object] | None = ...
+) -> _T | None: ...
+@overload
+def first_true(
+    iterable: Iterable[_T],
+    default: _U,
+    pred: Callable[[_T], object] | None = ...,
+) -> _T | _U: ...
+def random_product(
+    *args: Iterable[_T], repeat: int = ...
+) -> tuple[_T, ...]: ...
+def random_permutation(
+    iterable: Iterable[_T], r: int | None = ...
+) -> tuple[_T, ...]: ...
+def random_combination(iterable: Iterable[_T], r: int) -> tuple[_T, ...]: ...
+def random_combination_with_replacement(
+    iterable: Iterable[_T], r: int
+) -> tuple[_T, ...]: ...
+def nth_combination(
+    iterable: Iterable[_T], r: int, index: int
+) -> tuple[_T, ...]: ...
+def prepend(value: _T, iterator: Iterable[_U]) -> Iterator[_T | _U]: ...
+def convolve(signal: Iterable[_T], kernel: Iterable[_T]) -> Iterator[_T]: ...
+def before_and_after(
+    predicate: Callable[[_T], bool], it: Iterable[_T]
+) -> tuple[Iterator[_T], Iterator[_T]]: ...
+def triplewise(iterable: Iterable[_T]) -> Iterator[tuple[_T, _T, _T]]: ...
+def sliding_window(
+    iterable: Iterable[_T], n: int
+) -> Iterator[tuple[_T, ...]]: ...
+def subslices(iterable: Iterable[_T]) -> Iterator[list[_T]]: ...
+def polynomial_from_roots(roots: Sequence[_T]) -> list[_T]: ...
+def iter_index(
+    iterable: Iterable[object],
+    value: Any,
+    start: int | None = ...,
+) -> Iterator[int]: ...
+def sieve(n: int) -> Iterator[int]: ...
+def batched(
+    iterable: Iterable[_T],
+    n: int,
+) -> Iterator[tuple[_T]]: ...
+def transpose(
+    it: Iterable[Iterable[_T]],
+) -> Iterator[tuple[_T, ...]]: ...
+def matmul(m1: Sequence[_T], m2: Sequence[_T]) -> Iterator[tuple[_T]]: ...
+def factor(n: int) -> Iterator[int]: ...
+def polynomial_eval(coefficients: Sequence[_T], x: _U) -> _U: ...
+def sum_of_squares(it: Iterable[_T]) -> _T: ...
+def polynomial_derivative(coefficients: Sequence[_T]) -> list[_T]: ...

lib/python3.11/site-packages/mpmath/__init__.py

@@ -0,0 +1,468 @@
+__version__ = '1.3.0'
+
+from .usertools import monitor, timing
+
+from .ctx_fp import FPContext
+from .ctx_mp import MPContext
+from .ctx_iv import MPIntervalContext
+
+fp = FPContext()
+mp = MPContext()
+iv = MPIntervalContext()
+
+fp._mp = mp
+mp._mp = mp
+iv._mp = mp
+mp._fp = fp
+fp._fp = fp
+mp._iv = iv
+fp._iv = iv
+iv._iv = iv
+
+# XXX: extremely bad pickle hack
+from . import ctx_mp as _ctx_mp
+_ctx_mp._mpf_module.mpf = mp.mpf
+_ctx_mp._mpf_module.mpc = mp.mpc
+
+make_mpf = mp.make_mpf
+make_mpc = mp.make_mpc
+
+extraprec = mp.extraprec
+extradps = mp.extradps
+workprec = mp.workprec
+workdps = mp.workdps
+autoprec = mp.autoprec
+maxcalls = mp.maxcalls
+memoize = mp.memoize
+
+mag = mp.mag
+
+bernfrac = mp.bernfrac
+
+qfrom = mp.qfrom
+mfrom = mp.mfrom
+kfrom = mp.kfrom
+taufrom = mp.taufrom
+qbarfrom = mp.qbarfrom
+ellipfun = mp.ellipfun
+jtheta = mp.jtheta
+kleinj = mp.kleinj
+eta = mp.eta
+
+qp = mp.qp
+qhyper = mp.qhyper
+qgamma = mp.qgamma
+qfac = mp.qfac
+
+nint_distance = mp.nint_distance
+
+plot = mp.plot
+cplot = mp.cplot
+splot = mp.splot
+
+odefun = mp.odefun
+
+jacobian = mp.jacobian
+findroot = mp.findroot
+multiplicity = mp.multiplicity
+
+isinf = mp.isinf
+isnan = mp.isnan
+isnormal = mp.isnormal
+isint = mp.isint
+isfinite = mp.isfinite
+almosteq = mp.almosteq
+nan = mp.nan
+rand = mp.rand
+
+absmin = mp.absmin
+absmax = mp.absmax
+
+fraction = mp.fraction
+
+linspace = mp.linspace
+arange = mp.arange
+
+mpmathify = convert = mp.convert
+mpc = mp.mpc
+
+mpi = iv._mpi
+
+nstr = mp.nstr
+nprint = mp.nprint
+chop = mp.chop
+
+fneg = mp.fneg
+fadd = mp.fadd
+fsub = mp.fsub
+fmul = mp.fmul
+fdiv = mp.fdiv
+fprod = mp.fprod
+
+quad = mp.quad
+quadgl = mp.quadgl
+quadts = mp.quadts
+quadosc = mp.quadosc
+quadsubdiv = mp.quadsubdiv
+
+invertlaplace = mp.invertlaplace
+invlaptalbot = mp.invlaptalbot
+invlapstehfest = mp.invlapstehfest
+invlapdehoog = mp.invlapdehoog
+
+pslq = mp.pslq
+identify = mp.identify
+findpoly = mp.findpoly
+
+richardson = mp.richardson
+shanks = mp.shanks
+levin = mp.levin
+cohen_alt = mp.cohen_alt
+nsum = mp.nsum
+nprod = mp.nprod
+difference = mp.difference
+diff = mp.diff
+diffs = mp.diffs
+diffs_prod = mp.diffs_prod
+diffs_exp = mp.diffs_exp
+diffun = mp.diffun
+differint = mp.differint
+taylor = mp.taylor
+pade = mp.pade
+polyval = mp.polyval
+polyroots = mp.polyroots
+fourier = mp.fourier
+fourierval = mp.fourierval
+sumem = mp.sumem
+sumap = mp.sumap
+chebyfit = mp.chebyfit
+limit = mp.limit
+
+matrix = mp.matrix
+eye = mp.eye
+diag = mp.diag
+zeros = mp.zeros
+ones = mp.ones
+hilbert = mp.hilbert
+randmatrix = mp.randmatrix
+swap_row = mp.swap_row
+extend = mp.extend
+norm = mp.norm
+mnorm = mp.mnorm
+
+lu_solve = mp.lu_solve
+lu = mp.lu
+qr = mp.qr
+unitvector = mp.unitvector
+inverse = mp.inverse
+residual = mp.residual
+qr_solve = mp.qr_solve
+cholesky = mp.cholesky
+cholesky_solve = mp.cholesky_solve
+det = mp.det
+cond = mp.cond
+hessenberg = mp.hessenberg
+schur = mp.schur
+eig = mp.eig
+eig_sort = mp.eig_sort
+eigsy = mp.eigsy
+eighe = mp.eighe
+eigh = mp.eigh
+svd_r = mp.svd_r
+svd_c = mp.svd_c
+svd = mp.svd
+gauss_quadrature = mp.gauss_quadrature
+
+expm = mp.expm
+sqrtm = mp.sqrtm
+powm = mp.powm
+logm = mp.logm
+sinm = mp.sinm
+cosm = mp.cosm
+
+mpf = mp.mpf
+j = mp.j
+exp = mp.exp
+expj = mp.expj
+expjpi = mp.expjpi
+ln = mp.ln
+im = mp.im
+re = mp.re
+inf = mp.inf
+ninf = mp.ninf
+sign = mp.sign
+
+eps = mp.eps
+pi = mp.pi
+ln2 = mp.ln2
+ln10 = mp.ln10
+phi = mp.phi
+e = mp.e
+euler = mp.euler
+catalan = mp.catalan
+khinchin = mp.khinchin
+glaisher = mp.glaisher
+apery = mp.apery
+degree = mp.degree
+twinprime = mp.twinprime
+mertens = mp.mertens
+
+ldexp = mp.ldexp
+frexp = mp.frexp
+
+fsum = mp.fsum
+fdot = mp.fdot
+
+sqrt = mp.sqrt
+cbrt = mp.cbrt
+exp = mp.exp
+ln = mp.ln
+log = mp.log
+log10 = mp.log10
+power = mp.power
+cos = mp.cos
+sin = mp.sin
+tan = mp.tan
+cosh = mp.cosh
+sinh = mp.sinh
+tanh = mp.tanh
+acos = mp.acos
+asin = mp.asin
+atan = mp.atan
+asinh = mp.asinh
+acosh = mp.acosh
+atanh = mp.atanh
+sec = mp.sec
+csc = mp.csc
+cot = mp.cot
+sech = mp.sech
+csch = mp.csch
+coth = mp.coth
+asec = mp.asec
+acsc = mp.acsc
+acot = mp.acot
+asech = mp.asech
+acsch = mp.acsch
+acoth = mp.acoth
+cospi = mp.cospi
+sinpi = mp.sinpi
+sinc = mp.sinc
+sincpi = mp.sincpi
+cos_sin = mp.cos_sin
+cospi_sinpi = mp.cospi_sinpi
+fabs = mp.fabs
+re = mp.re
+im = mp.im
+conj = mp.conj
+floor = mp.floor
+ceil = mp.ceil
+nint = mp.nint
+frac = mp.frac
+root = mp.root
+nthroot = mp.nthroot
+hypot = mp.hypot
+fmod = mp.fmod
+ldexp = mp.ldexp
+frexp = mp.frexp
+sign = mp.sign
+arg = mp.arg
+phase = mp.phase
+polar = mp.polar
+rect = mp.rect
+degrees = mp.degrees
+radians = mp.radians
+atan2 = mp.atan2
+fib = mp.fib
+fibonacci = mp.fibonacci
+lambertw = mp.lambertw
+zeta = mp.zeta
+altzeta = mp.altzeta
+gamma = mp.gamma
+rgamma = mp.rgamma
+factorial = mp.factorial
+fac = mp.fac
+fac2 = mp.fac2
+beta = mp.beta
+betainc = mp.betainc
+psi = mp.psi
+#psi0 = mp.psi0
+#psi1 = mp.psi1
+#psi2 = mp.psi2
+#psi3 = mp.psi3
+polygamma = mp.polygamma
+digamma = mp.digamma
+#trigamma = mp.trigamma
+#tetragamma = mp.tetragamma
+#pentagamma = mp.pentagamma
+harmonic = mp.harmonic
+bernoulli = mp.bernoulli
+bernfrac = mp.bernfrac
+stieltjes = mp.stieltjes
+hurwitz = mp.hurwitz
+dirichlet = mp.dirichlet
+bernpoly = mp.bernpoly
+eulerpoly = mp.eulerpoly
+eulernum = mp.eulernum
+polylog = mp.polylog
+clsin = mp.clsin
+clcos = mp.clcos
+gammainc = mp.gammainc
+gammaprod = mp.gammaprod
+binomial = mp.binomial
+rf = mp.rf
+ff = mp.ff
+hyper = mp.hyper
+hyp0f1 = mp.hyp0f1
+hyp1f1 = mp.hyp1f1
+hyp1f2 = mp.hyp1f2
+hyp2f1 = mp.hyp2f1
+hyp2f2 = mp.hyp2f2
+hyp2f0 = mp.hyp2f0
+hyp2f3 = mp.hyp2f3
+hyp3f2 = mp.hyp3f2
+hyperu = mp.hyperu
+hypercomb = mp.hypercomb
+meijerg = mp.meijerg
+appellf1 = mp.appellf1
+appellf2 = mp.appellf2
+appellf3 = mp.appellf3
+appellf4 = mp.appellf4
+hyper2d = mp.hyper2d
+bihyper = mp.bihyper
+erf = mp.erf
+erfc = mp.erfc
+erfi = mp.erfi
+erfinv = mp.erfinv
+npdf = mp.npdf
+ncdf = mp.ncdf
+expint = mp.expint
+e1 = mp.e1
+ei = mp.ei
+li = mp.li
+ci = mp.ci
+si = mp.si
+chi = mp.chi
+shi = mp.shi
+fresnels = mp.fresnels
+fresnelc = mp.fresnelc
+airyai = mp.airyai
+airybi = mp.airybi
+airyaizero = mp.airyaizero
+airybizero = mp.airybizero
+scorergi = mp.scorergi
+scorerhi = mp.scorerhi
+ellipk = mp.ellipk
+ellipe = mp.ellipe
+ellipf = mp.ellipf
+ellippi = mp.ellippi
+elliprc = mp.elliprc
+elliprj = mp.elliprj
+elliprf = mp.elliprf
+elliprd = mp.elliprd
+elliprg = mp.elliprg
+agm = mp.agm
+jacobi = mp.jacobi
+chebyt = mp.chebyt
+chebyu = mp.chebyu
+legendre = mp.legendre
+legenp = mp.legenp
+legenq = mp.legenq
+hermite = mp.hermite
+pcfd = mp.pcfd
+pcfu = mp.pcfu
+pcfv = mp.pcfv
+pcfw = mp.pcfw
+gegenbauer = mp.gegenbauer
+laguerre = mp.laguerre
+spherharm = mp.spherharm
+besselj = mp.besselj
+j0 = mp.j0
+j1 = mp.j1
+besseli = mp.besseli
+bessely = mp.bessely
+besselk = mp.besselk
+besseljzero = mp.besseljzero
+besselyzero = mp.besselyzero
+hankel1 = mp.hankel1
+hankel2 = mp.hankel2
+struveh = mp.struveh
+struvel = mp.struvel
+angerj = mp.angerj
+webere = mp.webere
+lommels1 = mp.lommels1
+lommels2 = mp.lommels2
+whitm = mp.whitm
+whitw = mp.whitw
+ber = mp.ber
+bei = mp.bei
+ker = mp.ker
+kei = mp.kei
+coulombc = mp.coulombc
+coulombf = mp.coulombf
+coulombg = mp.coulombg
+barnesg = mp.barnesg
+superfac = mp.superfac
+hyperfac = mp.hyperfac
+loggamma = mp.loggamma
+siegeltheta = mp.siegeltheta
+siegelz = mp.siegelz
+grampoint = mp.grampoint
+zetazero = mp.zetazero
+riemannr = mp.riemannr
+primepi = mp.primepi
+primepi2 = mp.primepi2
+primezeta = mp.primezeta
+bell = mp.bell
+polyexp = mp.polyexp
+expm1 = mp.expm1
+log1p = mp.log1p
+powm1 = mp.powm1
+unitroots = mp.unitroots
+cyclotomic = mp.cyclotomic
+mangoldt = mp.mangoldt
+secondzeta = mp.secondzeta
+nzeros = mp.nzeros
+backlunds = mp.backlunds
+lerchphi = mp.lerchphi
+stirling1 = mp.stirling1
+stirling2 = mp.stirling2
+squarew = mp.squarew
+trianglew = mp.trianglew
+sawtoothw = mp.sawtoothw
+unit_triangle = mp.unit_triangle
+sigmoid = mp.sigmoid
+
+# be careful when changing this name, don't use test*!
+def runtests():
+    """
+    Run all mpmath tests and print output.
+    """
+    import os.path
+    from inspect import getsourcefile
+    from .tests import runtests as tests
+    testdir = os.path.dirname(os.path.abspath(getsourcefile(tests)))
+    importdir = os.path.abspath(testdir + '/../..')
+    tests.testit(importdir, testdir)
+
+def doctests(filter=[]):
+    import sys
+    from timeit import default_timer as clock
+    for i, arg in enumerate(sys.argv):
+        if '__init__.py' in arg:
+            filter = [sn for sn in sys.argv[i+1:] if not sn.startswith("-")]
+            break
+    import doctest
+    globs = globals().copy()
+    for obj in globs: #sorted(globs.keys()):
+        if filter:
+            if not sum([pat in obj for pat in filter]):
+                continue
+        sys.stdout.write(str(obj) + " ")
+        sys.stdout.flush()
+        t1 = clock()
+        doctest.run_docstring_examples(globs[obj], {}, verbose=("-v" in sys.argv))
+        t2 = clock()
+        print(round(t2-t1, 3))
+
+if __name__ == '__main__':
+    doctests()