2023-12-05 12:10:03 -08:00
2023-11-29 12:41:56 -08:00
<!DOCTYPE html>
2023-12-05 12:10:03 -08:00
2024-03-30 17:32:20 -07:00
< html lang = "en" data-content_root = "../" >
2023-12-05 12:10:03 -08:00
< head >
< meta charset = "utf-8" / >
< meta name = "viewport" content = "width=device-width, initial-scale=1.0" / > < meta name = "generator" content = "Docutils 0.18.1: http://docutils.sourceforge.net/" / >
2024-04-26 08:24:09 -07:00
< title > LLM inference — MLX 0.12.0 documentation< / title >
2023-12-05 12:10:03 -08:00
< script data-cfasync = "false" >
document.documentElement.dataset.mode = localStorage.getItem("mode") || "";
document.documentElement.dataset.theme = localStorage.getItem("theme") || "light";
< / script >
<!-- Loaded before other Sphinx assets -->
< link href = "../_static/styles/theme.css?digest=5b4479735964841361fd" rel = "stylesheet" / >
< link href = "../_static/styles/bootstrap.css?digest=5b4479735964841361fd" rel = "stylesheet" / >
< link href = "../_static/styles/pydata-sphinx-theme.css?digest=5b4479735964841361fd" rel = "stylesheet" / >
< link href = "../_static/vendor/fontawesome/6.1.2/css/all.min.css?digest=5b4479735964841361fd" rel = "stylesheet" / >
< link rel = "preload" as = "font" type = "font/woff2" crossorigin href = "../_static/vendor/fontawesome/6.1.2/webfonts/fa-solid-900.woff2" / >
< link rel = "preload" as = "font" type = "font/woff2" crossorigin href = "../_static/vendor/fontawesome/6.1.2/webfonts/fa-brands-400.woff2" / >
< link rel = "preload" as = "font" type = "font/woff2" crossorigin href = "../_static/vendor/fontawesome/6.1.2/webfonts/fa-regular-400.woff2" / >
2024-03-30 17:32:20 -07:00
< link rel = "stylesheet" type = "text/css" href = "../_static/pygments.css?v=fa44fd50" / >
< link rel = "stylesheet" type = "text/css" href = "../_static/styles/sphinx-book-theme.css?v=384b581d" / >
2023-12-05 12:10:03 -08:00
<!-- Pre - loaded scripts that we'll load fully later -->
< link rel = "preload" as = "script" href = "../_static/scripts/bootstrap.js?digest=5b4479735964841361fd" / >
< link rel = "preload" as = "script" href = "../_static/scripts/pydata-sphinx-theme.js?digest=5b4479735964841361fd" / >
< script src = "../_static/vendor/fontawesome/6.1.2/js/all.min.js?digest=5b4479735964841361fd" > < / script >
2024-04-26 08:24:09 -07:00
< script src = "../_static/documentation_options.js?v=e9e33cf0" > < / script >
2024-03-30 17:32:20 -07:00
< script src = "../_static/doctools.js?v=888ff710" > < / script >
< script src = "../_static/sphinx_highlight.js?v=dc90522c" > < / script >
< script src = "../_static/scripts/sphinx-book-theme.js?v=efea14e4" > < / script >
2023-12-05 12:10:03 -08:00
< script > DOCUMENTATION _OPTIONS . pagename = 'examples/llama-inference' ; < / script >
2023-11-29 12:41:56 -08:00
< link rel = "index" title = "Index" href = "../genindex.html" / >
< link rel = "search" title = "Search" href = "../search.html" / >
2023-12-05 12:10:03 -08:00
< link rel = "next" title = "Array" href = "../python/array.html" / >
< link rel = "prev" title = "Multi-Layer Perceptron" href = "mlp.html" / >
< meta name = "viewport" content = "width=device-width, initial-scale=1" / >
< meta name = "docsearch:language" content = "en" / >
< / head >
< body data-bs-spy = "scroll" data-bs-target = ".bd-toc-nav" data-offset = "180" data-bs-root-margin = "0px 0px -60%" data-default-mode = "" >
< a class = "skip-link" href = "#main-content" > Skip to main content< / a >
< div id = "pst-scroll-pixel-helper" > < / div >
< button type = "button" class = "btn rounded-pill" id = "pst-back-to-top" >
< i class = "fa-solid fa-arrow-up" > < / i >
Back to top
< / button >
< input type = "checkbox"
class="sidebar-toggle"
name="__primary"
id="__primary"/>
< label class = "overlay overlay-primary" for = "__primary" > < / label >
< input type = "checkbox"
class="sidebar-toggle"
name="__secondary"
id="__secondary"/>
< label class = "overlay overlay-secondary" for = "__secondary" > < / label >
< div class = "search-button__wrapper" >
< div class = "search-button__overlay" > < / div >
< div class = "search-button__search-container" >
< form class = "bd-search d-flex align-items-center"
action="../search.html"
method="get">
< i class = "fa-solid fa-magnifying-glass" > < / i >
< input type = "search"
class="form-control"
name="q"
id="search-input"
placeholder="Search..."
aria-label="Search..."
autocomplete="off"
autocorrect="off"
autocapitalize="off"
spellcheck="false"/>
< span class = "search-button__kbd-shortcut" > < kbd class = "kbd-shortcut__modifier" > Ctrl< / kbd > +< kbd > K< / kbd > < / span >
< / form > < / div >
< / div >
< nav class = "bd-header navbar navbar-expand-lg bd-navbar" >
< / nav >
< div class = "bd-container" >
< div class = "bd-container__inner bd-page-width" >
< div class = "bd-sidebar-primary bd-sidebar" >
< div class = "sidebar-header-items sidebar-primary__section" >
< / div >
< div class = "sidebar-primary-items__start sidebar-primary__section" >
< div class = "sidebar-primary-item" >
< a class = "navbar-brand logo" href = "../index.html" >
2024-04-26 08:24:09 -07:00
< img src = "../_static/mlx_logo.png" class = "logo__image only-light" alt = "MLX 0.12.0 documentation - Home" / >
< script > document . write ( ` <img src="../_static/mlx_logo_dark.png" class="logo__image only-dark" alt="MLX 0.12.0 documentation - Home"/> ` ) ; < / script >
2023-12-05 12:10:03 -08:00
< / a > < / div >
2024-03-30 17:32:20 -07:00
< div class = "sidebar-primary-item" >
< script >
document.write(`
< button class = "btn navbar-btn search-button-field search-button__button" title = "Search" aria-label = "Search" data-bs-placement = "bottom" data-bs-toggle = "tooltip" >
< i class = "fa-solid fa-magnifying-glass" > < / i >
< span class = "search-button__default-text" > Search< / span >
< span class = "search-button__kbd-shortcut" > < kbd class = "kbd-shortcut__modifier" > Ctrl< / kbd > +< kbd class = "kbd-shortcut__modifier" > K< / kbd > < / span >
< / button >
`);
< / script > < / div >
< div class = "sidebar-primary-item" > < nav class = "bd-links bd-docs-nav" aria-label = "Main" >
2023-12-05 12:10:03 -08:00
< div class = "bd-toc-item navbar-nav active" >
< p aria-level = "2" class = "caption" role = "heading" > < span class = "caption-text" > Install< / span > < / p >
< ul class = "nav bd-sidenav" >
2023-11-29 12:41:56 -08:00
< li class = "toctree-l1" > < a class = "reference internal" href = "../install.html" > Build and Install< / a > < / li >
< / ul >
2023-12-05 12:10:03 -08:00
< p aria-level = "2" class = "caption" role = "heading" > < span class = "caption-text" > Usage< / span > < / p >
< ul class = "nav bd-sidenav" >
2024-01-06 05:41:48 -08:00
< li class = "toctree-l1" > < a class = "reference internal" href = "../usage/quick_start.html" > Quick Start Guide< / a > < / li >
2024-01-10 14:14:38 -08:00
< li class = "toctree-l1" > < a class = "reference internal" href = "../usage/lazy_evaluation.html" > Lazy Evaluation< / a > < / li >
2024-01-06 05:41:48 -08:00
< li class = "toctree-l1" > < a class = "reference internal" href = "../usage/unified_memory.html" > Unified Memory< / a > < / li >
2024-01-10 14:14:38 -08:00
< li class = "toctree-l1" > < a class = "reference internal" href = "../usage/indexing.html" > Indexing Arrays< / a > < / li >
< li class = "toctree-l1" > < a class = "reference internal" href = "../usage/saving_and_loading.html" > Saving and Loading Arrays< / a > < / li >
2024-01-17 17:15:29 -08:00
< li class = "toctree-l1" > < a class = "reference internal" href = "../usage/function_transforms.html" > Function Transforms< / a > < / li >
2024-02-08 12:44:23 -08:00
< li class = "toctree-l1" > < a class = "reference internal" href = "../usage/compile.html" > Compilation< / a > < / li >
2024-01-06 05:41:48 -08:00
< li class = "toctree-l1" > < a class = "reference internal" href = "../usage/numpy.html" > Conversion to NumPy and Other Frameworks< / a > < / li >
2024-01-10 14:14:38 -08:00
< li class = "toctree-l1" > < a class = "reference internal" href = "../usage/using_streams.html" > Using Streams< / a > < / li >
2023-11-29 12:41:56 -08:00
< / ul >
2023-12-05 12:10:03 -08:00
< p aria-level = "2" class = "caption" role = "heading" > < span class = "caption-text" > Examples< / span > < / p >
< ul class = "current nav bd-sidenav" >
2023-11-29 12:41:56 -08:00
< li class = "toctree-l1" > < a class = "reference internal" href = "linear_regression.html" > Linear Regression< / a > < / li >
< li class = "toctree-l1" > < a class = "reference internal" href = "mlp.html" > Multi-Layer Perceptron< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l1 current active" > < a class = "current reference internal" href = "#" > LLM inference< / a > < / li >
< / ul >
< p aria-level = "2" class = "caption" role = "heading" > < span class = "caption-text" > Python API Reference< / span > < / p >
< ul class = "nav bd-sidenav" >
< li class = "toctree-l1 has-children" > < a class = "reference internal" href = "../python/array.html" > Array< / a > < input class = "toctree-checkbox" id = "toctree-checkbox-1" name = "toctree-checkbox-1" type = "checkbox" / > < label class = "toctree-toggle" for = "toctree-checkbox-1" > < i class = "fa-solid fa-chevron-down" > < / i > < / label > < ul >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.html" > mlx.core.array< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.astype.html" > mlx.core.array.astype< / a > < / li >
2024-03-30 17:32:20 -07:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.at.html" > mlx.core.array.at< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.item.html" > mlx.core.array.item< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.tolist.html" > mlx.core.array.tolist< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.dtype.html" > mlx.core.array.dtype< / a > < / li >
2024-03-30 17:32:20 -07:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.itemsize.html" > mlx.core.array.itemsize< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.nbytes.html" > mlx.core.array.nbytes< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.ndim.html" > mlx.core.array.ndim< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.shape.html" > mlx.core.array.shape< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.size.html" > mlx.core.array.size< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.abs.html" > mlx.core.array.abs< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.all.html" > mlx.core.array.all< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.any.html" > mlx.core.array.any< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.argmax.html" > mlx.core.array.argmax< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.argmin.html" > mlx.core.array.argmin< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.cos.html" > mlx.core.array.cos< / a > < / li >
2024-03-30 17:32:20 -07:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.cummax.html" > mlx.core.array.cummax< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.cummin.html" > mlx.core.array.cummin< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.cumprod.html" > mlx.core.array.cumprod< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.cumsum.html" > mlx.core.array.cumsum< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.diag.html" > mlx.core.array.diag< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.diagonal.html" > mlx.core.array.diagonal< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.exp.html" > mlx.core.array.exp< / a > < / li >
2024-03-30 17:32:20 -07:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.flatten.html" > mlx.core.array.flatten< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.log.html" > mlx.core.array.log< / a > < / li >
2024-03-30 17:32:20 -07:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.log10.html" > mlx.core.array.log10< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.log1p.html" > mlx.core.array.log1p< / a > < / li >
2024-03-30 17:32:20 -07:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.log2.html" > mlx.core.array.log2< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.logsumexp.html" > mlx.core.array.logsumexp< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.max.html" > mlx.core.array.max< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.mean.html" > mlx.core.array.mean< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.min.html" > mlx.core.array.min< / a > < / li >
2024-03-30 17:32:20 -07:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.moveaxis.html" > mlx.core.array.moveaxis< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.prod.html" > mlx.core.array.prod< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.reciprocal.html" > mlx.core.array.reciprocal< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.reshape.html" > mlx.core.array.reshape< / a > < / li >
2023-12-21 22:13:41 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.round.html" > mlx.core.array.round< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.rsqrt.html" > mlx.core.array.rsqrt< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.sin.html" > mlx.core.array.sin< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.split.html" > mlx.core.array.split< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.sqrt.html" > mlx.core.array.sqrt< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.square.html" > mlx.core.array.square< / a > < / li >
2024-03-30 17:32:20 -07:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.squeeze.html" > mlx.core.array.squeeze< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.swapaxes.html" > mlx.core.array.swapaxes< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.sum.html" > mlx.core.array.sum< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.transpose.html" > mlx.core.array.transpose< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.T.html" > mlx.core.array.T< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array.var.html" > mlx.core.array.var< / a > < / li >
2023-11-29 12:41:56 -08:00
< / ul >
< / li >
2024-03-30 17:32:20 -07:00
< li class = "toctree-l1 has-children" > < a class = "reference internal" href = "../python/data_types.html" > Data Types< / a > < input class = "toctree-checkbox" id = "toctree-checkbox-2" name = "toctree-checkbox-2" type = "checkbox" / > < label class = "toctree-toggle" for = "toctree-checkbox-2" > < i class = "fa-solid fa-chevron-down" > < / i > < / label > < ul >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.Dtype.html" > mlx.core.Dtype< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.DtypeCategory.html" > mlx.core.DtypeCategory< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.issubdtype.html" > mlx.core.issubdtype< / a > < / li >
< / ul >
< / li >
< li class = "toctree-l1 has-children" > < a class = "reference internal" href = "../python/devices_and_streams.html" > Devices and Streams< / a > < input class = "toctree-checkbox" id = "toctree-checkbox-3" name = "toctree-checkbox-3" type = "checkbox" / > < label class = "toctree-toggle" for = "toctree-checkbox-3" > < i class = "fa-solid fa-chevron-down" > < / i > < / label > < ul >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.Device.html" > mlx.core.Device< / a > < / li >
2024-02-17 13:25:37 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/stream_class.html" > mlx.core.Stream< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.default_device.html" > mlx.core.default_device< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.set_default_device.html" > mlx.core.set_default_device< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.default_stream.html" > mlx.core.default_stream< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.new_stream.html" > mlx.core.new_stream< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.set_default_stream.html" > mlx.core.set_default_stream< / a > < / li >
2024-02-17 13:25:37 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.stream.html" > mlx.core.stream< / a > < / li >
2024-04-26 08:24:09 -07:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.synchronize.html" > mlx.core.synchronize< / a > < / li >
2023-11-29 12:41:56 -08:00
< / ul >
< / li >
2024-03-30 17:32:20 -07:00
< li class = "toctree-l1 has-children" > < a class = "reference internal" href = "../python/ops.html" > Operations< / a > < input class = "toctree-checkbox" id = "toctree-checkbox-4" name = "toctree-checkbox-4" type = "checkbox" / > < label class = "toctree-toggle" for = "toctree-checkbox-4" > < i class = "fa-solid fa-chevron-down" > < / i > < / label > < ul >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.abs.html" > mlx.core.abs< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.add.html" > mlx.core.add< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.all.html" > mlx.core.all< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.allclose.html" > mlx.core.allclose< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.any.html" > mlx.core.any< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.arange.html" > mlx.core.arange< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.arccos.html" > mlx.core.arccos< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.arccosh.html" > mlx.core.arccosh< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.arcsin.html" > mlx.core.arcsin< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.arcsinh.html" > mlx.core.arcsinh< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.arctan.html" > mlx.core.arctan< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.arctanh.html" > mlx.core.arctanh< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.argmax.html" > mlx.core.argmax< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.argmin.html" > mlx.core.argmin< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.argpartition.html" > mlx.core.argpartition< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.argsort.html" > mlx.core.argsort< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.array_equal.html" > mlx.core.array_equal< / a > < / li >
2024-02-29 12:39:18 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.atleast_1d.html" > mlx.core.atleast_1d< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.atleast_2d.html" > mlx.core.atleast_2d< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.atleast_3d.html" > mlx.core.atleast_3d< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.broadcast_to.html" > mlx.core.broadcast_to< / a > < / li >
2024-04-26 08:24:09 -07:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.block_masked_mm.html" > mlx.core.block_masked_mm< / a > < / li >
2023-12-17 13:23:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.ceil.html" > mlx.core.ceil< / a > < / li >
2023-12-21 22:13:41 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.clip.html" > mlx.core.clip< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.concatenate.html" > mlx.core.concatenate< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.convolve.html" > mlx.core.convolve< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.conv1d.html" > mlx.core.conv1d< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.conv2d.html" > mlx.core.conv2d< / a > < / li >
2024-02-29 12:39:18 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.conv_general.html" > mlx.core.conv_general< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.cos.html" > mlx.core.cos< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.cosh.html" > mlx.core.cosh< / a > < / li >
2024-03-30 17:32:20 -07:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.cummax.html" > mlx.core.cummax< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.cummin.html" > mlx.core.cummin< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.cumprod.html" > mlx.core.cumprod< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.cumsum.html" > mlx.core.cumsum< / a > < / li >
2024-04-26 08:24:09 -07:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.degrees.html" > mlx.core.degrees< / a > < / li >
2023-12-21 22:13:41 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.dequantize.html" > mlx.core.dequantize< / a > < / li >
2024-02-01 13:08:29 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.diag.html" > mlx.core.diag< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.diagonal.html" > mlx.core.diagonal< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.divide.html" > mlx.core.divide< / a > < / li >
2024-01-10 14:14:38 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.divmod.html" > mlx.core.divmod< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.equal.html" > mlx.core.equal< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.erf.html" > mlx.core.erf< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.erfinv.html" > mlx.core.erfinv< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.exp.html" > mlx.core.exp< / a > < / li >
2024-04-11 17:33:33 -07:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.expm1.html" > mlx.core.expm1< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.expand_dims.html" > mlx.core.expand_dims< / a > < / li >
2023-12-13 14:46:24 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.eye.html" > mlx.core.eye< / a > < / li >
2023-12-17 13:23:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.flatten.html" > mlx.core.flatten< / a > < / li >
2023-12-21 22:13:41 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.floor.html" > mlx.core.floor< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.floor_divide.html" > mlx.core.floor_divide< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.full.html" > mlx.core.full< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.greater.html" > mlx.core.greater< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.greater_equal.html" > mlx.core.greater_equal< / a > < / li >
2023-12-13 14:46:24 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.identity.html" > mlx.core.identity< / a > < / li >
2024-01-10 14:14:38 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.inner.html" > mlx.core.inner< / a > < / li >
2024-03-14 12:46:45 -07:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.isclose.html" > mlx.core.isclose< / a > < / li >
2024-03-30 17:32:20 -07:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.isinf.html" > mlx.core.isinf< / a > < / li >
2024-01-17 17:15:29 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.isnan.html" > mlx.core.isnan< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.isneginf.html" > mlx.core.isneginf< / a > < / li >
2024-03-30 17:32:20 -07:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.isposinf.html" > mlx.core.isposinf< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.less.html" > mlx.core.less< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.less_equal.html" > mlx.core.less_equal< / a > < / li >
2023-12-21 22:13:41 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.linspace.html" > mlx.core.linspace< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.load.html" > mlx.core.load< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.log.html" > mlx.core.log< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.log2.html" > mlx.core.log2< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.log10.html" > mlx.core.log10< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.log1p.html" > mlx.core.log1p< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.logaddexp.html" > mlx.core.logaddexp< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.logical_not.html" > mlx.core.logical_not< / a > < / li >
2024-01-10 14:14:38 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.logical_and.html" > mlx.core.logical_and< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.logical_or.html" > mlx.core.logical_or< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.logsumexp.html" > mlx.core.logsumexp< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.matmul.html" > mlx.core.matmul< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.max.html" > mlx.core.max< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.maximum.html" > mlx.core.maximum< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.mean.html" > mlx.core.mean< / a > < / li >
2024-04-11 17:33:33 -07:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.meshgrid.html" > mlx.core.meshgrid< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.min.html" > mlx.core.min< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.minimum.html" > mlx.core.minimum< / a > < / li >
2023-12-17 13:23:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.moveaxis.html" > mlx.core.moveaxis< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.multiply.html" > mlx.core.multiply< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.negative.html" > mlx.core.negative< / a > < / li >
2024-04-26 08:24:09 -07:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.not_equal.html" > mlx.core.not_equal< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.ones.html" > mlx.core.ones< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.ones_like.html" > mlx.core.ones_like< / a > < / li >
2024-01-10 14:14:38 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.outer.html" > mlx.core.outer< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.partition.html" > mlx.core.partition< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.pad.html" > mlx.core.pad< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.prod.html" > mlx.core.prod< / a > < / li >
2023-12-21 22:13:41 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.quantize.html" > mlx.core.quantize< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.quantized_matmul.html" > mlx.core.quantized_matmul< / a > < / li >
2024-04-26 08:24:09 -07:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.radians.html" > mlx.core.radians< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.reciprocal.html" > mlx.core.reciprocal< / a > < / li >
2024-01-03 20:14:05 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.repeat.html" > mlx.core.repeat< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.reshape.html" > mlx.core.reshape< / a > < / li >
2023-12-21 22:13:41 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.round.html" > mlx.core.round< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.rsqrt.html" > mlx.core.rsqrt< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.save.html" > mlx.core.save< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.savez.html" > mlx.core.savez< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.savez_compressed.html" > mlx.core.savez_compressed< / a > < / li >
2024-01-10 14:14:38 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.save_gguf.html" > mlx.core.save_gguf< / a > < / li >
2024-01-03 20:14:05 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.save_safetensors.html" > mlx.core.save_safetensors< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.sigmoid.html" > mlx.core.sigmoid< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.sign.html" > mlx.core.sign< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.sin.html" > mlx.core.sin< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.sinh.html" > mlx.core.sinh< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.softmax.html" > mlx.core.softmax< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.sort.html" > mlx.core.sort< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.split.html" > mlx.core.split< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.sqrt.html" > mlx.core.sqrt< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.square.html" > mlx.core.square< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.squeeze.html" > mlx.core.squeeze< / a > < / li >
2023-12-17 13:23:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.stack.html" > mlx.core.stack< / a > < / li >
2024-04-11 17:33:33 -07:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.std.html" > mlx.core.std< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.stop_gradient.html" > mlx.core.stop_gradient< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.subtract.html" > mlx.core.subtract< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.sum.html" > mlx.core.sum< / a > < / li >
2023-12-17 13:23:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.swapaxes.html" > mlx.core.swapaxes< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.take.html" > mlx.core.take< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.take_along_axis.html" > mlx.core.take_along_axis< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.tan.html" > mlx.core.tan< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.tanh.html" > mlx.core.tanh< / a > < / li >
2024-01-03 20:14:05 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.tensordot.html" > mlx.core.tensordot< / a > < / li >
2024-03-14 12:46:45 -07:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.tile.html" > mlx.core.tile< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.topk.html" > mlx.core.topk< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.transpose.html" > mlx.core.transpose< / a > < / li >
2023-12-17 13:23:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.tri.html" > mlx.core.tri< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.tril.html" > mlx.core.tril< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.triu.html" > mlx.core.triu< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.var.html" > mlx.core.var< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.where.html" > mlx.core.where< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.zeros.html" > mlx.core.zeros< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.zeros_like.html" > mlx.core.zeros_like< / a > < / li >
2023-11-29 12:41:56 -08:00
< / ul >
2023-12-05 12:10:03 -08:00
< / li >
2024-03-30 17:32:20 -07:00
< li class = "toctree-l1 has-children" > < a class = "reference internal" href = "../python/random.html" > Random< / a > < input class = "toctree-checkbox" id = "toctree-checkbox-5" name = "toctree-checkbox-5" type = "checkbox" / > < label class = "toctree-toggle" for = "toctree-checkbox-5" > < i class = "fa-solid fa-chevron-down" > < / i > < / label > < ul >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.random.bernoulli.html" > mlx.core.random.bernoulli< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.random.categorical.html" > mlx.core.random.categorical< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.random.gumbel.html" > mlx.core.random.gumbel< / a > < / li >
2024-01-17 17:15:29 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.random.key.html" > mlx.core.random.key< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.random.normal.html" > mlx.core.random.normal< / a > < / li >
2024-04-11 17:33:33 -07:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.random.multivariate_normal.html" > mlx.core.random.multivariate_normal< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.random.randint.html" > mlx.core.random.randint< / a > < / li >
2024-01-17 17:15:29 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.random.seed.html" > mlx.core.random.seed< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.random.split.html" > mlx.core.random.split< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.random.truncated_normal.html" > mlx.core.random.truncated_normal< / a > < / li >
2024-01-17 17:15:29 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.random.uniform.html" > mlx.core.random.uniform< / a > < / li >
2023-12-05 12:10:03 -08:00
< / ul >
< / li >
2024-03-30 17:32:20 -07:00
< li class = "toctree-l1 has-children" > < a class = "reference internal" href = "../python/transforms.html" > Transforms< / a > < input class = "toctree-checkbox" id = "toctree-checkbox-6" name = "toctree-checkbox-6" type = "checkbox" / > < label class = "toctree-toggle" for = "toctree-checkbox-6" > < i class = "fa-solid fa-chevron-down" > < / i > < / label > < ul >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.eval.html" > mlx.core.eval< / a > < / li >
2024-02-08 12:44:23 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.compile.html" > mlx.core.compile< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.disable_compile.html" > mlx.core.disable_compile< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.enable_compile.html" > mlx.core.enable_compile< / a > < / li >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.grad.html" > mlx.core.grad< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.value_and_grad.html" > mlx.core.value_and_grad< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.jvp.html" > mlx.core.jvp< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.vjp.html" > mlx.core.vjp< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.vmap.html" > mlx.core.vmap< / a > < / li >
< / ul >
< / li >
2024-03-30 17:32:20 -07:00
< li class = "toctree-l1 has-children" > < a class = "reference internal" href = "../python/fast.html" > Fast< / a > < input class = "toctree-checkbox" id = "toctree-checkbox-7" name = "toctree-checkbox-7" type = "checkbox" / > < label class = "toctree-toggle" for = "toctree-checkbox-7" > < i class = "fa-solid fa-chevron-down" > < / i > < / label > < ul >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.fast.rms_norm.html" > mlx.core.fast.rms_norm< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.fast.layer_norm.html" > mlx.core.fast.layer_norm< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.fast.rope.html" > mlx.core.fast.rope< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.fast.scaled_dot_product_attention.html" > mlx.core.fast.scaled_dot_product_attention< / a > < / li >
< / ul >
< / li >
< li class = "toctree-l1 has-children" > < a class = "reference internal" href = "../python/fft.html" > FFT< / a > < input class = "toctree-checkbox" id = "toctree-checkbox-8" name = "toctree-checkbox-8" type = "checkbox" / > < label class = "toctree-toggle" for = "toctree-checkbox-8" > < i class = "fa-solid fa-chevron-down" > < / i > < / label > < ul >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.fft.fft.html" > mlx.core.fft.fft< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.fft.ifft.html" > mlx.core.fft.ifft< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.fft.fft2.html" > mlx.core.fft.fft2< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.fft.ifft2.html" > mlx.core.fft.ifft2< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.fft.fftn.html" > mlx.core.fft.fftn< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.fft.ifftn.html" > mlx.core.fft.ifftn< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.fft.rfft.html" > mlx.core.fft.rfft< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.fft.irfft.html" > mlx.core.fft.irfft< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.fft.rfft2.html" > mlx.core.fft.rfft2< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.fft.irfft2.html" > mlx.core.fft.irfft2< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.fft.rfftn.html" > mlx.core.fft.rfftn< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.fft.irfftn.html" > mlx.core.fft.irfftn< / a > < / li >
< / ul >
< / li >
2024-03-30 17:32:20 -07:00
< li class = "toctree-l1 has-children" > < a class = "reference internal" href = "../python/linalg.html" > Linear Algebra< / a > < input class = "toctree-checkbox" id = "toctree-checkbox-9" name = "toctree-checkbox-9" type = "checkbox" / > < label class = "toctree-toggle" for = "toctree-checkbox-9" > < i class = "fa-solid fa-chevron-down" > < / i > < / label > < ul >
2024-01-03 20:14:05 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.linalg.norm.html" > mlx.core.linalg.norm< / a > < / li >
2024-02-01 13:08:29 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.linalg.qr.html" > mlx.core.linalg.qr< / a > < / li >
2024-01-03 20:14:05 -08:00
< / ul >
< / li >
2024-03-30 17:32:20 -07:00
< li class = "toctree-l1 has-children" > < a class = "reference internal" href = "../python/metal.html" > Metal< / a > < input class = "toctree-checkbox" id = "toctree-checkbox-10" name = "toctree-checkbox-10" type = "checkbox" / > < label class = "toctree-toggle" for = "toctree-checkbox-10" > < i class = "fa-solid fa-chevron-down" > < / i > < / label > < ul >
2024-03-14 12:46:45 -07:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.metal.is_available.html" > mlx.core.metal.is_available< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.metal.get_active_memory.html" > mlx.core.metal.get_active_memory< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.metal.get_peak_memory.html" > mlx.core.metal.get_peak_memory< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.metal.get_cache_memory.html" > mlx.core.metal.get_cache_memory< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.metal.set_memory_limit.html" > mlx.core.metal.set_memory_limit< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.metal.set_cache_limit.html" > mlx.core.metal.set_cache_limit< / a > < / li >
2024-04-26 08:24:09 -07:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.metal.clear_cache.html" > mlx.core.metal.clear_cache< / a > < / li >
2024-04-11 17:33:33 -07:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.metal.start_capture.html" > mlx.core.metal.start_capture< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.core.metal.stop_capture.html" > mlx.core.metal.stop_capture< / a > < / li >
2024-03-14 12:46:45 -07:00
< / ul >
< / li >
2024-03-30 17:32:20 -07:00
< li class = "toctree-l1 has-children" > < a class = "reference internal" href = "../python/nn.html" > Neural Networks< / a > < input class = "toctree-checkbox" id = "toctree-checkbox-11" name = "toctree-checkbox-11" type = "checkbox" / > < label class = "toctree-toggle" for = "toctree-checkbox-11" > < i class = "fa-solid fa-chevron-down" > < / i > < / label > < ul >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.nn.value_and_grad.html" > mlx.nn.value_and_grad< / a > < / li >
2024-04-26 08:24:09 -07:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.nn.quantize.html" > mlx.nn.quantize< / a > < / li >
2024-03-30 17:32:20 -07:00
< li class = "toctree-l2 has-children" > < a class = "reference internal" href = "../python/nn/module.html" > Module< / a > < input class = "toctree-checkbox" id = "toctree-checkbox-12" name = "toctree-checkbox-12" type = "checkbox" / > < label class = "toctree-toggle" for = "toctree-checkbox-12" > < i class = "fa-solid fa-chevron-down" > < / i > < / label > < ul >
2024-01-03 20:14:05 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Module.training.html" > mlx.nn.Module.training< / a > < / li >
2024-02-08 12:44:23 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Module.state.html" > mlx.nn.Module.state< / a > < / li >
2024-01-03 20:14:05 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Module.apply.html" > mlx.nn.Module.apply< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Module.apply_to_modules.html" > mlx.nn.Module.apply_to_modules< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Module.children.html" > mlx.nn.Module.children< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Module.eval.html" > mlx.nn.Module.eval< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Module.filter_and_map.html" > mlx.nn.Module.filter_and_map< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Module.freeze.html" > mlx.nn.Module.freeze< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Module.leaf_modules.html" > mlx.nn.Module.leaf_modules< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Module.load_weights.html" > mlx.nn.Module.load_weights< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Module.modules.html" > mlx.nn.Module.modules< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Module.named_modules.html" > mlx.nn.Module.named_modules< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Module.parameters.html" > mlx.nn.Module.parameters< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Module.save_weights.html" > mlx.nn.Module.save_weights< / a > < / li >
2024-03-30 17:32:20 -07:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Module.set_dtype.html" > mlx.nn.Module.set_dtype< / a > < / li >
2024-01-03 20:14:05 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Module.train.html" > mlx.nn.Module.train< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Module.trainable_parameters.html" > mlx.nn.Module.trainable_parameters< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Module.unfreeze.html" > mlx.nn.Module.unfreeze< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Module.update.html" > mlx.nn.Module.update< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Module.update_modules.html" > mlx.nn.Module.update_modules< / a > < / li >
< / ul >
< / li >
2024-03-30 17:32:20 -07:00
< li class = "toctree-l2 has-children" > < a class = "reference internal" href = "../python/nn/layers.html" > Layers< / a > < input class = "toctree-checkbox" id = "toctree-checkbox-13" name = "toctree-checkbox-13" type = "checkbox" / > < label class = "toctree-toggle" for = "toctree-checkbox-13" > < i class = "fa-solid fa-chevron-down" > < / i > < / label > < ul >
2024-01-17 17:15:29 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.ALiBi.html" > mlx.nn.ALiBi< / a > < / li >
2024-02-17 13:25:37 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.AvgPool1d.html" > mlx.nn.AvgPool1d< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.AvgPool2d.html" > mlx.nn.AvgPool2d< / a > < / li >
2024-01-17 17:15:29 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.BatchNorm.html" > mlx.nn.BatchNorm< / a > < / li >
2023-12-17 13:23:03 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Conv1d.html" > mlx.nn.Conv1d< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Conv2d.html" > mlx.nn.Conv2d< / a > < / li >
2024-01-03 20:14:05 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Dropout.html" > mlx.nn.Dropout< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Dropout2d.html" > mlx.nn.Dropout2d< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Dropout3d.html" > mlx.nn.Dropout3d< / a > < / li >
2024-01-17 17:15:29 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Embedding.html" > mlx.nn.Embedding< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.GELU.html" > mlx.nn.GELU< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.GroupNorm.html" > mlx.nn.GroupNorm< / a > < / li >
2024-03-14 12:46:45 -07:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.GRU.html" > mlx.nn.GRU< / a > < / li >
2024-01-17 17:15:29 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.InstanceNorm.html" > mlx.nn.InstanceNorm< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.LayerNorm.html" > mlx.nn.LayerNorm< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Linear.html" > mlx.nn.Linear< / a > < / li >
2024-03-14 12:46:45 -07:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.LSTM.html" > mlx.nn.LSTM< / a > < / li >
2024-02-17 13:25:37 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.MaxPool1d.html" > mlx.nn.MaxPool1d< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.MaxPool2d.html" > mlx.nn.MaxPool2d< / a > < / li >
2024-01-17 17:15:29 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Mish.html" > mlx.nn.Mish< / a > < / li >
2023-12-17 13:23:03 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.MultiHeadAttention.html" > mlx.nn.MultiHeadAttention< / a > < / li >
2024-01-17 17:15:29 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.PReLU.html" > mlx.nn.PReLU< / a > < / li >
2024-04-26 08:24:09 -07:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.QuantizedEmbedding.html" > mlx.nn.QuantizedEmbedding< / a > < / li >
2024-01-17 17:15:29 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.QuantizedLinear.html" > mlx.nn.QuantizedLinear< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.RMSNorm.html" > mlx.nn.RMSNorm< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.ReLU.html" > mlx.nn.ReLU< / a > < / li >
2024-03-14 12:46:45 -07:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.RNN.html" > mlx.nn.RNN< / a > < / li >
2024-01-03 20:14:05 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.RoPE.html" > mlx.nn.RoPE< / a > < / li >
2024-01-17 17:15:29 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.SELU.html" > mlx.nn.SELU< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Sequential.html" > mlx.nn.Sequential< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.SiLU.html" > mlx.nn.SiLU< / a > < / li >
2024-01-03 20:14:05 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.SinusoidalPositionalEncoding.html" > mlx.nn.SinusoidalPositionalEncoding< / a > < / li >
2024-02-01 13:08:29 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Softshrink.html" > mlx.nn.Softshrink< / a > < / li >
2024-01-17 17:15:29 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Step.html" > mlx.nn.Step< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Transformer.html" > mlx.nn.Transformer< / a > < / li >
2024-02-29 12:39:18 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Upsample.html" > mlx.nn.Upsample< / a > < / li >
2023-11-29 12:41:56 -08:00
< / ul >
2023-12-05 12:10:03 -08:00
< / li >
2024-03-30 17:32:20 -07:00
< li class = "toctree-l2 has-children" > < a class = "reference internal" href = "../python/nn/functions.html" > Functions< / a > < input class = "toctree-checkbox" id = "toctree-checkbox-14" name = "toctree-checkbox-14" type = "checkbox" / > < label class = "toctree-toggle" for = "toctree-checkbox-14" > < i class = "fa-solid fa-chevron-down" > < / i > < / label > < ul >
2024-02-29 12:39:18 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.elu.html" > mlx.nn.elu< / a > < / li >
2023-12-17 13:23:03 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.gelu.html" > mlx.nn.gelu< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.gelu_approx.html" > mlx.nn.gelu_approx< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.gelu_fast_approx.html" > mlx.nn.gelu_fast_approx< / a > < / li >
2024-02-29 12:39:18 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.glu.html" > mlx.nn.glu< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.hardswish.html" > mlx.nn.hardswish< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.leaky_relu.html" > mlx.nn.leaky_relu< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.log_sigmoid.html" > mlx.nn.log_sigmoid< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.log_softmax.html" > mlx.nn.log_softmax< / a > < / li >
2024-01-17 17:15:29 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.mish.html" > mlx.nn.mish< / a > < / li >
2023-12-17 13:23:03 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.prelu.html" > mlx.nn.prelu< / a > < / li >
2024-01-17 17:15:29 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.relu.html" > mlx.nn.relu< / a > < / li >
2024-02-29 12:39:18 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.relu6.html" > mlx.nn.relu6< / a > < / li >
2024-01-17 17:15:29 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.selu.html" > mlx.nn.selu< / a > < / li >
2024-02-29 12:39:18 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.sigmoid.html" > mlx.nn.sigmoid< / a > < / li >
2023-12-17 13:23:03 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.silu.html" > mlx.nn.silu< / a > < / li >
2024-02-29 12:39:18 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.softmax.html" > mlx.nn.softmax< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.softplus.html" > mlx.nn.softplus< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.softshrink.html" > mlx.nn.softshrink< / a > < / li >
2023-12-17 13:23:03 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.step.html" > mlx.nn.step< / a > < / li >
2024-02-29 12:39:18 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.tanh.html" > mlx.nn.tanh< / a > < / li >
2023-12-17 13:23:03 -08:00
< / ul >
< / li >
2024-03-30 17:32:20 -07:00
< li class = "toctree-l2 has-children" > < a class = "reference internal" href = "../python/nn/losses.html" > Loss Functions< / a > < input class = "toctree-checkbox" id = "toctree-checkbox-15" name = "toctree-checkbox-15" type = "checkbox" / > < label class = "toctree-toggle" for = "toctree-checkbox-15" > < i class = "fa-solid fa-chevron-down" > < / i > < / label > < ul >
2023-12-17 13:23:03 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.losses.binary_cross_entropy.html" > mlx.nn.losses.binary_cross_entropy< / a > < / li >
2024-01-17 17:15:29 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.losses.cosine_similarity_loss.html" > mlx.nn.losses.cosine_similarity_loss< / a > < / li >
2023-12-21 22:13:41 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.losses.cross_entropy.html" > mlx.nn.losses.cross_entropy< / a > < / li >
2024-02-01 13:08:29 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.losses.gaussian_nll_loss.html" > mlx.nn.losses.gaussian_nll_loss< / a > < / li >
2024-01-17 17:15:29 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.losses.hinge_loss.html" > mlx.nn.losses.hinge_loss< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.losses.huber_loss.html" > mlx.nn.losses.huber_loss< / a > < / li >
2023-12-21 22:13:41 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.losses.kl_div_loss.html" > mlx.nn.losses.kl_div_loss< / a > < / li >
2023-12-17 13:23:03 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.losses.l1_loss.html" > mlx.nn.losses.l1_loss< / a > < / li >
2024-01-17 17:15:29 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.losses.log_cosh_loss.html" > mlx.nn.losses.log_cosh_loss< / a > < / li >
2024-02-08 12:44:23 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.losses.margin_ranking_loss.html" > mlx.nn.losses.margin_ranking_loss< / a > < / li >
2023-12-17 13:23:03 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.losses.mse_loss.html" > mlx.nn.losses.mse_loss< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.losses.nll_loss.html" > mlx.nn.losses.nll_loss< / a > < / li >
2023-12-21 22:13:41 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.losses.smooth_l1_loss.html" > mlx.nn.losses.smooth_l1_loss< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary_functions/mlx.nn.losses.triplet_loss.html" > mlx.nn.losses.triplet_loss< / a > < / li >
2023-12-17 13:23:03 -08:00
< / ul >
< / li >
2024-03-30 17:32:20 -07:00
< li class = "toctree-l2 has-children" > < a class = "reference internal" href = "../python/nn/init.html" > Initializers< / a > < input class = "toctree-checkbox" id = "toctree-checkbox-16" name = "toctree-checkbox-16" type = "checkbox" / > < label class = "toctree-toggle" for = "toctree-checkbox-16" > < i class = "fa-solid fa-chevron-down" > < / i > < / label > < ul >
2024-02-01 13:08:29 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.init.constant.html" > mlx.nn.init.constant< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.init.normal.html" > mlx.nn.init.normal< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.init.uniform.html" > mlx.nn.init.uniform< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.init.identity.html" > mlx.nn.init.identity< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.init.glorot_normal.html" > mlx.nn.init.glorot_normal< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.init.glorot_uniform.html" > mlx.nn.init.glorot_uniform< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.init.he_normal.html" > mlx.nn.init.he_normal< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.init.he_uniform.html" > mlx.nn.init.he_uniform< / a > < / li >
2023-12-17 13:23:03 -08:00
< / ul >
< / li >
2024-02-01 13:08:29 -08:00
< / ul >
< / li >
2024-03-30 17:32:20 -07:00
< li class = "toctree-l1 has-children" > < a class = "reference internal" href = "../python/optimizers.html" > Optimizers< / a > < input class = "toctree-checkbox" id = "toctree-checkbox-17" name = "toctree-checkbox-17" type = "checkbox" / > < label class = "toctree-toggle" for = "toctree-checkbox-17" > < i class = "fa-solid fa-chevron-down" > < / i > < / label > < ul >
< li class = "toctree-l2 has-children" > < a class = "reference internal" href = "../python/optimizers/optimizer.html" > Optimizer< / a > < input class = "toctree-checkbox" id = "toctree-checkbox-18" name = "toctree-checkbox-18" type = "checkbox" / > < label class = "toctree-toggle" for = "toctree-checkbox-18" > < i class = "fa-solid fa-chevron-down" > < / i > < / label > < ul >
2024-02-17 13:25:37 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/optimizers/_autosummary/mlx.optimizers.Optimizer.state.html" > mlx.optimizers.Optimizer.state< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/optimizers/_autosummary/mlx.optimizers.Optimizer.apply_gradients.html" > mlx.optimizers.Optimizer.apply_gradients< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/optimizers/_autosummary/mlx.optimizers.Optimizer.init.html" > mlx.optimizers.Optimizer.init< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/optimizers/_autosummary/mlx.optimizers.Optimizer.update.html" > mlx.optimizers.Optimizer.update< / a > < / li >
2024-02-08 12:44:23 -08:00
< / ul >
< / li >
2024-03-30 17:32:20 -07:00
< li class = "toctree-l2 has-children" > < a class = "reference internal" href = "../python/optimizers/common_optimizers.html" > Common Optimizers< / a > < input class = "toctree-checkbox" id = "toctree-checkbox-19" name = "toctree-checkbox-19" type = "checkbox" / > < label class = "toctree-toggle" for = "toctree-checkbox-19" > < i class = "fa-solid fa-chevron-down" > < / i > < / label > < ul >
2024-02-17 13:25:37 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/optimizers/_autosummary/mlx.optimizers.SGD.html" > mlx.optimizers.SGD< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/optimizers/_autosummary/mlx.optimizers.RMSprop.html" > mlx.optimizers.RMSprop< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/optimizers/_autosummary/mlx.optimizers.Adagrad.html" > mlx.optimizers.Adagrad< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/optimizers/_autosummary/mlx.optimizers.Adafactor.html" > mlx.optimizers.Adafactor< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/optimizers/_autosummary/mlx.optimizers.AdaDelta.html" > mlx.optimizers.AdaDelta< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/optimizers/_autosummary/mlx.optimizers.Adam.html" > mlx.optimizers.Adam< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/optimizers/_autosummary/mlx.optimizers.AdamW.html" > mlx.optimizers.AdamW< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/optimizers/_autosummary/mlx.optimizers.Adamax.html" > mlx.optimizers.Adamax< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/optimizers/_autosummary/mlx.optimizers.Lion.html" > mlx.optimizers.Lion< / a > < / li >
2023-12-05 12:10:03 -08:00
< / ul >
< / li >
2024-03-30 17:32:20 -07:00
< li class = "toctree-l2 has-children" > < a class = "reference internal" href = "../python/optimizers/schedulers.html" > Schedulers< / a > < input class = "toctree-checkbox" id = "toctree-checkbox-20" name = "toctree-checkbox-20" type = "checkbox" / > < label class = "toctree-toggle" for = "toctree-checkbox-20" > < i class = "fa-solid fa-chevron-down" > < / i > < / label > < ul >
2024-02-17 13:25:37 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/optimizers/_autosummary/mlx.optimizers.cosine_decay.html" > mlx.optimizers.cosine_decay< / a > < / li >
2024-02-29 12:39:18 -08:00
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/optimizers/_autosummary/mlx.optimizers.exponential_decay.html" > mlx.optimizers.exponential_decay< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/optimizers/_autosummary/mlx.optimizers.join_schedules.html" > mlx.optimizers.join_schedules< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/optimizers/_autosummary/mlx.optimizers.linear_schedule.html" > mlx.optimizers.linear_schedule< / a > < / li >
< li class = "toctree-l3" > < a class = "reference internal" href = "../python/optimizers/_autosummary/mlx.optimizers.step_decay.html" > mlx.optimizers.step_decay< / a > < / li >
2024-02-17 13:25:37 -08:00
< / ul >
< / li >
< / ul >
< / li >
2024-03-30 17:32:20 -07:00
< li class = "toctree-l1 has-children" > < a class = "reference internal" href = "../python/tree_utils.html" > Tree Utils< / a > < input class = "toctree-checkbox" id = "toctree-checkbox-21" name = "toctree-checkbox-21" type = "checkbox" / > < label class = "toctree-toggle" for = "toctree-checkbox-21" > < i class = "fa-solid fa-chevron-down" > < / i > < / label > < ul >
2023-12-05 12:10:03 -08:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.utils.tree_flatten.html" > mlx.utils.tree_flatten< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.utils.tree_unflatten.html" > mlx.utils.tree_unflatten< / a > < / li >
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.utils.tree_map.html" > mlx.utils.tree_map< / a > < / li >
2024-04-26 08:24:09 -07:00
< li class = "toctree-l2" > < a class = "reference internal" href = "../python/_autosummary/mlx.utils.tree_map_with_path.html" > mlx.utils.tree_map_with_path< / a > < / li >
2023-12-05 12:10:03 -08:00
< / ul >
< / li >
2023-11-29 12:41:56 -08:00
< / ul >
2023-12-05 12:10:03 -08:00
< p aria-level = "2" class = "caption" role = "heading" > < span class = "caption-text" > C++ API Reference< / span > < / p >
< ul class = "nav bd-sidenav" >
2023-11-29 12:41:56 -08:00
< li class = "toctree-l1" > < a class = "reference internal" href = "../cpp/ops.html" > Operations< / a > < / li >
2023-12-05 12:10:03 -08:00
< / ul >
< p aria-level = "2" class = "caption" role = "heading" > < span class = "caption-text" > Further Reading< / span > < / p >
< ul class = "nav bd-sidenav" >
< li class = "toctree-l1" > < a class = "reference internal" href = "../dev/extensions.html" > Developer Documentation< / a > < / li >
2024-03-30 17:32:20 -07:00
< li class = "toctree-l1" > < a class = "reference internal" href = "../dev/metal_debugger.html" > Metal Debugger< / a > < / li >
2023-11-29 12:41:56 -08:00
< / ul >
2023-12-05 12:10:03 -08:00
< / div >
< / nav > < / div >
< / div >
< div class = "sidebar-primary-items__end sidebar-primary__section" >
< / div >
< div id = "rtd-footer-container" > < / div >
2023-11-29 12:41:56 -08:00
< / div >
2023-12-05 12:10:03 -08:00
< main id = "main-content" class = "bd-main" >
< div class = "sbt-scroll-pixel-helper" > < / div >
2023-11-29 12:41:56 -08:00
2023-12-05 12:10:03 -08:00
< div class = "bd-content" >
< div class = "bd-article-container" >
< div class = "bd-header-article" >
< div class = "header-article-items header-article__inner" >
< div class = "header-article-items__start" >
< div class = "header-article-item" > < label class = "sidebar-toggle primary-toggle btn btn-sm" for = "__primary" title = "Toggle primary sidebar" data-bs-placement = "bottom" data-bs-toggle = "tooltip" >
< span class = "fa-solid fa-bars" > < / span >
< / label > < / div >
< / div >
< div class = "header-article-items__end" >
< div class = "header-article-item" >
< div class = "article-header-buttons" >
< a href = "https://github.com/ml-explore/mlx" target = "_blank"
class="btn btn-sm btn-source-repository-button"
title="Source repository"
data-bs-placement="bottom" data-bs-toggle="tooltip"
>
< span class = "btn__icon-container" >
< i class = "fab fa-github" > < / i >
< / span >
< / a >
< div class = "dropdown dropdown-download-buttons" >
< button class = "btn dropdown-toggle" type = "button" data-bs-toggle = "dropdown" aria-expanded = "false" aria-label = "Download this page" >
< i class = "fas fa-download" > < / i >
< / button >
< ul class = "dropdown-menu" >
< li > < a href = "../_sources/examples/llama-inference.rst" target = "_blank"
class="btn btn-sm btn-download-source-button dropdown-item"
title="Download source file"
data-bs-placement="left" data-bs-toggle="tooltip"
>
< span class = "btn__icon-container" >
< i class = "fas fa-file" > < / i >
< / span >
< span class = "btn__text-container" > .rst< / span >
< / a >
< / li >
< li >
< button onclick = "window.print()"
class="btn btn-sm btn-download-pdf-button dropdown-item"
title="Print to PDF"
data-bs-placement="left" data-bs-toggle="tooltip"
>
< span class = "btn__icon-container" >
< i class = "fas fa-file-pdf" > < / i >
< / span >
< span class = "btn__text-container" > .pdf< / span >
< / button >
< / li >
2023-11-29 12:41:56 -08:00
< / ul >
< / div >
2023-12-05 12:10:03 -08:00
< button onclick = "toggleFullScreen()"
class="btn btn-sm btn-fullscreen-button"
title="Fullscreen mode"
data-bs-placement="bottom" data-bs-toggle="tooltip"
>
< span class = "btn__icon-container" >
< i class = "fas fa-expand" > < / i >
< / span >
< / button >
< script >
document.write(`
< button class = "btn btn-sm navbar-btn theme-switch-button" title = "light/dark" aria-label = "light/dark" data-bs-placement = "bottom" data-bs-toggle = "tooltip" >
< span class = "theme-switch nav-link" data-mode = "light" > < i class = "fa-solid fa-sun fa-lg" > < / i > < / span >
< span class = "theme-switch nav-link" data-mode = "dark" > < i class = "fa-solid fa-moon fa-lg" > < / i > < / span >
< span class = "theme-switch nav-link" data-mode = "auto" > < i class = "fa-solid fa-circle-half-stroke fa-lg" > < / i > < / span >
< / button >
`);
< / script >
< script >
document.write(`
< button class = "btn btn-sm navbar-btn search-button search-button__button" title = "Search" aria-label = "Search" data-bs-placement = "bottom" data-bs-toggle = "tooltip" >
< i class = "fa-solid fa-magnifying-glass fa-lg" > < / i >
< / button >
`);
< / script >
< label class = "sidebar-toggle secondary-toggle btn btn-sm" for = "__secondary" title = "Toggle secondary sidebar" data-bs-placement = "bottom" data-bs-toggle = "tooltip" >
< span class = "fa-solid fa-list" > < / span >
< / label >
< / div > < / div >
< / div >
< / div >
< / div >
< div id = "jb-print-docs-body" class = "onlyprint" >
< h1 > LLM inference< / h1 >
<!-- Table of contents -->
< div id = "print-main-content" >
< div id = "jb-print-toc" >
< div >
< h2 > Contents < / h2 >
< / div >
< nav aria-label = "Page" >
< ul class = "visible nav section-nav flex-column" >
< li class = "toc-h2 nav-item toc-entry" > < a class = "reference internal nav-link" href = "#implementing-the-model" > Implementing the model< / a > < ul class = "visible nav section-nav flex-column" >
< li class = "toc-h3 nav-item toc-entry" > < a class = "reference internal nav-link" href = "#attention-layer" > Attention layer< / a > < / li >
< li class = "toc-h3 nav-item toc-entry" > < a class = "reference internal nav-link" href = "#encoder-layer" > Encoder layer< / a > < / li >
< li class = "toc-h3 nav-item toc-entry" > < a class = "reference internal nav-link" href = "#full-model" > Full model< / a > < / li >
< li class = "toc-h3 nav-item toc-entry" > < a class = "reference internal nav-link" href = "#generation" > Generation< / a > < / li >
< li class = "toc-h3 nav-item toc-entry" > < a class = "reference internal nav-link" href = "#putting-it-all-together" > Putting it all together< / a > < / li >
< / ul >
< / li >
< li class = "toc-h2 nav-item toc-entry" > < a class = "reference internal nav-link" href = "#converting-the-weights" > Converting the weights< / a > < / li >
< li class = "toc-h2 nav-item toc-entry" > < a class = "reference internal nav-link" href = "#weight-loading-and-benchmarking" > Weight loading and benchmarking< / a > < / li >
< li class = "toc-h2 nav-item toc-entry" > < a class = "reference internal nav-link" href = "#scripts" > Scripts< / a > < / li >
< / ul >
< / nav >
< / div >
< / div >
< / div >
< div id = "searchbox" > < / div >
< article class = "bd-article" role = "main" >
2023-11-29 12:41:56 -08:00
< section id = "llm-inference" >
2024-03-30 17:32:20 -07:00
< h1 > LLM inference< a class = "headerlink" href = "#llm-inference" title = "Link to this heading" > #< / a > < / h1 >
2023-11-29 12:41:56 -08:00
< p > MLX enables efficient inference of large-ish transformers on Apple silicon
without compromising on ease of use. In this example we will create an
inference script for the Llama family of transformer models in which the model
is defined in less than 200 lines of python.< / p >
< section id = "implementing-the-model" >
2024-03-30 17:32:20 -07:00
< h2 > Implementing the model< a class = "headerlink" href = "#implementing-the-model" title = "Link to this heading" > #< / a > < / h2 >
2023-11-29 12:41:56 -08:00
< p > We will use the neural network building blocks defined in the < code class = "xref py py-mod docutils literal notranslate" > < span class = "pre" > mlx.nn< / span > < / code >
module to concisely define the model architecture.< / p >
< section id = "attention-layer" >
2024-03-30 17:32:20 -07:00
< h3 > Attention layer< a class = "headerlink" href = "#attention-layer" title = "Link to this heading" > #< / a > < / h3 >
2023-11-29 12:41:56 -08:00
< p > We will start with the llama attention layer which notably uses the RoPE
2023-12-06 08:13:20 -08:00
positional encoding. < a class = "footnote-reference brackets" href = "#id4" id = "id1" role = "doc-noteref" > < span class = "fn-bracket" > [< / span > 1< span class = "fn-bracket" > ]< / span > < / a > In addition, our attention layer will optionally use a
2023-11-29 12:41:56 -08:00
key/value cache that will be concatenated with the provided keys and values to
support efficient inference.< / p >
2023-12-17 13:23:03 -08:00
< p > Our implementation uses < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Linear.html#mlx.nn.Linear" title = "mlx.nn.Linear" > < code class = "xref py py-class docutils literal notranslate" > < span class = "pre" > mlx.nn.Linear< / span > < / code > < / a > for all the projections and
< a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.RoPE.html#mlx.nn.RoPE" title = "mlx.nn.RoPE" > < code class = "xref py py-class docutils literal notranslate" > < span class = "pre" > mlx.nn.RoPE< / span > < / code > < / a > for the positional encoding.< / p >
2023-11-29 12:41:56 -08:00
< div class = "highlight-python notranslate" > < div class = "highlight" > < pre > < span > < / span > < span class = "kn" > import< / span > < span class = "nn" > mlx.core< / span > < span class = "k" > as< / span > < span class = "nn" > mx< / span >
< span class = "kn" > import< / span > < span class = "nn" > mlx.nn< / span > < span class = "k" > as< / span > < span class = "nn" > nn< / span >
< span class = "k" > class< / span > < span class = "nc" > LlamaAttention< / span > < span class = "p" > (< / span > < span class = "n" > nn< / span > < span class = "o" > .< / span > < span class = "n" > Module< / span > < span class = "p" > ):< / span >
< span class = "k" > def< / span > < span class = "fm" > __init__< / span > < span class = "p" > (< / span > < span class = "bp" > self< / span > < span class = "p" > ,< / span > < span class = "n" > dims< / span > < span class = "p" > :< / span > < span class = "nb" > int< / span > < span class = "p" > ,< / span > < span class = "n" > num_heads< / span > < span class = "p" > :< / span > < span class = "nb" > int< / span > < span class = "p" > ):< / span >
< span class = "nb" > super< / span > < span class = "p" > ()< / span > < span class = "o" > .< / span > < span class = "fm" > __init__< / span > < span class = "p" > ()< / span >
< span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > num_heads< / span > < span class = "o" > =< / span > < span class = "n" > num_heads< / span >
< span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > rope< / span > < span class = "o" > =< / span > < span class = "n" > nn< / span > < span class = "o" > .< / span > < span class = "n" > RoPE< / span > < span class = "p" > (< / span > < span class = "n" > dims< / span > < span class = "o" > //< / span > < span class = "n" > num_heads< / span > < span class = "p" > ,< / span > < span class = "n" > traditional< / span > < span class = "o" > =< / span > < span class = "kc" > True< / span > < span class = "p" > )< / span >
< span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > query_proj< / span > < span class = "o" > =< / span > < span class = "n" > nn< / span > < span class = "o" > .< / span > < span class = "n" > Linear< / span > < span class = "p" > (< / span > < span class = "n" > dims< / span > < span class = "p" > ,< / span > < span class = "n" > dims< / span > < span class = "p" > ,< / span > < span class = "n" > bias< / span > < span class = "o" > =< / span > < span class = "kc" > False< / span > < span class = "p" > )< / span >
< span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > key_proj< / span > < span class = "o" > =< / span > < span class = "n" > nn< / span > < span class = "o" > .< / span > < span class = "n" > Linear< / span > < span class = "p" > (< / span > < span class = "n" > dims< / span > < span class = "p" > ,< / span > < span class = "n" > dims< / span > < span class = "p" > ,< / span > < span class = "n" > bias< / span > < span class = "o" > =< / span > < span class = "kc" > False< / span > < span class = "p" > )< / span >
< span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > value_proj< / span > < span class = "o" > =< / span > < span class = "n" > nn< / span > < span class = "o" > .< / span > < span class = "n" > Linear< / span > < span class = "p" > (< / span > < span class = "n" > dims< / span > < span class = "p" > ,< / span > < span class = "n" > dims< / span > < span class = "p" > ,< / span > < span class = "n" > bias< / span > < span class = "o" > =< / span > < span class = "kc" > False< / span > < span class = "p" > )< / span >
< span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > out_proj< / span > < span class = "o" > =< / span > < span class = "n" > nn< / span > < span class = "o" > .< / span > < span class = "n" > Linear< / span > < span class = "p" > (< / span > < span class = "n" > dims< / span > < span class = "p" > ,< / span > < span class = "n" > dims< / span > < span class = "p" > ,< / span > < span class = "n" > bias< / span > < span class = "o" > =< / span > < span class = "kc" > False< / span > < span class = "p" > )< / span >
< span class = "k" > def< / span > < span class = "fm" > __call__< / span > < span class = "p" > (< / span > < span class = "bp" > self< / span > < span class = "p" > ,< / span > < span class = "n" > queries< / span > < span class = "p" > ,< / span > < span class = "n" > keys< / span > < span class = "p" > ,< / span > < span class = "n" > values< / span > < span class = "p" > ,< / span > < span class = "n" > mask< / span > < span class = "o" > =< / span > < span class = "kc" > None< / span > < span class = "p" > ,< / span > < span class = "n" > cache< / span > < span class = "o" > =< / span > < span class = "kc" > None< / span > < span class = "p" > ):< / span >
< span class = "n" > queries< / span > < span class = "o" > =< / span > < span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > query_proj< / span > < span class = "p" > (< / span > < span class = "n" > queries< / span > < span class = "p" > )< / span >
< span class = "n" > keys< / span > < span class = "o" > =< / span > < span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > key_proj< / span > < span class = "p" > (< / span > < span class = "n" > keys< / span > < span class = "p" > )< / span >
< span class = "n" > values< / span > < span class = "o" > =< / span > < span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > value_proj< / span > < span class = "p" > (< / span > < span class = "n" > values< / span > < span class = "p" > )< / span >
< span class = "c1" > # Extract some shapes< / span >
< span class = "n" > num_heads< / span > < span class = "o" > =< / span > < span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > num_heads< / span >
< span class = "n" > B< / span > < span class = "p" > ,< / span > < span class = "n" > L< / span > < span class = "p" > ,< / span > < span class = "n" > D< / span > < span class = "o" > =< / span > < span class = "n" > queries< / span > < span class = "o" > .< / span > < span class = "n" > shape< / span >
< span class = "c1" > # Prepare the queries, keys and values for the attention computation< / span >
< span class = "n" > queries< / span > < span class = "o" > =< / span > < span class = "n" > queries< / span > < span class = "o" > .< / span > < span class = "n" > reshape< / span > < span class = "p" > (< / span > < span class = "n" > B< / span > < span class = "p" > ,< / span > < span class = "n" > L< / span > < span class = "p" > ,< / span > < span class = "n" > num_heads< / span > < span class = "p" > ,< / span > < span class = "o" > -< / span > < span class = "mi" > 1< / span > < span class = "p" > )< / span > < span class = "o" > .< / span > < span class = "n" > transpose< / span > < span class = "p" > (< / span > < span class = "mi" > 0< / span > < span class = "p" > ,< / span > < span class = "mi" > 2< / span > < span class = "p" > ,< / span > < span class = "mi" > 1< / span > < span class = "p" > ,< / span > < span class = "mi" > 3< / span > < span class = "p" > )< / span >
< span class = "n" > keys< / span > < span class = "o" > =< / span > < span class = "n" > keys< / span > < span class = "o" > .< / span > < span class = "n" > reshape< / span > < span class = "p" > (< / span > < span class = "n" > B< / span > < span class = "p" > ,< / span > < span class = "n" > L< / span > < span class = "p" > ,< / span > < span class = "n" > num_heads< / span > < span class = "p" > ,< / span > < span class = "o" > -< / span > < span class = "mi" > 1< / span > < span class = "p" > )< / span > < span class = "o" > .< / span > < span class = "n" > transpose< / span > < span class = "p" > (< / span > < span class = "mi" > 0< / span > < span class = "p" > ,< / span > < span class = "mi" > 2< / span > < span class = "p" > ,< / span > < span class = "mi" > 1< / span > < span class = "p" > ,< / span > < span class = "mi" > 3< / span > < span class = "p" > )< / span >
< span class = "n" > values< / span > < span class = "o" > =< / span > < span class = "n" > values< / span > < span class = "o" > .< / span > < span class = "n" > reshape< / span > < span class = "p" > (< / span > < span class = "n" > B< / span > < span class = "p" > ,< / span > < span class = "n" > L< / span > < span class = "p" > ,< / span > < span class = "n" > num_heads< / span > < span class = "p" > ,< / span > < span class = "o" > -< / span > < span class = "mi" > 1< / span > < span class = "p" > )< / span > < span class = "o" > .< / span > < span class = "n" > transpose< / span > < span class = "p" > (< / span > < span class = "mi" > 0< / span > < span class = "p" > ,< / span > < span class = "mi" > 2< / span > < span class = "p" > ,< / span > < span class = "mi" > 1< / span > < span class = "p" > ,< / span > < span class = "mi" > 3< / span > < span class = "p" > )< / span >
< span class = "c1" > # Add RoPE to the queries and keys and combine them with the cache< / span >
< span class = "k" > if< / span > < span class = "n" > cache< / span > < span class = "ow" > is< / span > < span class = "ow" > not< / span > < span class = "kc" > None< / span > < span class = "p" > :< / span >
< span class = "n" > key_cache< / span > < span class = "p" > ,< / span > < span class = "n" > value_cache< / span > < span class = "o" > =< / span > < span class = "n" > cache< / span >
< span class = "n" > queries< / span > < span class = "o" > =< / span > < span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > rope< / span > < span class = "p" > (< / span > < span class = "n" > queries< / span > < span class = "p" > ,< / span > < span class = "n" > offset< / span > < span class = "o" > =< / span > < span class = "n" > key_cache< / span > < span class = "o" > .< / span > < span class = "n" > shape< / span > < span class = "p" > [< / span > < span class = "mi" > 2< / span > < span class = "p" > ])< / span >
< span class = "n" > keys< / span > < span class = "o" > =< / span > < span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > rope< / span > < span class = "p" > (< / span > < span class = "n" > keys< / span > < span class = "p" > ,< / span > < span class = "n" > offset< / span > < span class = "o" > =< / span > < span class = "n" > key_cache< / span > < span class = "o" > .< / span > < span class = "n" > shape< / span > < span class = "p" > [< / span > < span class = "mi" > 2< / span > < span class = "p" > ])< / span >
< span class = "n" > keys< / span > < span class = "o" > =< / span > < span class = "n" > mx< / span > < span class = "o" > .< / span > < span class = "n" > concatenate< / span > < span class = "p" > ([< / span > < span class = "n" > key_cache< / span > < span class = "p" > ,< / span > < span class = "n" > keys< / span > < span class = "p" > ],< / span > < span class = "n" > axis< / span > < span class = "o" > =< / span > < span class = "mi" > 2< / span > < span class = "p" > )< / span >
< span class = "n" > values< / span > < span class = "o" > =< / span > < span class = "n" > mx< / span > < span class = "o" > .< / span > < span class = "n" > concatenate< / span > < span class = "p" > ([< / span > < span class = "n" > value_cache< / span > < span class = "p" > ,< / span > < span class = "n" > values< / span > < span class = "p" > ],< / span > < span class = "n" > axis< / span > < span class = "o" > =< / span > < span class = "mi" > 2< / span > < span class = "p" > )< / span >
< span class = "k" > else< / span > < span class = "p" > :< / span >
< span class = "n" > queries< / span > < span class = "o" > =< / span > < span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > rope< / span > < span class = "p" > (< / span > < span class = "n" > queries< / span > < span class = "p" > )< / span >
< span class = "n" > keys< / span > < span class = "o" > =< / span > < span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > rope< / span > < span class = "p" > (< / span > < span class = "n" > keys< / span > < span class = "p" > )< / span >
< span class = "c1" > # Finally perform the attention computation< / span >
< span class = "n" > scale< / span > < span class = "o" > =< / span > < span class = "n" > math< / span > < span class = "o" > .< / span > < span class = "n" > sqrt< / span > < span class = "p" > (< / span > < span class = "mi" > 1< / span > < span class = "o" > /< / span > < span class = "n" > queries< / span > < span class = "o" > .< / span > < span class = "n" > shape< / span > < span class = "p" > [< / span > < span class = "o" > -< / span > < span class = "mi" > 1< / span > < span class = "p" > ])< / span >
< span class = "n" > scores< / span > < span class = "o" > =< / span > < span class = "p" > (< / span > < span class = "n" > queries< / span > < span class = "o" > *< / span > < span class = "n" > scale< / span > < span class = "p" > )< / span > < span class = "o" > @< / span > < span class = "n" > keys< / span > < span class = "o" > .< / span > < span class = "n" > transpose< / span > < span class = "p" > (< / span > < span class = "mi" > 0< / span > < span class = "p" > ,< / span > < span class = "mi" > 1< / span > < span class = "p" > ,< / span > < span class = "mi" > 3< / span > < span class = "p" > ,< / span > < span class = "mi" > 2< / span > < span class = "p" > )< / span >
< span class = "k" > if< / span > < span class = "n" > mask< / span > < span class = "ow" > is< / span > < span class = "ow" > not< / span > < span class = "kc" > None< / span > < span class = "p" > :< / span >
< span class = "n" > scores< / span > < span class = "o" > =< / span > < span class = "n" > scores< / span > < span class = "o" > +< / span > < span class = "n" > mask< / span >
< span class = "n" > scores< / span > < span class = "o" > =< / span > < span class = "n" > mx< / span > < span class = "o" > .< / span > < span class = "n" > softmax< / span > < span class = "p" > (< / span > < span class = "n" > scores< / span > < span class = "p" > ,< / span > < span class = "n" > axis< / span > < span class = "o" > =-< / span > < span class = "mi" > 1< / span > < span class = "p" > )< / span >
< span class = "n" > values_hat< / span > < span class = "o" > =< / span > < span class = "p" > (< / span > < span class = "n" > scores< / span > < span class = "o" > @< / span > < span class = "n" > values< / span > < span class = "p" > )< / span > < span class = "o" > .< / span > < span class = "n" > transpose< / span > < span class = "p" > (< / span > < span class = "mi" > 0< / span > < span class = "p" > ,< / span > < span class = "mi" > 2< / span > < span class = "p" > ,< / span > < span class = "mi" > 1< / span > < span class = "p" > ,< / span > < span class = "mi" > 3< / span > < span class = "p" > )< / span > < span class = "o" > .< / span > < span class = "n" > reshape< / span > < span class = "p" > (< / span > < span class = "n" > B< / span > < span class = "p" > ,< / span > < span class = "n" > L< / span > < span class = "p" > ,< / span > < span class = "o" > -< / span > < span class = "mi" > 1< / span > < span class = "p" > )< / span >
< span class = "c1" > # Note that we return the keys and values to possibly be used as a cache< / span >
< span class = "k" > return< / span > < span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > out_proj< / span > < span class = "p" > (< / span > < span class = "n" > values_hat< / span > < span class = "p" > ),< / span > < span class = "p" > (< / span > < span class = "n" > keys< / span > < span class = "p" > ,< / span > < span class = "n" > values< / span > < span class = "p" > )< / span >
< / pre > < / div >
< / div >
< / section >
< section id = "encoder-layer" >
2024-03-30 17:32:20 -07:00
< h3 > Encoder layer< a class = "headerlink" href = "#encoder-layer" title = "Link to this heading" > #< / a > < / h3 >
2023-11-29 12:41:56 -08:00
< p > The other component of the Llama model is the encoder layer which uses RMS
2023-12-06 08:13:20 -08:00
normalization < a class = "footnote-reference brackets" href = "#id5" id = "id2" role = "doc-noteref" > < span class = "fn-bracket" > [< / span > 2< span class = "fn-bracket" > ]< / span > < / a > and SwiGLU. < a class = "footnote-reference brackets" href = "#id6" id = "id3" role = "doc-noteref" > < span class = "fn-bracket" > [< / span > 3< span class = "fn-bracket" > ]< / span > < / a > For RMS normalization we will use
2023-12-17 13:23:03 -08:00
< a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.RMSNorm.html#mlx.nn.RMSNorm" title = "mlx.nn.RMSNorm" > < code class = "xref py py-class docutils literal notranslate" > < span class = "pre" > mlx.nn.RMSNorm< / span > < / code > < / a > that is already provided in < code class = "xref py py-mod docutils literal notranslate" > < span class = "pre" > mlx.nn< / span > < / code > .< / p >
2023-11-29 12:41:56 -08:00
< div class = "highlight-python notranslate" > < div class = "highlight" > < pre > < span > < / span > < span class = "k" > class< / span > < span class = "nc" > LlamaEncoderLayer< / span > < span class = "p" > (< / span > < span class = "n" > nn< / span > < span class = "o" > .< / span > < span class = "n" > Module< / span > < span class = "p" > ):< / span >
< span class = "k" > def< / span > < span class = "fm" > __init__< / span > < span class = "p" > (< / span > < span class = "bp" > self< / span > < span class = "p" > ,< / span > < span class = "n" > dims< / span > < span class = "p" > :< / span > < span class = "nb" > int< / span > < span class = "p" > ,< / span > < span class = "n" > mlp_dims< / span > < span class = "p" > :< / span > < span class = "nb" > int< / span > < span class = "p" > ,< / span > < span class = "n" > num_heads< / span > < span class = "p" > :< / span > < span class = "nb" > int< / span > < span class = "p" > ):< / span >
< span class = "nb" > super< / span > < span class = "p" > ()< / span > < span class = "o" > .< / span > < span class = "fm" > __init__< / span > < span class = "p" > ()< / span >
< span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > attention< / span > < span class = "o" > =< / span > < span class = "n" > LlamaAttention< / span > < span class = "p" > (< / span > < span class = "n" > dims< / span > < span class = "p" > ,< / span > < span class = "n" > num_heads< / span > < span class = "p" > )< / span >
< span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > norm1< / span > < span class = "o" > =< / span > < span class = "n" > nn< / span > < span class = "o" > .< / span > < span class = "n" > RMSNorm< / span > < span class = "p" > (< / span > < span class = "n" > dims< / span > < span class = "p" > )< / span >
< span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > norm2< / span > < span class = "o" > =< / span > < span class = "n" > nn< / span > < span class = "o" > .< / span > < span class = "n" > RMSNorm< / span > < span class = "p" > (< / span > < span class = "n" > dims< / span > < span class = "p" > )< / span >
< span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > linear1< / span > < span class = "o" > =< / span > < span class = "n" > nn< / span > < span class = "o" > .< / span > < span class = "n" > Linear< / span > < span class = "p" > (< / span > < span class = "n" > dims< / span > < span class = "p" > ,< / span > < span class = "n" > mlp_dims< / span > < span class = "p" > ,< / span > < span class = "n" > bias< / span > < span class = "o" > =< / span > < span class = "kc" > False< / span > < span class = "p" > )< / span >
< span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > linear2< / span > < span class = "o" > =< / span > < span class = "n" > nn< / span > < span class = "o" > .< / span > < span class = "n" > Linear< / span > < span class = "p" > (< / span > < span class = "n" > dims< / span > < span class = "p" > ,< / span > < span class = "n" > mlp_dims< / span > < span class = "p" > ,< / span > < span class = "n" > bias< / span > < span class = "o" > =< / span > < span class = "kc" > False< / span > < span class = "p" > )< / span >
< span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > linear3< / span > < span class = "o" > =< / span > < span class = "n" > nn< / span > < span class = "o" > .< / span > < span class = "n" > Linear< / span > < span class = "p" > (< / span > < span class = "n" > mlp_dims< / span > < span class = "p" > ,< / span > < span class = "n" > dims< / span > < span class = "p" > ,< / span > < span class = "n" > bias< / span > < span class = "o" > =< / span > < span class = "kc" > False< / span > < span class = "p" > )< / span >
< span class = "k" > def< / span > < span class = "fm" > __call__< / span > < span class = "p" > (< / span > < span class = "bp" > self< / span > < span class = "p" > ,< / span > < span class = "n" > x< / span > < span class = "p" > ,< / span > < span class = "n" > mask< / span > < span class = "o" > =< / span > < span class = "kc" > None< / span > < span class = "p" > ,< / span > < span class = "n" > cache< / span > < span class = "o" > =< / span > < span class = "kc" > None< / span > < span class = "p" > ):< / span >
< span class = "n" > y< / span > < span class = "o" > =< / span > < span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > norm1< / span > < span class = "p" > (< / span > < span class = "n" > x< / span > < span class = "p" > )< / span >
< span class = "n" > y< / span > < span class = "p" > ,< / span > < span class = "n" > cache< / span > < span class = "o" > =< / span > < span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > attention< / span > < span class = "p" > (< / span > < span class = "n" > y< / span > < span class = "p" > ,< / span > < span class = "n" > y< / span > < span class = "p" > ,< / span > < span class = "n" > y< / span > < span class = "p" > ,< / span > < span class = "n" > mask< / span > < span class = "p" > ,< / span > < span class = "n" > cache< / span > < span class = "p" > )< / span >
< span class = "n" > x< / span > < span class = "o" > =< / span > < span class = "n" > x< / span > < span class = "o" > +< / span > < span class = "n" > y< / span >
< span class = "n" > y< / span > < span class = "o" > =< / span > < span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > norm2< / span > < span class = "p" > (< / span > < span class = "n" > x< / span > < span class = "p" > )< / span >
< span class = "n" > a< / span > < span class = "o" > =< / span > < span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > linear1< / span > < span class = "p" > (< / span > < span class = "n" > y< / span > < span class = "p" > )< / span >
< span class = "n" > b< / span > < span class = "o" > =< / span > < span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > linear2< / span > < span class = "p" > (< / span > < span class = "n" > y< / span > < span class = "p" > )< / span >
< span class = "n" > y< / span > < span class = "o" > =< / span > < span class = "n" > a< / span > < span class = "o" > *< / span > < span class = "n" > mx< / span > < span class = "o" > .< / span > < span class = "n" > sigmoid< / span > < span class = "p" > (< / span > < span class = "n" > a< / span > < span class = "p" > )< / span > < span class = "o" > *< / span > < span class = "n" > b< / span >
< span class = "n" > y< / span > < span class = "o" > =< / span > < span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > linear3< / span > < span class = "p" > (< / span > < span class = "n" > y< / span > < span class = "p" > )< / span >
< span class = "n" > x< / span > < span class = "o" > =< / span > < span class = "n" > x< / span > < span class = "o" > +< / span > < span class = "n" > y< / span >
< span class = "k" > return< / span > < span class = "n" > x< / span > < span class = "p" > ,< / span > < span class = "n" > cache< / span >
< / pre > < / div >
< / div >
< / section >
< section id = "full-model" >
2024-03-30 17:32:20 -07:00
< h3 > Full model< a class = "headerlink" href = "#full-model" title = "Link to this heading" > #< / a > < / h3 >
2023-11-29 12:41:56 -08:00
< p > To implement any Llama model we simply have to combine < code class = "docutils literal notranslate" > < span class = "pre" > LlamaEncoderLayer< / span > < / code >
2023-12-17 13:23:03 -08:00
instances with an < a class = "reference internal" href = "../python/nn/_autosummary/mlx.nn.Embedding.html#mlx.nn.Embedding" title = "mlx.nn.Embedding" > < code class = "xref py py-class docutils literal notranslate" > < span class = "pre" > mlx.nn.Embedding< / span > < / code > < / a > to embed the input tokens.< / p >
2023-11-29 12:41:56 -08:00
< div class = "highlight-python notranslate" > < div class = "highlight" > < pre > < span > < / span > < span class = "k" > class< / span > < span class = "nc" > Llama< / span > < span class = "p" > (< / span > < span class = "n" > nn< / span > < span class = "o" > .< / span > < span class = "n" > Module< / span > < span class = "p" > ):< / span >
< span class = "k" > def< / span > < span class = "fm" > __init__< / span > < span class = "p" > (< / span >
< span class = "bp" > self< / span > < span class = "p" > ,< / span > < span class = "n" > num_layers< / span > < span class = "p" > :< / span > < span class = "nb" > int< / span > < span class = "p" > ,< / span > < span class = "n" > vocab_size< / span > < span class = "p" > :< / span > < span class = "nb" > int< / span > < span class = "p" > ,< / span > < span class = "n" > dims< / span > < span class = "p" > :< / span > < span class = "nb" > int< / span > < span class = "p" > ,< / span > < span class = "n" > mlp_dims< / span > < span class = "p" > :< / span > < span class = "nb" > int< / span > < span class = "p" > ,< / span > < span class = "n" > num_heads< / span > < span class = "p" > :< / span > < span class = "nb" > int< / span >
< span class = "p" > ):< / span >
< span class = "nb" > super< / span > < span class = "p" > ()< / span > < span class = "o" > .< / span > < span class = "fm" > __init__< / span > < span class = "p" > ()< / span >
< span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > embedding< / span > < span class = "o" > =< / span > < span class = "n" > nn< / span > < span class = "o" > .< / span > < span class = "n" > Embedding< / span > < span class = "p" > (< / span > < span class = "n" > vocab_size< / span > < span class = "p" > ,< / span > < span class = "n" > dims< / span > < span class = "p" > )< / span >
< span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > layers< / span > < span class = "o" > =< / span > < span class = "p" > [< / span >
< span class = "n" > LlamaEncoderLayer< / span > < span class = "p" > (< / span > < span class = "n" > dims< / span > < span class = "p" > ,< / span > < span class = "n" > mlp_dims< / span > < span class = "p" > ,< / span > < span class = "n" > num_heads< / span > < span class = "p" > )< / span > < span class = "k" > for< / span > < span class = "n" > _< / span > < span class = "ow" > in< / span > < span class = "nb" > range< / span > < span class = "p" > (< / span > < span class = "n" > num_layers< / span > < span class = "p" > )< / span >
< span class = "p" > ]< / span >
< span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > norm< / span > < span class = "o" > =< / span > < span class = "n" > nn< / span > < span class = "o" > .< / span > < span class = "n" > RMSNorm< / span > < span class = "p" > (< / span > < span class = "n" > dims< / span > < span class = "p" > )< / span >
< span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > out_proj< / span > < span class = "o" > =< / span > < span class = "n" > nn< / span > < span class = "o" > .< / span > < span class = "n" > Linear< / span > < span class = "p" > (< / span > < span class = "n" > dims< / span > < span class = "p" > ,< / span > < span class = "n" > vocab_size< / span > < span class = "p" > ,< / span > < span class = "n" > bias< / span > < span class = "o" > =< / span > < span class = "kc" > False< / span > < span class = "p" > )< / span >
< span class = "k" > def< / span > < span class = "fm" > __call__< / span > < span class = "p" > (< / span > < span class = "bp" > self< / span > < span class = "p" > ,< / span > < span class = "n" > x< / span > < span class = "p" > ):< / span >
< span class = "n" > mask< / span > < span class = "o" > =< / span > < span class = "n" > nn< / span > < span class = "o" > .< / span > < span class = "n" > MultiHeadAttention< / span > < span class = "o" > .< / span > < span class = "n" > create_additive_causal_mask< / span > < span class = "p" > (< / span > < span class = "n" > x< / span > < span class = "o" > .< / span > < span class = "n" > shape< / span > < span class = "p" > [< / span > < span class = "mi" > 1< / span > < span class = "p" > ])< / span >
< span class = "n" > mask< / span > < span class = "o" > =< / span > < span class = "n" > mask< / span > < span class = "o" > .< / span > < span class = "n" > astype< / span > < span class = "p" > (< / span > < span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > embedding< / span > < span class = "o" > .< / span > < span class = "n" > weight< / span > < span class = "o" > .< / span > < span class = "n" > dtype< / span > < span class = "p" > )< / span >
< span class = "n" > x< / span > < span class = "o" > =< / span > < span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > embedding< / span > < span class = "p" > (< / span > < span class = "n" > x< / span > < span class = "p" > )< / span >
< span class = "k" > for< / span > < span class = "n" > l< / span > < span class = "ow" > in< / span > < span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > layers< / span > < span class = "p" > :< / span >
< span class = "n" > x< / span > < span class = "p" > ,< / span > < span class = "n" > _< / span > < span class = "o" > =< / span > < span class = "n" > l< / span > < span class = "p" > (< / span > < span class = "n" > x< / span > < span class = "p" > ,< / span > < span class = "n" > mask< / span > < span class = "p" > )< / span >
< span class = "n" > x< / span > < span class = "o" > =< / span > < span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > norm< / span > < span class = "p" > (< / span > < span class = "n" > x< / span > < span class = "p" > )< / span >
< span class = "k" > return< / span > < span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > out_proj< / span > < span class = "p" > (< / span > < span class = "n" > x< / span > < span class = "p" > )< / span >
< / pre > < / div >
< / div >
< p > Note that in the implementation above we use a simple list to hold the encoder
layers but using < code class = "docutils literal notranslate" > < span class = "pre" > model.parameters()< / span > < / code > will still consider these layers.< / p >
< / section >
< section id = "generation" >
2024-03-30 17:32:20 -07:00
< h3 > Generation< a class = "headerlink" href = "#generation" title = "Link to this heading" > #< / a > < / h3 >
2023-11-29 12:41:56 -08:00
< p > Our < code class = "docutils literal notranslate" > < span class = "pre" > Llama< / span > < / code > module can be used for training but not inference as the
< code class = "docutils literal notranslate" > < span class = "pre" > __call__< / span > < / code > method above processes one input, completely ignores the cache and
performs no sampling whatsoever. In the rest of this subsection, we will
implement the inference function as a python generator that processes the
prompt and then autoregressively yields tokens one at a time.< / p >
< div class = "highlight-python notranslate" > < div class = "highlight" > < pre > < span > < / span > < span class = "k" > class< / span > < span class = "nc" > Llama< / span > < span class = "p" > (< / span > < span class = "n" > nn< / span > < span class = "o" > .< / span > < span class = "n" > Module< / span > < span class = "p" > ):< / span >
< span class = "o" > ...< / span >
< span class = "k" > def< / span > < span class = "nf" > generate< / span > < span class = "p" > (< / span > < span class = "bp" > self< / span > < span class = "p" > ,< / span > < span class = "n" > x< / span > < span class = "p" > ,< / span > < span class = "n" > temp< / span > < span class = "o" > =< / span > < span class = "mf" > 1.0< / span > < span class = "p" > ):< / span >
< span class = "n" > cache< / span > < span class = "o" > =< / span > < span class = "p" > []< / span >
< span class = "c1" > # Make an additive causal mask. We will need that to process the prompt.< / span >
< span class = "n" > mask< / span > < span class = "o" > =< / span > < span class = "n" > nn< / span > < span class = "o" > .< / span > < span class = "n" > MultiHeadAttention< / span > < span class = "o" > .< / span > < span class = "n" > create_additive_causal_mask< / span > < span class = "p" > (< / span > < span class = "n" > x< / span > < span class = "o" > .< / span > < span class = "n" > shape< / span > < span class = "p" > [< / span > < span class = "mi" > 1< / span > < span class = "p" > ])< / span >
< span class = "n" > mask< / span > < span class = "o" > =< / span > < span class = "n" > mask< / span > < span class = "o" > .< / span > < span class = "n" > astype< / span > < span class = "p" > (< / span > < span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > embedding< / span > < span class = "o" > .< / span > < span class = "n" > weight< / span > < span class = "o" > .< / span > < span class = "n" > dtype< / span > < span class = "p" > )< / span >
< span class = "c1" > # First we process the prompt x the same way as in __call__ but< / span >
< span class = "c1" > # save the caches in cache< / span >
< span class = "n" > x< / span > < span class = "o" > =< / span > < span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > embedding< / span > < span class = "p" > (< / span > < span class = "n" > x< / span > < span class = "p" > )< / span >
< span class = "k" > for< / span > < span class = "n" > l< / span > < span class = "ow" > in< / span > < span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > layers< / span > < span class = "p" > :< / span >
< span class = "n" > x< / span > < span class = "p" > ,< / span > < span class = "n" > c< / span > < span class = "o" > =< / span > < span class = "n" > l< / span > < span class = "p" > (< / span > < span class = "n" > x< / span > < span class = "p" > ,< / span > < span class = "n" > mask< / span > < span class = "o" > =< / span > < span class = "n" > mask< / span > < span class = "p" > )< / span >
< span class = "n" > cache< / span > < span class = "o" > .< / span > < span class = "n" > append< / span > < span class = "p" > (< / span > < span class = "n" > c< / span > < span class = "p" > )< / span > < span class = "c1" > # < --- we store the per layer cache in a< / span >
< span class = "c1" > # simple python list< / span >
< span class = "n" > x< / span > < span class = "o" > =< / span > < span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > norm< / span > < span class = "p" > (< / span > < span class = "n" > x< / span > < span class = "p" > )< / span >
< span class = "n" > y< / span > < span class = "o" > =< / span > < span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > out_proj< / span > < span class = "p" > (< / span > < span class = "n" > x< / span > < span class = "p" > [:,< / span > < span class = "o" > -< / span > < span class = "mi" > 1< / span > < span class = "p" > ])< / span > < span class = "c1" > # < --- we only care about the last logits< / span >
< span class = "c1" > # that generate the next token< / span >
< span class = "n" > y< / span > < span class = "o" > =< / span > < span class = "n" > mx< / span > < span class = "o" > .< / span > < span class = "n" > random< / span > < span class = "o" > .< / span > < span class = "n" > categorical< / span > < span class = "p" > (< / span > < span class = "n" > y< / span > < span class = "o" > *< / span > < span class = "p" > (< / span > < span class = "mi" > 1< / span > < span class = "o" > /< / span > < span class = "n" > temp< / span > < span class = "p" > ))< / span >
< span class = "c1" > # y now has size [1]< / span >
< span class = "c1" > # Since MLX is lazily evaluated nothing is computed yet.< / span >
< span class = "c1" > # Calling y.item() would force the computation to happen at< / span >
< span class = "c1" > # this point but we can also choose not to do that and let the< / span >
< span class = "c1" > # user choose when to start the computation.< / span >
< span class = "k" > yield< / span > < span class = "n" > y< / span >
< span class = "c1" > # Now we parsed the prompt and generated the first token we< / span >
< span class = "c1" > # need to feed it back into the model and loop to generate the< / span >
< span class = "c1" > # rest.< / span >
< span class = "k" > while< / span > < span class = "kc" > True< / span > < span class = "p" > :< / span >
< span class = "c1" > # Unsqueezing the last dimension to add a sequence length< / span >
< span class = "c1" > # dimension of 1< / span >
< span class = "n" > x< / span > < span class = "o" > =< / span > < span class = "n" > y< / span > < span class = "p" > [:,< / span > < span class = "kc" > None< / span > < span class = "p" > ]< / span >
< span class = "n" > x< / span > < span class = "o" > =< / span > < span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > embedding< / span > < span class = "p" > (< / span > < span class = "n" > x< / span > < span class = "p" > )< / span >
< span class = "k" > for< / span > < span class = "n" > i< / span > < span class = "ow" > in< / span > < span class = "nb" > range< / span > < span class = "p" > (< / span > < span class = "nb" > len< / span > < span class = "p" > (< / span > < span class = "n" > cache< / span > < span class = "p" > )):< / span >
< span class = "c1" > # We are overwriting the arrays in the cache list. When< / span >
< span class = "c1" > # the computation will happen, MLX will be discarding the< / span >
< span class = "c1" > # old cache the moment it is not needed anymore.< / span >
< span class = "n" > x< / span > < span class = "p" > ,< / span > < span class = "n" > cache< / span > < span class = "p" > [< / span > < span class = "n" > i< / span > < span class = "p" > ]< / span > < span class = "o" > =< / span > < span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > layers< / span > < span class = "p" > [< / span > < span class = "n" > i< / span > < span class = "p" > ](< / span > < span class = "n" > x< / span > < span class = "p" > ,< / span > < span class = "n" > mask< / span > < span class = "o" > =< / span > < span class = "kc" > None< / span > < span class = "p" > ,< / span > < span class = "n" > cache< / span > < span class = "o" > =< / span > < span class = "n" > cache< / span > < span class = "p" > [< / span > < span class = "n" > i< / span > < span class = "p" > ])< / span >
< span class = "n" > x< / span > < span class = "o" > =< / span > < span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > norm< / span > < span class = "p" > (< / span > < span class = "n" > x< / span > < span class = "p" > )< / span >
< span class = "n" > y< / span > < span class = "o" > =< / span > < span class = "bp" > self< / span > < span class = "o" > .< / span > < span class = "n" > out_proj< / span > < span class = "p" > (< / span > < span class = "n" > x< / span > < span class = "p" > [:,< / span > < span class = "o" > -< / span > < span class = "mi" > 1< / span > < span class = "p" > ])< / span >
< span class = "n" > y< / span > < span class = "o" > =< / span > < span class = "n" > mx< / span > < span class = "o" > .< / span > < span class = "n" > random< / span > < span class = "o" > .< / span > < span class = "n" > categorical< / span > < span class = "p" > (< / span > < span class = "n" > y< / span > < span class = "o" > *< / span > < span class = "p" > (< / span > < span class = "mi" > 1< / span > < span class = "o" > /< / span > < span class = "n" > temp< / span > < span class = "p" > ))< / span >
< span class = "k" > yield< / span > < span class = "n" > y< / span >
< / pre > < / div >
< / div >
< / section >
< section id = "putting-it-all-together" >
2024-03-30 17:32:20 -07:00
< h3 > Putting it all together< a class = "headerlink" href = "#putting-it-all-together" title = "Link to this heading" > #< / a > < / h3 >
2023-11-29 12:41:56 -08:00
< p > We now have everything we need to create a Llama model and sample tokens from
it. In the following code, we randomly initialize a small Llama model, process
6 tokens of prompt and generate 10 tokens.< / p >
< div class = "highlight-python notranslate" > < div class = "highlight" > < pre > < span > < / span > < span class = "n" > model< / span > < span class = "o" > =< / span > < span class = "n" > Llama< / span > < span class = "p" > (< / span > < span class = "n" > num_layers< / span > < span class = "o" > =< / span > < span class = "mi" > 12< / span > < span class = "p" > ,< / span > < span class = "n" > vocab_size< / span > < span class = "o" > =< / span > < span class = "mi" > 8192< / span > < span class = "p" > ,< / span > < span class = "n" > dims< / span > < span class = "o" > =< / span > < span class = "mi" > 512< / span > < span class = "p" > ,< / span > < span class = "n" > mlp_dims< / span > < span class = "o" > =< / span > < span class = "mi" > 1024< / span > < span class = "p" > ,< / span > < span class = "n" > num_heads< / span > < span class = "o" > =< / span > < span class = "mi" > 8< / span > < span class = "p" > )< / span >
< span class = "c1" > # Since MLX is lazily evaluated nothing has actually been materialized yet.< / span >
< span class = "c1" > # We could have set the `dims` to 20_000 on a machine with 8GB of RAM and the< / span >
< span class = "c1" > # code above would still run. Let' s actually materialize the model.< / span >
< span class = "n" > mx< / span > < span class = "o" > .< / span > < span class = "n" > eval< / span > < span class = "p" > (< / span > < span class = "n" > model< / span > < span class = "o" > .< / span > < span class = "n" > parameters< / span > < span class = "p" > ())< / span >
< span class = "n" > prompt< / span > < span class = "o" > =< / span > < span class = "n" > mx< / span > < span class = "o" > .< / span > < span class = "n" > array< / span > < span class = "p" > ([[< / span > < span class = "mi" > 1< / span > < span class = "p" > ,< / span > < span class = "mi" > 10< / span > < span class = "p" > ,< / span > < span class = "mi" > 8< / span > < span class = "p" > ,< / span > < span class = "mi" > 32< / span > < span class = "p" > ,< / span > < span class = "mi" > 44< / span > < span class = "p" > ,< / span > < span class = "mi" > 7< / span > < span class = "p" > ]])< / span > < span class = "c1" > # < -- Note the double brackets because we< / span >
< span class = "c1" > # have a batch dimension even< / span >
< span class = "c1" > # though it is 1 in this case< / span >
< span class = "n" > generated< / span > < span class = "o" > =< / span > < span class = "p" > [< / span > < span class = "n" > t< / span > < span class = "k" > for< / span > < span class = "n" > i< / span > < span class = "p" > ,< / span > < span class = "n" > t< / span > < span class = "ow" > in< / span > < span class = "nb" > zip< / span > < span class = "p" > (< / span > < span class = "nb" > range< / span > < span class = "p" > (< / span > < span class = "mi" > 10< / span > < span class = "p" > ),< / span > < span class = "n" > model< / span > < span class = "o" > .< / span > < span class = "n" > generate< / span > < span class = "p" > (< / span > < span class = "n" > prompt< / span > < span class = "p" > ,< / span > < span class = "mf" > 0.8< / span > < span class = "p" > ))]< / span >
< span class = "c1" > # Since we haven' t evaluated anything, nothing is computed yet. The list< / span >
< span class = "c1" > # `generated` contains the arrays that hold the computation graph for the< / span >
< span class = "c1" > # full processing of the prompt and the generation of 10 tokens.< / span >
< span class = "c1" > #< / span >
< span class = "c1" > # We can evaluate them one at a time, or all together. Concatenate them or< / span >
< span class = "c1" > # print them. They would all result in very similar runtimes and give exactly< / span >
< span class = "c1" > # the same results.< / span >
< span class = "n" > mx< / span > < span class = "o" > .< / span > < span class = "n" > eval< / span > < span class = "p" > (< / span > < span class = "n" > generated< / span > < span class = "p" > )< / span >
< / pre > < / div >
< / div >
< / section >
< / section >
< section id = "converting-the-weights" >
2024-03-30 17:32:20 -07:00
< h2 > Converting the weights< a class = "headerlink" href = "#converting-the-weights" title = "Link to this heading" > #< / a > < / h2 >
2023-11-29 12:41:56 -08:00
< p > This section assumes that you have access to the original Llama weights and the
SentencePiece model that comes with them. We will write a small script to
convert the PyTorch weights to MLX compatible ones and write them in a NPZ file
that can be loaded directly by MLX.< / p >
< div class = "highlight-python notranslate" > < div class = "highlight" > < pre > < span > < / span > < span class = "kn" > import< / span > < span class = "nn" > argparse< / span >
< span class = "kn" > from< / span > < span class = "nn" > itertools< / span > < span class = "kn" > import< / span > < span class = "n" > starmap< / span >
< span class = "kn" > import< / span > < span class = "nn" > numpy< / span > < span class = "k" > as< / span > < span class = "nn" > np< / span >
< span class = "kn" > import< / span > < span class = "nn" > torch< / span >
< span class = "k" > def< / span > < span class = "nf" > map_torch_to_mlx< / span > < span class = "p" > (< / span > < span class = "n" > key< / span > < span class = "p" > ,< / span > < span class = "n" > value< / span > < span class = "p" > ):< / span >
< span class = "k" > if< / span > < span class = "s2" > " tok_embedding" < / span > < span class = "ow" > in< / span > < span class = "n" > key< / span > < span class = "p" > :< / span >
< span class = "n" > key< / span > < span class = "o" > =< / span > < span class = "s2" > " embedding.weight" < / span >
< span class = "k" > elif< / span > < span class = "s2" > " norm" < / span > < span class = "ow" > in< / span > < span class = "n" > key< / span > < span class = "p" > :< / span >
< span class = "n" > key< / span > < span class = "o" > =< / span > < span class = "n" > key< / span > < span class = "o" > .< / span > < span class = "n" > replace< / span > < span class = "p" > (< / span > < span class = "s2" > " attention_norm" < / span > < span class = "p" > ,< / span > < span class = "s2" > " norm1" < / span > < span class = "p" > )< / span > < span class = "o" > .< / span > < span class = "n" > replace< / span > < span class = "p" > (< / span > < span class = "s2" > " ffn_norm" < / span > < span class = "p" > ,< / span > < span class = "s2" > " norm2" < / span > < span class = "p" > )< / span >
< span class = "k" > elif< / span > < span class = "s2" > " wq" < / span > < span class = "ow" > in< / span > < span class = "n" > key< / span > < span class = "ow" > or< / span > < span class = "s2" > " wk" < / span > < span class = "ow" > in< / span > < span class = "n" > key< / span > < span class = "ow" > or< / span > < span class = "s2" > " wv" < / span > < span class = "ow" > in< / span > < span class = "n" > key< / span > < span class = "ow" > or< / span > < span class = "s2" > " wo" < / span > < span class = "ow" > in< / span > < span class = "n" > key< / span > < span class = "p" > :< / span >
< span class = "n" > key< / span > < span class = "o" > =< / span > < span class = "n" > key< / span > < span class = "o" > .< / span > < span class = "n" > replace< / span > < span class = "p" > (< / span > < span class = "s2" > " wq" < / span > < span class = "p" > ,< / span > < span class = "s2" > " query_proj" < / span > < span class = "p" > )< / span >
< span class = "n" > key< / span > < span class = "o" > =< / span > < span class = "n" > key< / span > < span class = "o" > .< / span > < span class = "n" > replace< / span > < span class = "p" > (< / span > < span class = "s2" > " wk" < / span > < span class = "p" > ,< / span > < span class = "s2" > " key_proj" < / span > < span class = "p" > )< / span >
< span class = "n" > key< / span > < span class = "o" > =< / span > < span class = "n" > key< / span > < span class = "o" > .< / span > < span class = "n" > replace< / span > < span class = "p" > (< / span > < span class = "s2" > " wv" < / span > < span class = "p" > ,< / span > < span class = "s2" > " value_proj" < / span > < span class = "p" > )< / span >
< span class = "n" > key< / span > < span class = "o" > =< / span > < span class = "n" > key< / span > < span class = "o" > .< / span > < span class = "n" > replace< / span > < span class = "p" > (< / span > < span class = "s2" > " wo" < / span > < span class = "p" > ,< / span > < span class = "s2" > " out_proj" < / span > < span class = "p" > )< / span >
< span class = "k" > elif< / span > < span class = "s2" > " w1" < / span > < span class = "ow" > in< / span > < span class = "n" > key< / span > < span class = "ow" > or< / span > < span class = "s2" > " w2" < / span > < span class = "ow" > in< / span > < span class = "n" > key< / span > < span class = "ow" > or< / span > < span class = "s2" > " w3" < / span > < span class = "ow" > in< / span > < span class = "n" > key< / span > < span class = "p" > :< / span >
< span class = "c1" > # The FFN is a separate submodule in PyTorch< / span >
< span class = "n" > key< / span > < span class = "o" > =< / span > < span class = "n" > key< / span > < span class = "o" > .< / span > < span class = "n" > replace< / span > < span class = "p" > (< / span > < span class = "s2" > " feed_forward.w1" < / span > < span class = "p" > ,< / span > < span class = "s2" > " linear1" < / span > < span class = "p" > )< / span >
< span class = "n" > key< / span > < span class = "o" > =< / span > < span class = "n" > key< / span > < span class = "o" > .< / span > < span class = "n" > replace< / span > < span class = "p" > (< / span > < span class = "s2" > " feed_forward.w3" < / span > < span class = "p" > ,< / span > < span class = "s2" > " linear2" < / span > < span class = "p" > )< / span >
< span class = "n" > key< / span > < span class = "o" > =< / span > < span class = "n" > key< / span > < span class = "o" > .< / span > < span class = "n" > replace< / span > < span class = "p" > (< / span > < span class = "s2" > " feed_forward.w2" < / span > < span class = "p" > ,< / span > < span class = "s2" > " linear3" < / span > < span class = "p" > )< / span >
< span class = "k" > elif< / span > < span class = "s2" > " output" < / span > < span class = "ow" > in< / span > < span class = "n" > key< / span > < span class = "p" > :< / span >
< span class = "n" > key< / span > < span class = "o" > =< / span > < span class = "n" > key< / span > < span class = "o" > .< / span > < span class = "n" > replace< / span > < span class = "p" > (< / span > < span class = "s2" > " output" < / span > < span class = "p" > ,< / span > < span class = "s2" > " out_proj" < / span > < span class = "p" > )< / span >
< span class = "k" > elif< / span > < span class = "s2" > " rope" < / span > < span class = "ow" > in< / span > < span class = "n" > key< / span > < span class = "p" > :< / span >
< span class = "k" > return< / span > < span class = "kc" > None< / span > < span class = "p" > ,< / span > < span class = "kc" > None< / span >
< span class = "k" > return< / span > < span class = "n" > key< / span > < span class = "p" > ,< / span > < span class = "n" > value< / span > < span class = "o" > .< / span > < span class = "n" > numpy< / span > < span class = "p" > ()< / span >
< span class = "k" > if< / span > < span class = "vm" > __name__< / span > < span class = "o" > ==< / span > < span class = "s2" > " __main__" < / span > < span class = "p" > :< / span >
< span class = "n" > parser< / span > < span class = "o" > =< / span > < span class = "n" > argparse< / span > < span class = "o" > .< / span > < span class = "n" > ArgumentParser< / span > < span class = "p" > (< / span > < span class = "n" > description< / span > < span class = "o" > =< / span > < span class = "s2" > " Convert Llama weights to MLX" < / span > < span class = "p" > )< / span >
< span class = "n" > parser< / span > < span class = "o" > .< / span > < span class = "n" > add_argument< / span > < span class = "p" > (< / span > < span class = "s2" > " torch_weights" < / span > < span class = "p" > )< / span >
< span class = "n" > parser< / span > < span class = "o" > .< / span > < span class = "n" > add_argument< / span > < span class = "p" > (< / span > < span class = "s2" > " output_file" < / span > < span class = "p" > )< / span >
< span class = "n" > args< / span > < span class = "o" > =< / span > < span class = "n" > parser< / span > < span class = "o" > .< / span > < span class = "n" > parse_args< / span > < span class = "p" > ()< / span >
< span class = "n" > state< / span > < span class = "o" > =< / span > < span class = "n" > torch< / span > < span class = "o" > .< / span > < span class = "n" > load< / span > < span class = "p" > (< / span > < span class = "n" > args< / span > < span class = "o" > .< / span > < span class = "n" > torch_weights< / span > < span class = "p" > )< / span >
< span class = "n" > np< / span > < span class = "o" > .< / span > < span class = "n" > savez< / span > < span class = "p" > (< / span >
< span class = "n" > args< / span > < span class = "o" > .< / span > < span class = "n" > output_file< / span > < span class = "p" > ,< / span >
< span class = "o" > **< / span > < span class = "p" > {< / span > < span class = "n" > k< / span > < span class = "p" > :< / span > < span class = "n" > v< / span > < span class = "k" > for< / span > < span class = "n" > k< / span > < span class = "p" > ,< / span > < span class = "n" > v< / span > < span class = "ow" > in< / span > < span class = "n" > starmap< / span > < span class = "p" > (< / span > < span class = "n" > map_torch_to_mlx< / span > < span class = "p" > ,< / span > < span class = "n" > state< / span > < span class = "o" > .< / span > < span class = "n" > items< / span > < span class = "p" > ())< / span > < span class = "k" > if< / span > < span class = "n" > k< / span > < span class = "ow" > is< / span > < span class = "ow" > not< / span > < span class = "kc" > None< / span > < span class = "p" > }< / span >
< span class = "p" > )< / span >
< / pre > < / div >
< / div >
< / section >
< section id = "weight-loading-and-benchmarking" >
2024-03-30 17:32:20 -07:00
< h2 > Weight loading and benchmarking< a class = "headerlink" href = "#weight-loading-and-benchmarking" title = "Link to this heading" > #< / a > < / h2 >
2023-11-29 12:41:56 -08:00
< p > After converting the weights to be compatible to our implementation, all that is
left is to load them from disk and we can finally use the LLM to generate text.
We can load numpy format files using the < a class = "reference internal" href = "../python/_autosummary/mlx.core.load.html#mlx.core.load" title = "mlx.core.load" > < code class = "xref py py-func docutils literal notranslate" > < span class = "pre" > mlx.core.load()< / span > < / code > < / a > operation.< / p >
< p > To create a parameter dictionary from the key/value representation of NPZ files
we will use the < a class = "reference internal" href = "../python/_autosummary/mlx.utils.tree_unflatten.html#mlx.utils.tree_unflatten" title = "mlx.utils.tree_unflatten" > < code class = "xref py py-func docutils literal notranslate" > < span class = "pre" > mlx.utils.tree_unflatten()< / span > < / code > < / a > helper method as follows:< / p >
< div class = "highlight-python notranslate" > < div class = "highlight" > < pre > < span > < / span > < span class = "kn" > from< / span > < span class = "nn" > mlx.utils< / span > < span class = "kn" > import< / span > < span class = "n" > tree_unflatten< / span >
< span class = "n" > model< / span > < span class = "o" > .< / span > < span class = "n" > update< / span > < span class = "p" > (< / span > < span class = "n" > tree_unflatten< / span > < span class = "p" > (< / span > < span class = "nb" > list< / span > < span class = "p" > (< / span > < span class = "n" > mx< / span > < span class = "o" > .< / span > < span class = "n" > load< / span > < span class = "p" > (< / span > < span class = "n" > weight_file< / span > < span class = "p" > )< / span > < span class = "o" > .< / span > < span class = "n" > items< / span > < span class = "p" > ())))< / span >
< / pre > < / div >
< / div >
< p > < a class = "reference internal" href = "../python/_autosummary/mlx.utils.tree_unflatten.html#mlx.utils.tree_unflatten" title = "mlx.utils.tree_unflatten" > < code class = "xref py py-meth docutils literal notranslate" > < span class = "pre" > mlx.utils.tree_unflatten()< / span > < / code > < / a > will take keys from the NPZ file that look
like < code class = "docutils literal notranslate" > < span class = "pre" > layers.2.attention.query_proj.weight< / span > < / code > and will transform them to< / p >
< div class = "highlight-python notranslate" > < div class = "highlight" > < pre > < span > < / span > < span class = "p" > {< / span > < span class = "s2" > " layers" < / span > < span class = "p" > :< / span > < span class = "p" > [< / span > < span class = "o" > ...< / span > < span class = "p" > ,< / span > < span class = "o" > ...< / span > < span class = "p" > ,< / span > < span class = "p" > {< / span > < span class = "s2" > " attention" < / span > < span class = "p" > :< / span > < span class = "p" > {< / span > < span class = "s2" > " query_proj" < / span > < span class = "p" > :< / span > < span class = "p" > {< / span > < span class = "s2" > " weight" < / span > < span class = "p" > :< / span > < span class = "o" > ...< / span > < span class = "p" > }}}]}< / span >
< / pre > < / div >
< / div >
< p > which can then be used to update the model. Note that the method above incurs
several unnecessary copies from disk to numpy and then from numpy to MLX. It
will be replaced in the future with direct loading to MLX.< / p >
2024-01-10 14:14:38 -08:00
< p > You can download the full example code in < a class = "reference external" href = "https://github.com/ml-explore/mlx-examples/tree/main/llms/llama" > mlx-examples< / a > . Assuming, the
2023-11-29 12:41:56 -08:00
existence of < code class = "docutils literal notranslate" > < span class = "pre" > weights.pth< / span > < / code > and < code class = "docutils literal notranslate" > < span class = "pre" > tokenizer.model< / span > < / code > in the current working
directory we can play around with our inference script as follows (the timings
are representative of an M1 Ultra and the 7B parameter Llama model):< / p >
2024-03-30 17:32:20 -07:00
< div class = "highlight-bash notranslate" > < div class = "highlight" > < pre > < span > < / span > $< span class = "w" > < / span > python< span class = "w" > < / span > convert.py< span class = "w" > < / span > weights.pth< span class = "w" > < / span > llama-7B.mlx.npz
$< span class = "w" > < / span > python< span class = "w" > < / span > llama.py< span class = "w" > < / span > llama-7B.mlx.npz< span class = "w" > < / span > tokenizer.model< span class = "w" > < / span > < span class = "s1" > ' Call me Ishmael. Some years ago never mind how long precisely' < / span >
< span class = "o" > [< / span > INFO< span class = "o" > ]< / span > < span class = "w" > < / span > Loading< span class = "w" > < / span > model< span class = "w" > < / span > from< span class = "w" > < / span > disk:< span class = "w" > < / span > < span class = "m" > 5< / span > .247< span class = "w" > < / span > s
Press< span class = "w" > < / span > enter< span class = "w" > < / span > to< span class = "w" > < / span > start< span class = "w" > < / span > generation
2023-11-29 12:41:56 -08:00
------
2024-03-30 17:32:20 -07:00
,< span class = "w" > < / span > having< span class = "w" > < / span > little< span class = "w" > < / span > or< span class = "w" > < / span > no< span class = "w" > < / span > money< span class = "w" > < / span > < span class = "k" > in< / span > < span class = "w" > < / span > my< span class = "w" > < / span > purse,< span class = "w" > < / span > and< span class = "w" > < / span > nothing< span class = "w" > < / span > of< span class = "w" > < / span > greater< span class = "w" > < / span > consequence< span class = "w" > < / span > < span class = "k" > in< / span > < span class = "w" > < / span > my< span class = "w" > < / span > mind,< span class = "w" > < / span > I< span class = "w" > < / span > happened< span class = "w" > < / span > to< span class = "w" > < / span > be< span class = "w" > < / span > walking< span class = "w" > < / span > down< span class = "w" > < / span > Gower< span class = "w" > < / span > Street< span class = "w" > < / span > < span class = "k" > in< / span > < span class = "w" > < / span > the< span class = "w" > < / span > afternoon,< span class = "w" > < / span > < span class = "k" > in< / span > < span class = "w" > < / span > the< span class = "w" > < / span > heavy< span class = "w" > < / span > rain,< span class = "w" > < / span > and< span class = "w" > < / span > I< span class = "w" > < / span > saw< span class = "w" > < / span > a< span class = "w" > < / span > few< span class = "w" > < / span > steps< span class = "w" > < / span > off,< span class = "w" > < / span > a< span class = "w" > < / span > man< span class = "w" > < / span > < span class = "k" > in< / span > < span class = "w" > < / span > rags,< span class = "w" > < / span > who< span class = "w" > < / span > sat< span class = "w" > < / span > upon< span class = "w" > < / span > his< span class = "w" > < / span > bundle< span class = "w" > < / span > and< span class = "w" > < / span > looked< span class = "w" > < / span > hard< span class = "w" > < / span > into< span class = "w" > < / span > the< span class = "w" > < / span > wet< span class = "w" > < / span > as< span class = "w" > < / span > < span class = "k" > if< / span > < span class = "w" > < / span > he< span class = "w" > < / span > were< span class = "w" > < / span > going< span class = "w" > < / span > to< span class = "w" > < / span > cry.< span class = "w" > < / span > I< span class = "w" > < / span > watched< span class = "w" > < / span > him< span class = "w" > < / span > attentively< span class = "w" > < / span > < span class = "k" > for< / span > < span class = "w" > < / span > some< span class = "w" > < / span > time,< span class = "w" > < / span > and< span class = "w" > < / span > could< span class = "w" > < / span > not< span class = "w" > < / span > but< span class = "w" > < / span > observe< span class = "w" > < / span > that,< span class = "w" > < / span > though< span class = "w" > < / span > a< span class = "w" > < / span > numerous< span class = "w" > < / span > crowd< span class = "w" > < / span > was< span class = "w" > < / span > hurrying< span class = "w" > < / span > up< span class = "w" > < / span > and< span class = "w" > < / span > down,
2023-11-29 12:41:56 -08:00
------
2024-03-30 17:32:20 -07:00
< span class = "o" > [< / span > INFO< span class = "o" > ]< / span > < span class = "w" > < / span > Prompt< span class = "w" > < / span > processing:< span class = "w" > < / span > < span class = "m" > 0< / span > .437< span class = "w" > < / span > s
< span class = "o" > [< / span > INFO< span class = "o" > ]< / span > < span class = "w" > < / span > Full< span class = "w" > < / span > generation:< span class = "w" > < / span > < span class = "m" > 4< / span > .330< span class = "w" > < / span > s
2023-11-29 12:41:56 -08:00
< / pre > < / div >
< / div >
< p > We observe that 4.3 seconds are required to generate 100 tokens and 0.4 seconds
of those are spent processing the prompt. This amounts to a little over < strong > 39 ms
per token< / strong > .< / p >
< p > By running with a much bigger prompt we can see that the per token generation
time as well as the prompt processing time remains almost constant.< / p >
2024-03-30 17:32:20 -07:00
< div class = "highlight-bash notranslate" > < div class = "highlight" > < pre > < span > < / span > $< span class = "w" > < / span > python< span class = "w" > < / span > llama.py< span class = "w" > < / span > llama-7B.mlx.npz< span class = "w" > < / span > tokenizer.model< span class = "w" > < / span > < span class = "s1" > ' Call me Ishmael. Some years ago never mind how long precisely, having little or no money in my purse, and nothing of greater consequence in my mind, I happened to be walking down Gower Street in the afternoon, in the heavy rain, and I saw a few steps off, a man in rags, who sat upon his bundle and looked hard into the wet as if he were going to cry. I watched him attentively for some time, and could not but observe that, though a numerous crowd was hurrying up and down, nobody took the least notice of him. I stopped at last, at a little distance, as if I had been in doubt, and after looking on a few minutes, walked straight up to him. He slowly raised his eyes, and fixed them upon me for a moment, without speaking, and then resumed his place and posture as before. I stood looking at him for a while, feeling very much pain at heart, and then said to him, “What are you doing there?” Something like a smile passed over his face, as he said slowly, “I am waiting for someone; but it has been three quarters of an hour now, and he has not come.” “What is it you are waiting for?” said I. Still he made no immediate reply, but again put his face down upon his hands, and did not' < / span >
< span class = "o" > [< / span > INFO< span class = "o" > ]< / span > < span class = "w" > < / span > Loading< span class = "w" > < / span > model< span class = "w" > < / span > from< span class = "w" > < / span > disk:< span class = "w" > < / span > < span class = "m" > 5< / span > .247< span class = "w" > < / span > s
Press< span class = "w" > < / span > enter< span class = "w" > < / span > to< span class = "w" > < / span > start< span class = "w" > < / span > generation
2023-11-29 12:41:56 -08:00
------
2024-03-30 17:32:20 -07:00
take< span class = "w" > < / span > his< span class = "w" > < / span > eyes< span class = "w" > < / span > from< span class = "w" > < / span > the< span class = "w" > < / span > ground.< span class = "w" > < / span > “What< span class = "w" > < / span > is< span class = "w" > < / span > it< span class = "w" > < / span > you< span class = "w" > < / span > are< span class = "w" > < / span > waiting< span class = "w" > < / span > < span class = "k" > for< / span > ?”< span class = "w" > < / span > said< span class = "w" > < / span > I.< span class = "w" > < / span > “I< span class = "w" > < / span > am< span class = "w" > < / span > not< span class = "w" > < / span > accustomed< span class = "w" > < / span > to< span class = "w" > < / span > be< span class = "w" > < / span > thus< span class = "w" > < / span > questioned,”< span class = "w" > < / span > said< span class = "w" > < / span > he.< span class = "w" > < / span > “You< span class = "w" > < / span > look< span class = "w" > < / span > like< span class = "w" > < / span > a< span class = "w" > < / span > reasonable< span class = "w" > < / span > man—tell< span class = "w" > < / span > me,< span class = "w" > < / span > < span class = "k" > then< / span > ,< span class = "w" > < / span > what< span class = "w" > < / span > are< span class = "w" > < / span > you< span class = "w" > < / span > waiting< span class = "w" > < / span > < span class = "k" > for< / span > ?”< span class = "w" > < / span > “You< span class = "w" > < / span > would< span class = "w" > < / span > not< span class = "w" > < / span > understand,”< span class = "w" > < / span > he< span class = "w" > < / span > replied< span class = "p" > ;< / span > < span class = "w" > < / span > “and< span class = "w" > < / span > how< span class = "w" > < / span > could< span class = "w" > < / span > you< span class = "w" > < / span > < span class = "nb" > help< / span > < span class = "w" > < / span > me,< span class = "w" > < / span > < span class = "k" > if< / span > < span class = "w" > < / span > I< span class = "w" > < / span > were< span class = "w" > < / span > to< span class = "w" > < / span > tell< span class = "w" > < / span > you?”< span class = "w" > < / span > “I< span class = "w" > < / span > should< span class = "w" > < / span > not< span class = "w" > < / span > only< span class = "w" > < / span > understand,< span class = "w" > < / span > but< span class = "w" > < / span > would< span class = "w" > < / span > < span class = "k" > do< / span > < span class = "w" > < / span > all< span class = "w" > < / span > that< span class = "w" > < / span > I< span class = "w" > < / span > could,”< span class = "w" > < / span > said< span class = "w" > < / span > I.< span class = "w" > < / span > He< span class = "w" > < / span > did< span class = "w" > < / span > not
2023-11-29 12:41:56 -08:00
------
2024-03-30 17:32:20 -07:00
< span class = "o" > [< / span > INFO< span class = "o" > ]< / span > < span class = "w" > < / span > Prompt< span class = "w" > < / span > processing:< span class = "w" > < / span > < span class = "m" > 0< / span > .579< span class = "w" > < / span > s
< span class = "o" > [< / span > INFO< span class = "o" > ]< / span > < span class = "w" > < / span > Full< span class = "w" > < / span > generation:< span class = "w" > < / span > < span class = "m" > 4< / span > .690< span class = "w" > < / span > s
$< span class = "w" > < / span > python< span class = "w" > < / span > llama.py< span class = "w" > < / span > --num-tokens< span class = "w" > < / span > < span class = "m" > 500< / span > < span class = "w" > < / span > llama-7B.mlx.npz< span class = "w" > < / span > tokenizer.model< span class = "w" > < / span > < span class = "s1" > ' Call me Ishmael. Some years ago never mind how long precisely, having little or no money in my purse, and nothing of greater consequence in my mind, I happened to be walking down Gower Street in the afternoon, in the heavy rain, and I saw a few steps off, a man in rags, who sat upon his bundle and looked hard into the wet as if he were going to cry. I watched him attentively for some time, and could not but observe that, though a numerous crowd was hurrying up and down, nobody took the least notice of him. I stopped at last, at a little distance, as if I had been in doubt, and after looking on a few minutes, walked straight up to him. He slowly raised his eyes, and fixed them upon me for a moment, without speaking, and then resumed his place and posture as before. I stood looking at him for a while, feeling very much pain at heart, and then said to him, “What are you doing there?” Something like a smile passed over his face, as he said slowly, “I am waiting for someone; but it has been three quarters of an hour now, and he has not come.” “What is it you are waiting for?” said I. Still he made no immediate reply, but again put his face down upon his hands, and did not' < / span >
< span class = "o" > [< / span > INFO< span class = "o" > ]< / span > < span class = "w" > < / span > Loading< span class = "w" > < / span > model< span class = "w" > < / span > from< span class = "w" > < / span > disk:< span class = "w" > < / span > < span class = "m" > 5< / span > .628< span class = "w" > < / span > s
Press< span class = "w" > < / span > enter< span class = "w" > < / span > to< span class = "w" > < / span > start< span class = "w" > < / span > generation
2023-11-29 12:41:56 -08:00
------
2024-03-30 17:32:20 -07:00
take< span class = "w" > < / span > his< span class = "w" > < / span > eyes< span class = "w" > < / span > from< span class = "w" > < / span > the< span class = "w" > < / span > ground.< span class = "w" > < / span > “What< span class = "w" > < / span > is< span class = "w" > < / span > it< span class = "w" > < / span > you< span class = "w" > < / span > are< span class = "w" > < / span > waiting< span class = "w" > < / span > < span class = "k" > for< / span > ?”< span class = "w" > < / span > said< span class = "w" > < / span > I.< span class = "w" > < / span > “I< span class = "w" > < / span > am< span class = "w" > < / span > not< span class = "w" > < / span > accustomed< span class = "w" > < / span > to< span class = "w" > < / span > be< span class = "w" > < / span > thus< span class = "w" > < / span > questioned,”< span class = "w" > < / span > said< span class = "w" > < / span > he.< span class = "w" > < / span > “You< span class = "w" > < / span > look< span class = "w" > < / span > like< span class = "w" > < / span > a< span class = "w" > < / span > reasonable< span class = "w" > < / span > man—tell< span class = "w" > < / span > me,< span class = "w" > < / span > < span class = "k" > then< / span > ,< span class = "w" > < / span > what< span class = "w" > < / span > are< span class = "w" > < / span > you< span class = "w" > < / span > waiting< span class = "w" > < / span > < span class = "k" > for< / span > ?”< span class = "w" > < / span > “You< span class = "w" > < / span > would< span class = "w" > < / span > not< span class = "w" > < / span > understand,”< span class = "w" > < / span > he< span class = "w" > < / span > replied< span class = "p" > ;< / span > < span class = "w" > < / span > “and< span class = "w" > < / span > how< span class = "w" > < / span > could< span class = "w" > < / span > you< span class = "w" > < / span > < span class = "nb" > help< / span > < span class = "w" > < / span > me,< span class = "w" > < / span > < span class = "k" > if< / span > < span class = "w" > < / span > I< span class = "w" > < / span > were< span class = "w" > < / span > to< span class = "w" > < / span > tell< span class = "w" > < / span > you?”< span class = "w" > < / span > “I< span class = "w" > < / span > should< span class = "w" > < / span > not< span class = "w" > < / span > only< span class = "w" > < / span > understand,< span class = "w" > < / span > but< span class = "w" > < / span > would< span class = "w" > < / span > < span class = "k" > do< / span > < span class = "w" > < / span > all< span class = "w" > < / span > that< span class = "w" > < / span > I< span class = "w" > < / span > could,”< span class = "w" > < / span > said< span class = "w" > < / span > I.< span class = "w" > < / span > He< span class = "w" > < / span > did< span class = "w" > < / span > not< span class = "w" > < / span > reply,< span class = "w" > < / span > but< span class = "w" > < / span > still< span class = "w" > < / span > went< span class = "w" > < / span > on< span class = "w" > < / span > looking< span class = "w" > < / span > at< span class = "w" > < / span > the< span class = "w" > < / span > ground,< span class = "w" > < / span > and< span class = "w" > < / span > took< span class = "w" > < / span > hold< span class = "w" > < / span > of< span class = "w" > < / span > his< span class = "w" > < / span > bundle< span class = "w" > < / span > with< span class = "w" > < / span > a< span class = "w" > < / span > nervous< span class = "w" > < / span > trembling.< span class = "w" > < / span > I< span class = "w" > < / span > waited< span class = "w" > < / span > some< span class = "w" > < / span > time,< span class = "w" > < / span > and< span class = "w" > < / span > < span class = "k" > then< / span > < span class = "w" > < / span > resumed.< span class = "w" > < / span > “It< span class = "w" > < / span > is< span class = "w" > < / span > of< span class = "w" > < / span > no< span class = "w" > < / span > use< span class = "w" > < / span > to< span class = "w" > < / span > say< span class = "w" > < / span > you< span class = "w" > < / span > would< span class = "w" > < / span > not< span class = "w" > < / span > understand,< span class = "w" > < / span > < span class = "k" > if< / span > < span class = "w" > < / span > I< span class = "w" > < / span > were< span class = "w" > < / span > to< span class = "w" > < / span > tell< span class = "w" > < / span > you,”< span class = "w" > < / span > said< span class = "w" > < / span > he.< span class = "w" > < / span > “I< span class = "w" > < / span > have< span class = "w" > < / span > not< span class = "w" > < / span > told< span class = "w" > < / span > you< span class = "w" > < / span > why< span class = "w" > < / span > I< span class = "w" > < / span > am< span class = "w" > < / span > waiting< span class = "w" > < / span > < span class = "k" > for< / span > < span class = "w" > < / span > him,”< span class = "w" > < / span > said< span class = "w" > < / span > I.< span class = "w" > < / span > “And< span class = "w" > < / span > I< span class = "w" > < / span > am< span class = "w" > < / span > sure< span class = "w" > < / span > I< span class = "w" > < / span > should< span class = "w" > < /span
2023-11-29 12:41:56 -08:00
------
2024-03-30 17:32:20 -07:00
< span class = "o" > [< / span > INFO< span class = "o" > ]< / span > < span class = "w" > < / span > Prompt< span class = "w" > < / span > processing:< span class = "w" > < / span > < span class = "m" > 0< / span > .633< span class = "w" > < / span > s
< span class = "o" > [< / span > INFO< span class = "o" > ]< / span > < span class = "w" > < / span > Full< span class = "w" > < / span > generation:< span class = "w" > < / span > < span class = "m" > 21< / span > .475< span class = "w" > < / span > s
2023-11-29 12:41:56 -08:00
< / pre > < / div >
< / div >
< / section >
< section id = "scripts" >
2024-03-30 17:32:20 -07:00
< h2 > Scripts< a class = "headerlink" href = "#scripts" title = "Link to this heading" > #< / a > < / h2 >
2023-11-29 12:41:56 -08:00
< div class = "admonition-download-the-code admonition" >
< p class = "admonition-title" > Download the code< / p >
2024-01-10 14:14:38 -08:00
< p > The full example code is available in < a class = "reference external" href = "https://github.com/ml-explore/mlx-examples/tree/main/llms/llama" > mlx-examples< / a > .< / p >
2023-11-29 12:41:56 -08:00
< / div >
2024-03-30 17:32:20 -07:00
< aside class = "footnote-list brackets" >
2023-12-06 08:13:20 -08:00
< aside class = "footnote brackets" id = "id4" role = "note" >
2023-11-29 12:41:56 -08:00
< span class = "label" > < span class = "fn-bracket" > [< / span > < a role = "doc-backlink" href = "#id1" > 1< / a > < span class = "fn-bracket" > ]< / span > < / span >
< p > Su, J., Lu, Y., Pan, S., Murtadha, A., Wen, B. and Liu, Y., 2021.
Roformer: Enhanced transformer with rotary position embedding. arXiv
preprint arXiv:2104.09864.< / p >
< / aside >
2023-12-06 08:13:20 -08:00
< aside class = "footnote brackets" id = "id5" role = "note" >
2023-11-29 12:41:56 -08:00
< span class = "label" > < span class = "fn-bracket" > [< / span > < a role = "doc-backlink" href = "#id2" > 2< / a > < span class = "fn-bracket" > ]< / span > < / span >
< p > Zhang, B. and Sennrich, R., 2019. Root mean square layer normalization.
Advances in Neural Information Processing Systems, 32.< / p >
< / aside >
2023-12-06 08:13:20 -08:00
< aside class = "footnote brackets" id = "id6" role = "note" >
2023-11-29 12:41:56 -08:00
< span class = "label" > < span class = "fn-bracket" > [< / span > < a role = "doc-backlink" href = "#id3" > 3< / a > < span class = "fn-bracket" > ]< / span > < / span >
< p > Shazeer, N., 2020. Glu variants improve transformer. arXiv preprint
arXiv:2002.05202.< / p >
< / aside >
2024-03-30 17:32:20 -07:00
< / aside >
2023-11-29 12:41:56 -08:00
< / section >
< / section >
2023-12-05 12:10:03 -08:00
< / article >
< footer class = "prev-next-footer" >
< div class = "prev-next-area" >
< a class = "left-prev"
href="mlp.html"
title="previous page">
< i class = "fa-solid fa-angle-left" > < / i >
< div class = "prev-next-info" >
< p class = "prev-next-subtitle" > previous< / p >
< p class = "prev-next-title" > Multi-Layer Perceptron< / p >
< / div >
< / a >
< a class = "right-next"
href="../python/array.html"
title="next page">
< div class = "prev-next-info" >
< p class = "prev-next-subtitle" > next< / p >
< p class = "prev-next-title" > Array< / p >
< / div >
< i class = "fa-solid fa-angle-right" > < / i >
< / a >
< / div >
< / footer >
< / div >
< div class = "bd-sidebar-secondary bd-toc" > < div class = "sidebar-secondary-items sidebar-secondary__inner" >
< div class = "sidebar-secondary-item" >
< div class = "page-toc tocsection onthispage" >
< i class = "fa-solid fa-list" > < / i > Contents
< / div >
< nav class = "bd-toc-nav page-toc" >
< ul class = "visible nav section-nav flex-column" >
< li class = "toc-h2 nav-item toc-entry" > < a class = "reference internal nav-link" href = "#implementing-the-model" > Implementing the model< / a > < ul class = "visible nav section-nav flex-column" >
< li class = "toc-h3 nav-item toc-entry" > < a class = "reference internal nav-link" href = "#attention-layer" > Attention layer< / a > < / li >
< li class = "toc-h3 nav-item toc-entry" > < a class = "reference internal nav-link" href = "#encoder-layer" > Encoder layer< / a > < / li >
< li class = "toc-h3 nav-item toc-entry" > < a class = "reference internal nav-link" href = "#full-model" > Full model< / a > < / li >
< li class = "toc-h3 nav-item toc-entry" > < a class = "reference internal nav-link" href = "#generation" > Generation< / a > < / li >
< li class = "toc-h3 nav-item toc-entry" > < a class = "reference internal nav-link" href = "#putting-it-all-together" > Putting it all together< / a > < / li >
< / ul >
< / li >
< li class = "toc-h2 nav-item toc-entry" > < a class = "reference internal nav-link" href = "#converting-the-weights" > Converting the weights< / a > < / li >
< li class = "toc-h2 nav-item toc-entry" > < a class = "reference internal nav-link" href = "#weight-loading-and-benchmarking" > Weight loading and benchmarking< / a > < / li >
< li class = "toc-h2 nav-item toc-entry" > < a class = "reference internal nav-link" href = "#scripts" > Scripts< / a > < / li >
< / ul >
< / nav > < / div >
2023-11-29 12:41:56 -08:00
2023-12-05 12:10:03 -08:00
< / div > < / div >
< / div >
< footer class = "bd-footer-content" >
< div class = "bd-footer-content__inner container" >
< div class = "footer-item" >
< p class = "component-author" >
By MLX Contributors
< / p >
2023-11-29 12:41:56 -08:00
< / div >
2023-12-05 12:10:03 -08:00
< div class = "footer-item" >
2023-11-29 12:41:56 -08:00
2023-12-05 12:10:03 -08:00
< p class = "copyright" >
© Copyright 2023, MLX Contributors.
< br / >
< / p >
2023-11-29 12:41:56 -08:00
< / div >
2023-12-05 12:10:03 -08:00
< div class = "footer-item" >
< / div >
< div class = "footer-item" >
< / div >
< / div >
< / footer >
< / main >
< / div >
< / div >
<!-- Scripts loaded after <body> so the DOM is not blocked -->
< script src = "../_static/scripts/bootstrap.js?digest=5b4479735964841361fd" > < / script >
< script src = "../_static/scripts/pydata-sphinx-theme.js?digest=5b4479735964841361fd" > < / script >
2023-11-29 12:41:56 -08:00
2023-12-05 12:10:03 -08:00
< footer class = "bd-footer" >
< / footer >
< / body >
2023-11-29 12:41:56 -08:00
< / html >
