zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-06-25 18:11:15 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity
a08f30df65
mlx/docs/build/html/python/nn.html
CircleCI Docs a08f30df65 rebase
2025-06-04 01:01:56 +00:00

1700 lines
147 KiB
HTML
Raw Blame History

This file contains invisible Unicode characters

This file contains invisible Unicode characters that are indistinguishable to humans but may be processed differently by a computer. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

<!DOCTYPE html>
<html lang="en" data-content_root="../" >
<head>
<meta charset="utf-8" />
<meta name="viewport" content="width=device-width, initial-scale=1.0" /><meta name="viewport" content="width=device-width, initial-scale=1" />
<title>Neural Networks &#8212; MLX 0.24.0 documentation</title>
<script data-cfasync="false">
document.documentElement.dataset.mode = localStorage.getItem("mode") || "";
document.documentElement.dataset.theme = localStorage.getItem("theme") || "";
</script>
<!-- Loaded before other Sphinx assets -->
<link href="../_static/styles/theme.css?digest=dfe6caa3a7d634c4db9b" rel="stylesheet" />
<link href="../_static/styles/bootstrap.css?digest=dfe6caa3a7d634c4db9b" rel="stylesheet" />
<link href="../_static/styles/pydata-sphinx-theme.css?digest=dfe6caa3a7d634c4db9b" rel="stylesheet" />
<link href="../_static/vendor/fontawesome/6.5.2/css/all.min.css?digest=dfe6caa3a7d634c4db9b" rel="stylesheet" />
<link rel="preload" as="font" type="font/woff2" crossorigin href="../_static/vendor/fontawesome/6.5.2/webfonts/fa-solid-900.woff2" />
<link rel="preload" as="font" type="font/woff2" crossorigin href="../_static/vendor/fontawesome/6.5.2/webfonts/fa-brands-400.woff2" />
<link rel="preload" as="font" type="font/woff2" crossorigin href="../_static/vendor/fontawesome/6.5.2/webfonts/fa-regular-400.woff2" />
<link rel="stylesheet" type="text/css" href="../_static/pygments.css?v=03e43079" />
<link rel="stylesheet" type="text/css" href="../_static/styles/sphinx-book-theme.css?v=eba8b062" />
<!-- Pre-loaded scripts that we'll load fully later -->
<link rel="preload" as="script" href="../_static/scripts/bootstrap.js?digest=dfe6caa3a7d634c4db9b" />
<link rel="preload" as="script" href="../_static/scripts/pydata-sphinx-theme.js?digest=dfe6caa3a7d634c4db9b" />
<script src="../_static/vendor/fontawesome/6.5.2/js/all.min.js?digest=dfe6caa3a7d634c4db9b"></script>
<script src="../_static/documentation_options.js?v=ae1d10b0"></script>
<script src="../_static/doctools.js?v=9a2dae69"></script>
<script src="../_static/sphinx_highlight.js?v=dc90522c"></script>
<script src="../_static/scripts/sphinx-book-theme.js?v=887ef09a"></script>
<script async="async" src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js"></script>
<script>DOCUMENTATION_OPTIONS.pagename = 'python/nn';</script>
<link rel="icon" href="../_static/mlx_logo.png"/>
<link rel="index" title="Index" href="../genindex.html" />
<link rel="search" title="Search" href="../search.html" />
<link rel="next" title="mlx.nn.value_and_grad" href="_autosummary/mlx.nn.value_and_grad.html" />
<link rel="prev" title="mlx.core.metal.stop_capture" href="_autosummary/mlx.core.metal.stop_capture.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="">
<div id="pst-skip-link" class="skip-link d-print-none"><a href="#main-content">Skip to main content</a></div>
<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"
id="pst-primary-sidebar-checkbox"/>
<label class="overlay overlay-primary" for="pst-primary-sidebar-checkbox"></label>
<input type="checkbox"
class="sidebar-toggle"
id="pst-secondary-sidebar-checkbox"/>
<label class="overlay overlay-secondary" for="pst-secondary-sidebar-checkbox"></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>
<div class="pst-async-banner-revealer d-none">
<aside id="bd-header-version-warning" class="d-none d-print-none" aria-label="Version warning"></aside>
</div>
<header class="bd-header navbar navbar-expand-lg bd-navbar d-print-none">
</header>
<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">
<img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.24.0 documentation - Home"/>
<script>document.write(`<img src="../_static/mlx_logo_dark.png" class="logo__image only-dark" alt="MLX 0.24.0 documentation - Home"/>`);</script>
</a></div>
<div class="sidebar-primary-item">
<script>
document.write(`
<button class="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">
<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">
<li class="toctree-l1"><a class="reference internal" href="../install.html">Build and Install</a></li>
</ul>
<p aria-level="2" class="caption" role="heading"><span class="caption-text">Usage</span></p>
<ul class="nav bd-sidenav">
<li class="toctree-l1"><a class="reference internal" href="../usage/quick_start.html">Quick Start Guide</a></li>
<li class="toctree-l1"><a class="reference internal" href="../usage/lazy_evaluation.html">Lazy Evaluation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../usage/unified_memory.html">Unified Memory</a></li>
<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>
<li class="toctree-l1"><a class="reference internal" href="../usage/function_transforms.html">Function Transforms</a></li>
<li class="toctree-l1"><a class="reference internal" href="../usage/compile.html">Compilation</a></li>
<li class="toctree-l1"><a class="reference internal" href="../usage/numpy.html">Conversion to NumPy and Other Frameworks</a></li>
<li class="toctree-l1"><a class="reference internal" href="../usage/distributed.html">Distributed Communication</a></li>
<li class="toctree-l1"><a class="reference internal" href="../usage/using_streams.html">Using Streams</a></li>
<li class="toctree-l1"><a class="reference internal" href="../usage/export.html">Exporting Functions</a></li>
</ul>
<p aria-level="2" class="caption" role="heading"><span class="caption-text">Examples</span></p>
<ul class="nav bd-sidenav">
<li class="toctree-l1"><a class="reference internal" href="../examples/linear_regression.html">Linear Regression</a></li>
<li class="toctree-l1"><a class="reference internal" href="../examples/mlp.html">Multi-Layer Perceptron</a></li>
<li class="toctree-l1"><a class="reference internal" href="../examples/llama-inference.html">LLM inference</a></li>
</ul>
<p aria-level="2" class="caption" role="heading"><span class="caption-text">Python API Reference</span></p>
<ul class="current nav bd-sidenav">
<li class="toctree-l1 has-children"><a class="reference internal" href="array.html">Array</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.html">mlx.core.array</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.astype.html">mlx.core.array.astype</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.at.html">mlx.core.array.at</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.item.html">mlx.core.array.item</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.tolist.html">mlx.core.array.tolist</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.dtype.html">mlx.core.array.dtype</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.itemsize.html">mlx.core.array.itemsize</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.nbytes.html">mlx.core.array.nbytes</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.ndim.html">mlx.core.array.ndim</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.shape.html">mlx.core.array.shape</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.size.html">mlx.core.array.size</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.abs.html">mlx.core.array.abs</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.all.html">mlx.core.array.all</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.any.html">mlx.core.array.any</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.argmax.html">mlx.core.array.argmax</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.argmin.html">mlx.core.array.argmin</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.conj.html">mlx.core.array.conj</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.cos.html">mlx.core.array.cos</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.cummax.html">mlx.core.array.cummax</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.cummin.html">mlx.core.array.cummin</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.cumprod.html">mlx.core.array.cumprod</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.cumsum.html">mlx.core.array.cumsum</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.diag.html">mlx.core.array.diag</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.diagonal.html">mlx.core.array.diagonal</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.exp.html">mlx.core.array.exp</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.flatten.html">mlx.core.array.flatten</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.log.html">mlx.core.array.log</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.log10.html">mlx.core.array.log10</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.log1p.html">mlx.core.array.log1p</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.log2.html">mlx.core.array.log2</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.logsumexp.html">mlx.core.array.logsumexp</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.max.html">mlx.core.array.max</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.mean.html">mlx.core.array.mean</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.min.html">mlx.core.array.min</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.moveaxis.html">mlx.core.array.moveaxis</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.prod.html">mlx.core.array.prod</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.reciprocal.html">mlx.core.array.reciprocal</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.reshape.html">mlx.core.array.reshape</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.round.html">mlx.core.array.round</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.rsqrt.html">mlx.core.array.rsqrt</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.sin.html">mlx.core.array.sin</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.split.html">mlx.core.array.split</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.sqrt.html">mlx.core.array.sqrt</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.square.html">mlx.core.array.square</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.squeeze.html">mlx.core.array.squeeze</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.std.html">mlx.core.array.std</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.sum.html">mlx.core.array.sum</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.swapaxes.html">mlx.core.array.swapaxes</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.transpose.html">mlx.core.array.transpose</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.T.html">mlx.core.array.T</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.var.html">mlx.core.array.var</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array.view.html">mlx.core.array.view</a></li>
</ul>
</details></li>
<li class="toctree-l1 has-children"><a class="reference internal" href="data_types.html">Data Types</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.Dtype.html">mlx.core.Dtype</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.DtypeCategory.html">mlx.core.DtypeCategory</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.issubdtype.html">mlx.core.issubdtype</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.finfo.html">mlx.core.finfo</a></li>
</ul>
</details></li>
<li class="toctree-l1 has-children"><a class="reference internal" href="devices_and_streams.html">Devices and Streams</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.Device.html">mlx.core.Device</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/stream_class.html">mlx.core.Stream</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.default_device.html">mlx.core.default_device</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.set_default_device.html">mlx.core.set_default_device</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.default_stream.html">mlx.core.default_stream</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.new_stream.html">mlx.core.new_stream</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.set_default_stream.html">mlx.core.set_default_stream</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.stream.html">mlx.core.stream</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.synchronize.html">mlx.core.synchronize</a></li>
</ul>
</details></li>
<li class="toctree-l1 has-children"><a class="reference internal" href="export.html">Export Functions</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.export_function.html">mlx.core.export_function</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.import_function.html">mlx.core.import_function</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.exporter.html">mlx.core.exporter</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.export_to_dot.html">mlx.core.export_to_dot</a></li>
</ul>
</details></li>
<li class="toctree-l1 has-children"><a class="reference internal" href="ops.html">Operations</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.abs.html">mlx.core.abs</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.add.html">mlx.core.add</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.addmm.html">mlx.core.addmm</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.all.html">mlx.core.all</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.allclose.html">mlx.core.allclose</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.any.html">mlx.core.any</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.arange.html">mlx.core.arange</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.arccos.html">mlx.core.arccos</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.arccosh.html">mlx.core.arccosh</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.arcsin.html">mlx.core.arcsin</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.arcsinh.html">mlx.core.arcsinh</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.arctan.html">mlx.core.arctan</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.arctan2.html">mlx.core.arctan2</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.arctanh.html">mlx.core.arctanh</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.argmax.html">mlx.core.argmax</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.argmin.html">mlx.core.argmin</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.argpartition.html">mlx.core.argpartition</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.argsort.html">mlx.core.argsort</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.array_equal.html">mlx.core.array_equal</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.as_strided.html">mlx.core.as_strided</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.atleast_1d.html">mlx.core.atleast_1d</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.atleast_2d.html">mlx.core.atleast_2d</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.atleast_3d.html">mlx.core.atleast_3d</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.bitwise_and.html">mlx.core.bitwise_and</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.bitwise_invert.html">mlx.core.bitwise_invert</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.bitwise_or.html">mlx.core.bitwise_or</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.bitwise_xor.html">mlx.core.bitwise_xor</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.block_masked_mm.html">mlx.core.block_masked_mm</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.broadcast_to.html">mlx.core.broadcast_to</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.ceil.html">mlx.core.ceil</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.clip.html">mlx.core.clip</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.concatenate.html">mlx.core.concatenate</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.conj.html">mlx.core.conj</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.conjugate.html">mlx.core.conjugate</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.convolve.html">mlx.core.convolve</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.conv1d.html">mlx.core.conv1d</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.conv2d.html">mlx.core.conv2d</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.conv3d.html">mlx.core.conv3d</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.conv_transpose1d.html">mlx.core.conv_transpose1d</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.conv_transpose2d.html">mlx.core.conv_transpose2d</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.conv_transpose3d.html">mlx.core.conv_transpose3d</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.conv_general.html">mlx.core.conv_general</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.cos.html">mlx.core.cos</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.cosh.html">mlx.core.cosh</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.cummax.html">mlx.core.cummax</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.cummin.html">mlx.core.cummin</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.cumprod.html">mlx.core.cumprod</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.cumsum.html">mlx.core.cumsum</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.degrees.html">mlx.core.degrees</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.dequantize.html">mlx.core.dequantize</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.diag.html">mlx.core.diag</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.diagonal.html">mlx.core.diagonal</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.divide.html">mlx.core.divide</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.divmod.html">mlx.core.divmod</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.einsum.html">mlx.core.einsum</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.einsum_path.html">mlx.core.einsum_path</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.equal.html">mlx.core.equal</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.erf.html">mlx.core.erf</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.erfinv.html">mlx.core.erfinv</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.exp.html">mlx.core.exp</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.expm1.html">mlx.core.expm1</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.expand_dims.html">mlx.core.expand_dims</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.eye.html">mlx.core.eye</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.flatten.html">mlx.core.flatten</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.floor.html">mlx.core.floor</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.floor_divide.html">mlx.core.floor_divide</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.full.html">mlx.core.full</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.gather_mm.html">mlx.core.gather_mm</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.gather_qmm.html">mlx.core.gather_qmm</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.greater.html">mlx.core.greater</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.greater_equal.html">mlx.core.greater_equal</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.hadamard_transform.html">mlx.core.hadamard_transform</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.identity.html">mlx.core.identity</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.imag.html">mlx.core.imag</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.inner.html">mlx.core.inner</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.isfinite.html">mlx.core.isfinite</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.isclose.html">mlx.core.isclose</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.isinf.html">mlx.core.isinf</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.isnan.html">mlx.core.isnan</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.isneginf.html">mlx.core.isneginf</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.isposinf.html">mlx.core.isposinf</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.issubdtype.html">mlx.core.issubdtype</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.kron.html">mlx.core.kron</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.left_shift.html">mlx.core.left_shift</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.less.html">mlx.core.less</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.less_equal.html">mlx.core.less_equal</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.linspace.html">mlx.core.linspace</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.load.html">mlx.core.load</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.log.html">mlx.core.log</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.log2.html">mlx.core.log2</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.log10.html">mlx.core.log10</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.log1p.html">mlx.core.log1p</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.logaddexp.html">mlx.core.logaddexp</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.logical_not.html">mlx.core.logical_not</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.logical_and.html">mlx.core.logical_and</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.logical_or.html">mlx.core.logical_or</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.logsumexp.html">mlx.core.logsumexp</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.matmul.html">mlx.core.matmul</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.max.html">mlx.core.max</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.maximum.html">mlx.core.maximum</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.mean.html">mlx.core.mean</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.meshgrid.html">mlx.core.meshgrid</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.min.html">mlx.core.min</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.minimum.html">mlx.core.minimum</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.moveaxis.html">mlx.core.moveaxis</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.multiply.html">mlx.core.multiply</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.nan_to_num.html">mlx.core.nan_to_num</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.negative.html">mlx.core.negative</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.not_equal.html">mlx.core.not_equal</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.ones.html">mlx.core.ones</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.ones_like.html">mlx.core.ones_like</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.outer.html">mlx.core.outer</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.partition.html">mlx.core.partition</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.pad.html">mlx.core.pad</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.power.html">mlx.core.power</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.prod.html">mlx.core.prod</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.put_along_axis.html">mlx.core.put_along_axis</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.quantize.html">mlx.core.quantize</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.quantized_matmul.html">mlx.core.quantized_matmul</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.radians.html">mlx.core.radians</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.real.html">mlx.core.real</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.reciprocal.html">mlx.core.reciprocal</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.remainder.html">mlx.core.remainder</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.repeat.html">mlx.core.repeat</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.reshape.html">mlx.core.reshape</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.right_shift.html">mlx.core.right_shift</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.roll.html">mlx.core.roll</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.round.html">mlx.core.round</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.rsqrt.html">mlx.core.rsqrt</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.save.html">mlx.core.save</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.savez.html">mlx.core.savez</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.savez_compressed.html">mlx.core.savez_compressed</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.save_gguf.html">mlx.core.save_gguf</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.save_safetensors.html">mlx.core.save_safetensors</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.sigmoid.html">mlx.core.sigmoid</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.sign.html">mlx.core.sign</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.sin.html">mlx.core.sin</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.sinh.html">mlx.core.sinh</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.slice.html">mlx.core.slice</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.slice_update.html">mlx.core.slice_update</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.softmax.html">mlx.core.softmax</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.sort.html">mlx.core.sort</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.split.html">mlx.core.split</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.sqrt.html">mlx.core.sqrt</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.square.html">mlx.core.square</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.squeeze.html">mlx.core.squeeze</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.stack.html">mlx.core.stack</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.std.html">mlx.core.std</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.stop_gradient.html">mlx.core.stop_gradient</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.subtract.html">mlx.core.subtract</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.sum.html">mlx.core.sum</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.swapaxes.html">mlx.core.swapaxes</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.take.html">mlx.core.take</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.take_along_axis.html">mlx.core.take_along_axis</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.tan.html">mlx.core.tan</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.tanh.html">mlx.core.tanh</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.tensordot.html">mlx.core.tensordot</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.tile.html">mlx.core.tile</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.topk.html">mlx.core.topk</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.trace.html">mlx.core.trace</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.transpose.html">mlx.core.transpose</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.tri.html">mlx.core.tri</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.tril.html">mlx.core.tril</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.triu.html">mlx.core.triu</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.unflatten.html">mlx.core.unflatten</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.var.html">mlx.core.var</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.view.html">mlx.core.view</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.where.html">mlx.core.where</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.zeros.html">mlx.core.zeros</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.zeros_like.html">mlx.core.zeros_like</a></li>
</ul>
</details></li>
<li class="toctree-l1 has-children"><a class="reference internal" href="random.html">Random</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.random.bernoulli.html">mlx.core.random.bernoulli</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.random.categorical.html">mlx.core.random.categorical</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.random.gumbel.html">mlx.core.random.gumbel</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.random.key.html">mlx.core.random.key</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.random.normal.html">mlx.core.random.normal</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.random.multivariate_normal.html">mlx.core.random.multivariate_normal</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.random.randint.html">mlx.core.random.randint</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.random.seed.html">mlx.core.random.seed</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.random.split.html">mlx.core.random.split</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.random.truncated_normal.html">mlx.core.random.truncated_normal</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.random.uniform.html">mlx.core.random.uniform</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.random.laplace.html">mlx.core.random.laplace</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.random.permutation.html">mlx.core.random.permutation</a></li>
</ul>
</details></li>
<li class="toctree-l1 has-children"><a class="reference internal" href="transforms.html">Transforms</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.eval.html">mlx.core.eval</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.compile.html">mlx.core.compile</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.custom_function.html">mlx.core.custom_function</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.disable_compile.html">mlx.core.disable_compile</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.enable_compile.html">mlx.core.enable_compile</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.grad.html">mlx.core.grad</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.value_and_grad.html">mlx.core.value_and_grad</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.jvp.html">mlx.core.jvp</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.vjp.html">mlx.core.vjp</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.vmap.html">mlx.core.vmap</a></li>
</ul>
</details></li>
<li class="toctree-l1 has-children"><a class="reference internal" href="fast.html">Fast</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.fast.rms_norm.html">mlx.core.fast.rms_norm</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.fast.layer_norm.html">mlx.core.fast.layer_norm</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.fast.rope.html">mlx.core.fast.rope</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.fast.scaled_dot_product_attention.html">mlx.core.fast.scaled_dot_product_attention</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.fast.metal_kernel.html">mlx.core.fast.metal_kernel</a></li>
</ul>
</details></li>
<li class="toctree-l1 has-children"><a class="reference internal" href="fft.html">FFT</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.fft.fft.html">mlx.core.fft.fft</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.fft.ifft.html">mlx.core.fft.ifft</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.fft.fft2.html">mlx.core.fft.fft2</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.fft.ifft2.html">mlx.core.fft.ifft2</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.fft.fftn.html">mlx.core.fft.fftn</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.fft.ifftn.html">mlx.core.fft.ifftn</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.fft.rfft.html">mlx.core.fft.rfft</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.fft.irfft.html">mlx.core.fft.irfft</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.fft.rfft2.html">mlx.core.fft.rfft2</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.fft.irfft2.html">mlx.core.fft.irfft2</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.fft.rfftn.html">mlx.core.fft.rfftn</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.fft.irfftn.html">mlx.core.fft.irfftn</a></li>
</ul>
</details></li>
<li class="toctree-l1 has-children"><a class="reference internal" href="linalg.html">Linear Algebra</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.linalg.inv.html">mlx.core.linalg.inv</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.linalg.tri_inv.html">mlx.core.linalg.tri_inv</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.linalg.norm.html">mlx.core.linalg.norm</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.linalg.cholesky.html">mlx.core.linalg.cholesky</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.linalg.cholesky_inv.html">mlx.core.linalg.cholesky_inv</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.linalg.cross.html">mlx.core.linalg.cross</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.linalg.qr.html">mlx.core.linalg.qr</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.linalg.svd.html">mlx.core.linalg.svd</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.linalg.eigvalsh.html">mlx.core.linalg.eigvalsh</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.linalg.eigh.html">mlx.core.linalg.eigh</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.linalg.lu.html">mlx.core.linalg.lu</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.linalg.lu_factor.html">mlx.core.linalg.lu_factor</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.linalg.solve.html">mlx.core.linalg.solve</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.linalg.solve_triangular.html">mlx.core.linalg.solve_triangular</a></li>
</ul>
</details></li>
<li class="toctree-l1 has-children"><a class="reference internal" href="metal.html">Metal</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.metal.is_available.html">mlx.core.metal.is_available</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.metal.device_info.html">mlx.core.metal.device_info</a></li>
<li class="toctree-l2"><a class="reference internal" href="_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="_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="_autosummary/mlx.core.metal.reset_peak_memory.html">mlx.core.metal.reset_peak_memory</a></li>
<li class="toctree-l2"><a class="reference internal" href="_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="_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="_autosummary/mlx.core.metal.set_cache_limit.html">mlx.core.metal.set_cache_limit</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.metal.set_wired_limit.html">mlx.core.metal.set_wired_limit</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.metal.clear_cache.html">mlx.core.metal.clear_cache</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.metal.start_capture.html">mlx.core.metal.start_capture</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.metal.stop_capture.html">mlx.core.metal.stop_capture</a></li>
</ul>
</details></li>
<li class="toctree-l1 current active has-children"><a class="current reference internal" href="#">Neural Networks</a><details open="open"><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.nn.value_and_grad.html">mlx.nn.value_and_grad</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.nn.quantize.html">mlx.nn.quantize</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.nn.average_gradients.html">mlx.nn.average_gradients</a></li>
<li class="toctree-l2 has-children"><a class="reference internal" href="nn/module.html">Module</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.training.html">mlx.nn.Module.training</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.state.html">mlx.nn.Module.state</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.apply.html">mlx.nn.Module.apply</a></li>
<li class="toctree-l3"><a class="reference internal" href="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="nn/_autosummary/mlx.nn.Module.children.html">mlx.nn.Module.children</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.eval.html">mlx.nn.Module.eval</a></li>
<li class="toctree-l3"><a class="reference internal" href="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="nn/_autosummary/mlx.nn.Module.freeze.html">mlx.nn.Module.freeze</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.leaf_modules.html">mlx.nn.Module.leaf_modules</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.load_weights.html">mlx.nn.Module.load_weights</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.modules.html">mlx.nn.Module.modules</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.named_modules.html">mlx.nn.Module.named_modules</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.parameters.html">mlx.nn.Module.parameters</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.save_weights.html">mlx.nn.Module.save_weights</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.set_dtype.html">mlx.nn.Module.set_dtype</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.train.html">mlx.nn.Module.train</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.trainable_parameters.html">mlx.nn.Module.trainable_parameters</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.unfreeze.html">mlx.nn.Module.unfreeze</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.update.html">mlx.nn.Module.update</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.update_modules.html">mlx.nn.Module.update_modules</a></li>
</ul>
</details></li>
<li class="toctree-l2 has-children"><a class="reference internal" href="nn/layers.html">Layers</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.ALiBi.html">mlx.nn.ALiBi</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.AvgPool1d.html">mlx.nn.AvgPool1d</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.AvgPool2d.html">mlx.nn.AvgPool2d</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.AvgPool3d.html">mlx.nn.AvgPool3d</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.BatchNorm.html">mlx.nn.BatchNorm</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.CELU.html">mlx.nn.CELU</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Conv1d.html">mlx.nn.Conv1d</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Conv2d.html">mlx.nn.Conv2d</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Conv3d.html">mlx.nn.Conv3d</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.ConvTranspose1d.html">mlx.nn.ConvTranspose1d</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.ConvTranspose2d.html">mlx.nn.ConvTranspose2d</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.ConvTranspose3d.html">mlx.nn.ConvTranspose3d</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Dropout.html">mlx.nn.Dropout</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Dropout2d.html">mlx.nn.Dropout2d</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Dropout3d.html">mlx.nn.Dropout3d</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Embedding.html">mlx.nn.Embedding</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.ELU.html">mlx.nn.ELU</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.GELU.html">mlx.nn.GELU</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.GLU.html">mlx.nn.GLU</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.GroupNorm.html">mlx.nn.GroupNorm</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.GRU.html">mlx.nn.GRU</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.HardShrink.html">mlx.nn.HardShrink</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.HardTanh.html">mlx.nn.HardTanh</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Hardswish.html">mlx.nn.Hardswish</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.InstanceNorm.html">mlx.nn.InstanceNorm</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.LayerNorm.html">mlx.nn.LayerNorm</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.LeakyReLU.html">mlx.nn.LeakyReLU</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Linear.html">mlx.nn.Linear</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.LogSigmoid.html">mlx.nn.LogSigmoid</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.LogSoftmax.html">mlx.nn.LogSoftmax</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.LSTM.html">mlx.nn.LSTM</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.MaxPool1d.html">mlx.nn.MaxPool1d</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.MaxPool2d.html">mlx.nn.MaxPool2d</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.MaxPool3d.html">mlx.nn.MaxPool3d</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Mish.html">mlx.nn.Mish</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.MultiHeadAttention.html">mlx.nn.MultiHeadAttention</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.PReLU.html">mlx.nn.PReLU</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.QuantizedEmbedding.html">mlx.nn.QuantizedEmbedding</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.QuantizedLinear.html">mlx.nn.QuantizedLinear</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.RMSNorm.html">mlx.nn.RMSNorm</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.ReLU.html">mlx.nn.ReLU</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.ReLU6.html">mlx.nn.ReLU6</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.RNN.html">mlx.nn.RNN</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.RoPE.html">mlx.nn.RoPE</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.SELU.html">mlx.nn.SELU</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Sequential.html">mlx.nn.Sequential</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Sigmoid.html">mlx.nn.Sigmoid</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.SiLU.html">mlx.nn.SiLU</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.SinusoidalPositionalEncoding.html">mlx.nn.SinusoidalPositionalEncoding</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Softmin.html">mlx.nn.Softmin</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Softshrink.html">mlx.nn.Softshrink</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Softsign.html">mlx.nn.Softsign</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Softmax.html">mlx.nn.Softmax</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Softplus.html">mlx.nn.Softplus</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Step.html">mlx.nn.Step</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Tanh.html">mlx.nn.Tanh</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Transformer.html">mlx.nn.Transformer</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Upsample.html">mlx.nn.Upsample</a></li>
</ul>
</details></li>
<li class="toctree-l2 has-children"><a class="reference internal" href="nn/functions.html">Functions</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.elu.html">mlx.nn.elu</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.celu.html">mlx.nn.celu</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.gelu.html">mlx.nn.gelu</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.gelu_approx.html">mlx.nn.gelu_approx</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.gelu_fast_approx.html">mlx.nn.gelu_fast_approx</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.glu.html">mlx.nn.glu</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.hard_shrink.html">mlx.nn.hard_shrink</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.hard_tanh.html">mlx.nn.hard_tanh</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.hardswish.html">mlx.nn.hardswish</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.leaky_relu.html">mlx.nn.leaky_relu</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.log_sigmoid.html">mlx.nn.log_sigmoid</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.log_softmax.html">mlx.nn.log_softmax</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.mish.html">mlx.nn.mish</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.prelu.html">mlx.nn.prelu</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.relu.html">mlx.nn.relu</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.relu6.html">mlx.nn.relu6</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.selu.html">mlx.nn.selu</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.sigmoid.html">mlx.nn.sigmoid</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.silu.html">mlx.nn.silu</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.softmax.html">mlx.nn.softmax</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.softmin.html">mlx.nn.softmin</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.softplus.html">mlx.nn.softplus</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.softshrink.html">mlx.nn.softshrink</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.step.html">mlx.nn.step</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.tanh.html">mlx.nn.tanh</a></li>
</ul>
</details></li>
<li class="toctree-l2 has-children"><a class="reference internal" href="nn/losses.html">Loss Functions</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.binary_cross_entropy.html">mlx.nn.losses.binary_cross_entropy</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.cosine_similarity_loss.html">mlx.nn.losses.cosine_similarity_loss</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.cross_entropy.html">mlx.nn.losses.cross_entropy</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.gaussian_nll_loss.html">mlx.nn.losses.gaussian_nll_loss</a></li>
<li class="toctree-l3"><a class="reference internal" href="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="nn/_autosummary_functions/mlx.nn.losses.huber_loss.html">mlx.nn.losses.huber_loss</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.kl_div_loss.html">mlx.nn.losses.kl_div_loss</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.l1_loss.html">mlx.nn.losses.l1_loss</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.log_cosh_loss.html">mlx.nn.losses.log_cosh_loss</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.margin_ranking_loss.html">mlx.nn.losses.margin_ranking_loss</a></li>
<li class="toctree-l3"><a class="reference internal" href="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="nn/_autosummary_functions/mlx.nn.losses.nll_loss.html">mlx.nn.losses.nll_loss</a></li>
<li class="toctree-l3"><a class="reference internal" href="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="nn/_autosummary_functions/mlx.nn.losses.triplet_loss.html">mlx.nn.losses.triplet_loss</a></li>
</ul>
</details></li>
<li class="toctree-l2 has-children"><a class="reference internal" href="nn/init.html">Initializers</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.init.constant.html">mlx.nn.init.constant</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.init.normal.html">mlx.nn.init.normal</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.init.uniform.html">mlx.nn.init.uniform</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.init.identity.html">mlx.nn.init.identity</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.init.glorot_normal.html">mlx.nn.init.glorot_normal</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.init.glorot_uniform.html">mlx.nn.init.glorot_uniform</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.init.he_normal.html">mlx.nn.init.he_normal</a></li>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.init.he_uniform.html">mlx.nn.init.he_uniform</a></li>
</ul>
</details></li>
</ul>
</details></li>
<li class="toctree-l1 has-children"><a class="reference internal" href="optimizers.html">Optimizers</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul>
<li class="toctree-l2 has-children"><a class="reference internal" href="optimizers/optimizer.html">Optimizer</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul>
<li class="toctree-l3"><a class="reference internal" href="optimizers/_autosummary/mlx.optimizers.Optimizer.state.html">mlx.optimizers.Optimizer.state</a></li>
<li class="toctree-l3"><a class="reference internal" href="optimizers/_autosummary/mlx.optimizers.Optimizer.apply_gradients.html">mlx.optimizers.Optimizer.apply_gradients</a></li>
<li class="toctree-l3"><a class="reference internal" href="optimizers/_autosummary/mlx.optimizers.Optimizer.init.html">mlx.optimizers.Optimizer.init</a></li>
<li class="toctree-l3"><a class="reference internal" href="optimizers/_autosummary/mlx.optimizers.Optimizer.update.html">mlx.optimizers.Optimizer.update</a></li>
</ul>
</details></li>
<li class="toctree-l2 has-children"><a class="reference internal" href="optimizers/common_optimizers.html">Common Optimizers</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul>
<li class="toctree-l3"><a class="reference internal" href="optimizers/_autosummary/mlx.optimizers.SGD.html">mlx.optimizers.SGD</a></li>
<li class="toctree-l3"><a class="reference internal" href="optimizers/_autosummary/mlx.optimizers.RMSprop.html">mlx.optimizers.RMSprop</a></li>
<li class="toctree-l3"><a class="reference internal" href="optimizers/_autosummary/mlx.optimizers.Adagrad.html">mlx.optimizers.Adagrad</a></li>
<li class="toctree-l3"><a class="reference internal" href="optimizers/_autosummary/mlx.optimizers.Adafactor.html">mlx.optimizers.Adafactor</a></li>
<li class="toctree-l3"><a class="reference internal" href="optimizers/_autosummary/mlx.optimizers.AdaDelta.html">mlx.optimizers.AdaDelta</a></li>
<li class="toctree-l3"><a class="reference internal" href="optimizers/_autosummary/mlx.optimizers.Adam.html">mlx.optimizers.Adam</a></li>
<li class="toctree-l3"><a class="reference internal" href="optimizers/_autosummary/mlx.optimizers.AdamW.html">mlx.optimizers.AdamW</a></li>
<li class="toctree-l3"><a class="reference internal" href="optimizers/_autosummary/mlx.optimizers.Adamax.html">mlx.optimizers.Adamax</a></li>
<li class="toctree-l3"><a class="reference internal" href="optimizers/_autosummary/mlx.optimizers.Lion.html">mlx.optimizers.Lion</a></li>
</ul>
</details></li>
<li class="toctree-l2 has-children"><a class="reference internal" href="optimizers/schedulers.html">Schedulers</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul>
<li class="toctree-l3"><a class="reference internal" href="optimizers/_autosummary/mlx.optimizers.cosine_decay.html">mlx.optimizers.cosine_decay</a></li>
<li class="toctree-l3"><a class="reference internal" href="optimizers/_autosummary/mlx.optimizers.exponential_decay.html">mlx.optimizers.exponential_decay</a></li>
<li class="toctree-l3"><a class="reference internal" href="optimizers/_autosummary/mlx.optimizers.join_schedules.html">mlx.optimizers.join_schedules</a></li>
<li class="toctree-l3"><a class="reference internal" href="optimizers/_autosummary/mlx.optimizers.linear_schedule.html">mlx.optimizers.linear_schedule</a></li>
<li class="toctree-l3"><a class="reference internal" href="optimizers/_autosummary/mlx.optimizers.step_decay.html">mlx.optimizers.step_decay</a></li>
</ul>
</details></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.optimizers.clip_grad_norm.html">mlx.optimizers.clip_grad_norm</a></li>
</ul>
</details></li>
<li class="toctree-l1 has-children"><a class="reference internal" href="distributed.html">Distributed Communication</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.distributed.Group.html">mlx.core.distributed.Group</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.distributed.is_available.html">mlx.core.distributed.is_available</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.distributed.init.html">mlx.core.distributed.init</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.distributed.all_sum.html">mlx.core.distributed.all_sum</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.distributed.all_gather.html">mlx.core.distributed.all_gather</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.distributed.send.html">mlx.core.distributed.send</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.distributed.recv.html">mlx.core.distributed.recv</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.core.distributed.recv_like.html">mlx.core.distributed.recv_like</a></li>
</ul>
</details></li>
<li class="toctree-l1 has-children"><a class="reference internal" href="tree_utils.html">Tree Utils</a><details><summary><span class="toctree-toggle" role="presentation"><i class="fa-solid fa-chevron-down"></i></span></summary><ul>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.utils.tree_flatten.html">mlx.utils.tree_flatten</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.utils.tree_unflatten.html">mlx.utils.tree_unflatten</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.utils.tree_map.html">mlx.utils.tree_map</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.utils.tree_map_with_path.html">mlx.utils.tree_map_with_path</a></li>
<li class="toctree-l2"><a class="reference internal" href="_autosummary/mlx.utils.tree_reduce.html">mlx.utils.tree_reduce</a></li>
</ul>
</details></li>
</ul>
<p aria-level="2" class="caption" role="heading"><span class="caption-text">C++ API Reference</span></p>
<ul class="nav bd-sidenav">
<li class="toctree-l1"><a class="reference internal" href="../cpp/ops.html">Operations</a></li>
</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">Custom Extensions in MLX</a></li>
<li class="toctree-l1"><a class="reference internal" href="../dev/metal_debugger.html">Metal Debugger</a></li>
<li class="toctree-l1"><a class="reference internal" href="../dev/custom_metal_kernels.html">Custom Metal Kernels</a></li>
<li class="toctree-l1"><a class="reference internal" href="../dev/mlx_in_cpp.html">Using MLX in C++</a></li>
</ul>
</div>
</nav></div>
</div>
<div class="sidebar-primary-items__end sidebar-primary__section">
</div>
<div id="rtd-footer-container"></div>
</div>
<main id="main-content" class="bd-main" role="main">
<div class="sbt-scroll-pixel-helper"></div>
<div class="bd-content">
<div class="bd-article-container">
<div class="bd-header-article d-print-none">
<div class="header-article-items header-article__inner">
<div class="header-article-items__start">
<div class="header-article-item"><button class="sidebar-toggle primary-toggle btn btn-sm" title="Toggle primary sidebar" data-bs-placement="bottom" data-bs-toggle="tooltip">
<span class="fa-solid fa-bars"></span>
</button></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/python/nn.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>
</ul>
</div>
<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 nav-link pst-navbar-icon theme-switch-button" title="light/dark" aria-label="light/dark" data-bs-placement="bottom" data-bs-toggle="tooltip">
<i class="theme-switch fa-solid fa-sun fa-lg" data-mode="light"></i>
<i class="theme-switch fa-solid fa-moon fa-lg" data-mode="dark"></i>
<i class="theme-switch fa-solid fa-circle-half-stroke fa-lg" data-mode="auto"></i>
</button>
`);
</script>
<script>
document.write(`
<button class="btn btn-sm pst-navbar-icon 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>
<button class="sidebar-toggle secondary-toggle btn btn-sm" title="Toggle secondary sidebar" data-bs-placement="bottom" data-bs-toggle="tooltip">
<span class="fa-solid fa-list"></span>
</button>
</div></div>
</div>
</div>
</div>
<div id="jb-print-docs-body" class="onlyprint">
<h1>Neural Networks</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="#quick-start-with-neural-networks">Quick Start with Neural Networks</a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#the-module-class">The Module Class</a><ul class="visible nav section-nav flex-column">
<li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#parameters">Parameters</a></li>
<li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#updating-the-parameters">Updating the Parameters</a></li>
<li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#inspecting-modules">Inspecting Modules</a></li>
</ul>
</li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#value-and-grad">Value and Grad</a><ul class="visible nav section-nav flex-column">
</ul>
</li>
</ul>
</nav>
</div>
</div>
</div>
<div id="searchbox"></div>
<article class="bd-article">
<section id="neural-networks">
<span id="nn"></span><h1>Neural Networks<a class="headerlink" href="#neural-networks" title="Link to this heading">#</a></h1>
<p>Writing arbitrarily complex neural networks in MLX can be done using only
<a class="reference internal" href="_autosummary/mlx.core.array.html#mlx.core.array" title="mlx.core.array"><code class="xref py py-class docutils literal notranslate"><span class="pre">mlx.core.array</span></code></a> and <a class="reference internal" href="_autosummary/mlx.core.value_and_grad.html#mlx.core.value_and_grad" title="mlx.core.value_and_grad"><code class="xref py py-meth docutils literal notranslate"><span class="pre">mlx.core.value_and_grad()</span></code></a>. However, this requires the
user to write again and again the same simple neural network operations as well
as handle all the parameter state and initialization manually and explicitly.</p>
<p>The module <code class="xref py py-mod docutils literal notranslate"><span class="pre">mlx.nn</span></code> solves this problem by providing an intuitive way of
composing neural network layers, initializing their parameters, freezing them
for finetuning and more.</p>
<section id="quick-start-with-neural-networks">
<h2>Quick Start with Neural Networks<a class="headerlink" href="#quick-start-with-neural-networks" title="Link to this heading">#</a></h2>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="kn">import</span><span class="w"> </span><span class="nn">mlx.core</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">mx</span>
<span class="kn">import</span><span class="w"> </span><span class="nn">mlx.nn</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">nn</span>
<span class="k">class</span><span class="w"> </span><span class="nc">MLP</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="w"> </span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">in_dims</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span> <span class="n">out_dims</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">layers</span> <span class="o">=</span> <span class="p">[</span>
<span class="n">nn</span><span class="o">.</span><span class="n">Linear</span><span class="p">(</span><span class="n">in_dims</span><span class="p">,</span> <span class="mi">128</span><span class="p">),</span>
<span class="n">nn</span><span class="o">.</span><span class="n">Linear</span><span class="p">(</span><span class="mi">128</span><span class="p">,</span> <span class="mi">128</span><span class="p">),</span>
<span class="n">nn</span><span class="o">.</span><span class="n">Linear</span><span class="p">(</span><span class="mi">128</span><span class="p">,</span> <span class="n">out_dims</span><span class="p">),</span>
<span class="p">]</span>
<span class="k">def</span><span class="w"> </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="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">l</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</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="o">=</span> <span class="n">mx</span><span class="o">.</span><span class="n">maximum</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span> <span class="k">if</span> <span class="n">i</span> <span class="o">&gt;</span> <span class="mi">0</span> <span class="k">else</span> <span class="n">x</span>
<span class="n">x</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="k">return</span> <span class="n">x</span>
<span class="c1"># The model is created with all its parameters but nothing is initialized</span>
<span class="c1"># yet because MLX is lazily evaluated</span>
<span class="n">mlp</span> <span class="o">=</span> <span class="n">MLP</span><span class="p">(</span><span class="mi">2</span><span class="p">,</span> <span class="mi">10</span><span class="p">)</span>
<span class="c1"># We can access its parameters by calling mlp.parameters()</span>
<span class="n">params</span> <span class="o">=</span> <span class="n">mlp</span><span class="o">.</span><span class="n">parameters</span><span class="p">()</span>
<span class="nb">print</span><span class="p">(</span><span class="n">params</span><span class="p">[</span><span class="s2">&quot;layers&quot;</span><span class="p">][</span><span class="mi">0</span><span class="p">][</span><span class="s2">&quot;weight&quot;</span><span class="p">]</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
<span class="c1"># Printing a parameter will cause it to be evaluated and thus initialized</span>
<span class="nb">print</span><span class="p">(</span><span class="n">params</span><span class="p">[</span><span class="s2">&quot;layers&quot;</span><span class="p">][</span><span class="mi">0</span><span class="p">])</span>
<span class="c1"># We can also force evaluate all parameters to initialize the model</span>
<span class="n">mx</span><span class="o">.</span><span class="n">eval</span><span class="p">(</span><span class="n">mlp</span><span class="o">.</span><span class="n">parameters</span><span class="p">())</span>
<span class="c1"># A simple loss function.</span>
<span class="c1"># NOTE: It doesn&#39;t matter how it uses the mlp model. It currently captures</span>
<span class="c1"># it from the local scope. It could be a positional argument or a</span>
<span class="c1"># keyword argument.</span>
<span class="k">def</span><span class="w"> </span><span class="nf">l2_loss</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">y_hat</span> <span class="o">=</span> <span class="n">mlp</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
<span class="k">return</span> <span class="p">(</span><span class="n">y_hat</span> <span class="o">-</span> <span class="n">y</span><span class="p">)</span><span class="o">.</span><span class="n">square</span><span class="p">()</span><span class="o">.</span><span class="n">mean</span><span class="p">()</span>
<span class="c1"># Calling `nn.value_and_grad` instead of `mx.value_and_grad` returns the</span>
<span class="c1"># gradient with respect to `mlp.trainable_parameters()`</span>
<span class="n">loss_and_grad</span> <span class="o">=</span> <span class="n">nn</span><span class="o">.</span><span class="n">value_and_grad</span><span class="p">(</span><span class="n">mlp</span><span class="p">,</span> <span class="n">l2_loss</span><span class="p">)</span>
</pre></div>
</div>
</section>
<section id="the-module-class">
<span id="module-class"></span><h2>The Module Class<a class="headerlink" href="#the-module-class" title="Link to this heading">#</a></h2>
<p>The workhorse of any neural network library is the <a class="reference internal" href="nn/module.html#mlx.nn.Module" title="mlx.nn.Module"><code class="xref py py-class docutils literal notranslate"><span class="pre">Module</span></code></a> class. In
MLX the <a class="reference internal" href="nn/module.html#mlx.nn.Module" title="mlx.nn.Module"><code class="xref py py-class docutils literal notranslate"><span class="pre">Module</span></code></a> class is a container of <a class="reference internal" href="_autosummary/mlx.core.array.html#mlx.core.array" title="mlx.core.array"><code class="xref py py-class docutils literal notranslate"><span class="pre">mlx.core.array</span></code></a> or
<a class="reference internal" href="nn/module.html#mlx.nn.Module" title="mlx.nn.Module"><code class="xref py py-class docutils literal notranslate"><span class="pre">Module</span></code></a> instances. Its main function is to provide a way to
recursively <strong>access</strong> and <strong>update</strong> its parameters and those of its
submodules.</p>
<section id="parameters">
<h3>Parameters<a class="headerlink" href="#parameters" title="Link to this heading">#</a></h3>
<p>A parameter of a module is any public member of type <a class="reference internal" href="_autosummary/mlx.core.array.html#mlx.core.array" title="mlx.core.array"><code class="xref py py-class docutils literal notranslate"><span class="pre">mlx.core.array</span></code></a> (its
name should not start with <code class="docutils literal notranslate"><span class="pre">_</span></code>). It can be arbitrarily nested in other
<a class="reference internal" href="nn/module.html#mlx.nn.Module" title="mlx.nn.Module"><code class="xref py py-class docutils literal notranslate"><span class="pre">Module</span></code></a> instances or lists and dictionaries.</p>
<p><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.parameters.html#mlx.nn.Module.parameters" title="mlx.nn.Module.parameters"><code class="xref py py-meth docutils literal notranslate"><span class="pre">Module.parameters()</span></code></a> can be used to extract a nested dictionary with all
the parameters of a module and its submodules.</p>
<p>A <a class="reference internal" href="nn/module.html#mlx.nn.Module" title="mlx.nn.Module"><code class="xref py py-class docutils literal notranslate"><span class="pre">Module</span></code></a> can also keep track of “frozen” parameters. See the
<a class="reference internal" href="nn/_autosummary/mlx.nn.Module.freeze.html#mlx.nn.Module.freeze" title="mlx.nn.Module.freeze"><code class="xref py py-meth docutils literal notranslate"><span class="pre">Module.freeze()</span></code></a> method for more details. <a class="reference internal" href="_autosummary/mlx.nn.value_and_grad.html#mlx.nn.value_and_grad" title="mlx.nn.value_and_grad"><code class="xref py py-meth docutils literal notranslate"><span class="pre">mlx.nn.value_and_grad()</span></code></a>
the gradients returned will be with respect to these trainable parameters.</p>
</section>
<section id="updating-the-parameters">
<h3>Updating the Parameters<a class="headerlink" href="#updating-the-parameters" title="Link to this heading">#</a></h3>
<p>MLX modules allow accessing and updating individual parameters. However, most
times we need to update large subsets of a module’s parameters. This action is
performed by <a class="reference internal" href="nn/_autosummary/mlx.nn.Module.update.html#mlx.nn.Module.update" title="mlx.nn.Module.update"><code class="xref py py-meth docutils literal notranslate"><span class="pre">Module.update()</span></code></a>.</p>
</section>
<section id="inspecting-modules">
<h3>Inspecting Modules<a class="headerlink" href="#inspecting-modules" title="Link to this heading">#</a></h3>
<p>The simplest way to see the model architecture is to print it. Following along with
the above example, you can print the <code class="docutils literal notranslate"><span class="pre">MLP</span></code> with:</p>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="nb">print</span><span class="p">(</span><span class="n">mlp</span><span class="p">)</span>
</pre></div>
</div>
<p>This will display:</p>
<div class="highlight-shell notranslate"><div class="highlight"><pre><span></span>MLP<span class="o">(</span>
<span class="w"> </span><span class="o">(</span>layers.0<span class="o">)</span>:<span class="w"> </span>Linear<span class="o">(</span><span class="nv">input_dims</span><span class="o">=</span><span class="m">2</span>,<span class="w"> </span><span class="nv">output_dims</span><span class="o">=</span><span class="m">128</span>,<span class="w"> </span><span class="nv">bias</span><span class="o">=</span>True<span class="o">)</span>
<span class="w"> </span><span class="o">(</span>layers.1<span class="o">)</span>:<span class="w"> </span>Linear<span class="o">(</span><span class="nv">input_dims</span><span class="o">=</span><span class="m">128</span>,<span class="w"> </span><span class="nv">output_dims</span><span class="o">=</span><span class="m">128</span>,<span class="w"> </span><span class="nv">bias</span><span class="o">=</span>True<span class="o">)</span>
<span class="w"> </span><span class="o">(</span>layers.2<span class="o">)</span>:<span class="w"> </span>Linear<span class="o">(</span><span class="nv">input_dims</span><span class="o">=</span><span class="m">128</span>,<span class="w"> </span><span class="nv">output_dims</span><span class="o">=</span><span class="m">10</span>,<span class="w"> </span><span class="nv">bias</span><span class="o">=</span>True<span class="o">)</span>
<span class="o">)</span>
</pre></div>
</div>
<p>To get more detailed information on the arrays in a <a class="reference internal" href="nn/module.html#mlx.nn.Module" title="mlx.nn.Module"><code class="xref py py-class docutils literal notranslate"><span class="pre">Module</span></code></a> you can use
<a class="reference internal" href="_autosummary/mlx.utils.tree_map.html#mlx.utils.tree_map" title="mlx.utils.tree_map"><code class="xref py py-func docutils literal notranslate"><span class="pre">mlx.utils.tree_map()</span></code></a> on the parameters. For example, to see the shapes of
all the parameters in a <a class="reference internal" href="nn/module.html#mlx.nn.Module" title="mlx.nn.Module"><code class="xref py py-class docutils literal notranslate"><span class="pre">Module</span></code></a> do:</p>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="kn">from</span><span class="w"> </span><span class="nn">mlx.utils</span><span class="w"> </span><span class="kn">import</span> <span class="n">tree_map</span>
<span class="n">shapes</span> <span class="o">=</span> <span class="n">tree_map</span><span class="p">(</span><span class="k">lambda</span> <span class="n">p</span><span class="p">:</span> <span class="n">p</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="n">mlp</span><span class="o">.</span><span class="n">parameters</span><span class="p">())</span>
</pre></div>
</div>
<p>As another example, you can count the number of parameters in a <a class="reference internal" href="nn/module.html#mlx.nn.Module" title="mlx.nn.Module"><code class="xref py py-class docutils literal notranslate"><span class="pre">Module</span></code></a>
with:</p>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="kn">from</span><span class="w"> </span><span class="nn">mlx.utils</span><span class="w"> </span><span class="kn">import</span> <span class="n">tree_flatten</span>
<span class="n">num_params</span> <span class="o">=</span> <span class="nb">sum</span><span class="p">(</span><span class="n">v</span><span class="o">.</span><span class="n">size</span> <span class="k">for</span> <span class="n">_</span><span class="p">,</span> <span class="n">v</span> <span class="ow">in</span> <span class="n">tree_flatten</span><span class="p">(</span><span class="n">mlp</span><span class="o">.</span><span class="n">parameters</span><span class="p">()))</span>
</pre></div>
</div>
</section>
</section>
<section id="value-and-grad">
<h2>Value and Grad<a class="headerlink" href="#value-and-grad" title="Link to this heading">#</a></h2>
<p>Using a <a class="reference internal" href="nn/module.html#mlx.nn.Module" title="mlx.nn.Module"><code class="xref py py-class docutils literal notranslate"><span class="pre">Module</span></code></a> does not preclude using MLX’s high order function
transformations (<a class="reference internal" href="_autosummary/mlx.core.value_and_grad.html#mlx.core.value_and_grad" title="mlx.core.value_and_grad"><code class="xref py py-meth docutils literal notranslate"><span class="pre">mlx.core.value_and_grad()</span></code></a>, <a class="reference internal" href="_autosummary/mlx.core.grad.html#mlx.core.grad" title="mlx.core.grad"><code class="xref py py-meth docutils literal notranslate"><span class="pre">mlx.core.grad()</span></code></a>, etc.). However,
these function transformations assume pure functions, namely the parameters
should be passed as an argument to the function being transformed.</p>
<p>There is an easy pattern to achieve that with MLX modules</p>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="n">model</span> <span class="o">=</span> <span class="o">...</span>
<span class="k">def</span><span class="w"> </span><span class="nf">f</span><span class="p">(</span><span class="n">params</span><span class="p">,</span> <span class="n">other_inputs</span><span class="p">):</span>
<span class="n">model</span><span class="o">.</span><span class="n">update</span><span class="p">(</span><span class="n">params</span><span class="p">)</span> <span class="c1"># &lt;---- Necessary to make the model use the passed parameters</span>
<span class="k">return</span> <span class="n">model</span><span class="p">(</span><span class="n">other_inputs</span><span class="p">)</span>
<span class="n">f</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">trainable_parameters</span><span class="p">(),</span> <span class="n">mx</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="mi">10</span><span class="p">,)))</span>
</pre></div>
</div>
<p>However, <a class="reference internal" href="_autosummary/mlx.nn.value_and_grad.html#mlx.nn.value_and_grad" title="mlx.nn.value_and_grad"><code class="xref py py-meth docutils literal notranslate"><span class="pre">mlx.nn.value_and_grad()</span></code></a> provides precisely this pattern and only
computes the gradients with respect to the trainable parameters of the model.</p>
<p>In detail:</p>
<ul class="simple">
<li><p>it wraps the passed function with a function that calls <a class="reference internal" href="nn/_autosummary/mlx.nn.Module.update.html#mlx.nn.Module.update" title="mlx.nn.Module.update"><code class="xref py py-meth docutils literal notranslate"><span class="pre">Module.update()</span></code></a>
to make sure the model is using the provided parameters.</p></li>
<li><p>it calls <a class="reference internal" href="_autosummary/mlx.core.value_and_grad.html#mlx.core.value_and_grad" title="mlx.core.value_and_grad"><code class="xref py py-meth docutils literal notranslate"><span class="pre">mlx.core.value_and_grad()</span></code></a> to transform the function into a function
that also computes the gradients with respect to the passed parameters.</p></li>
<li><p>it wraps the returned function with a function that passes the trainable
parameters as the first argument to the function returned by
<a class="reference internal" href="_autosummary/mlx.core.value_and_grad.html#mlx.core.value_and_grad" title="mlx.core.value_and_grad"><code class="xref py py-meth docutils literal notranslate"><span class="pre">mlx.core.value_and_grad()</span></code></a></p></li>
</ul>
<div class="pst-scrollable-table-container"><table class="autosummary longtable table autosummary">
<tbody>
<tr class="row-odd"><td><p><a class="reference internal" href="_autosummary/mlx.nn.value_and_grad.html#mlx.nn.value_and_grad" title="mlx.nn.value_and_grad"><code class="xref py py-obj docutils literal notranslate"><span class="pre">value_and_grad</span></code></a>(model, fn)</p></td>
<td><p>Transform the passed function <code class="docutils literal notranslate"><span class="pre">fn</span></code> to a function that computes the gradients of <code class="docutils literal notranslate"><span class="pre">fn</span></code> wrt the model's trainable parameters and also its value.</p></td>
</tr>
<tr class="row-even"><td><p><a class="reference internal" href="_autosummary/mlx.nn.quantize.html#mlx.nn.quantize" title="mlx.nn.quantize"><code class="xref py py-obj docutils literal notranslate"><span class="pre">quantize</span></code></a>(model[, group_size, bits, ...])</p></td>
<td><p>Quantize the sub-modules of a module according to a predicate.</p></td>
</tr>
<tr class="row-odd"><td><p><a class="reference internal" href="_autosummary/mlx.nn.average_gradients.html#mlx.nn.average_gradients" title="mlx.nn.average_gradients"><code class="xref py py-obj docutils literal notranslate"><span class="pre">average_gradients</span></code></a>(gradients[, group, ...])</p></td>
<td><p>Average the gradients across the distributed processes in the passed group.</p></td>
</tr>
</tbody>
</table>
</div>
<div class="toctree-wrapper compound">
<ul>
<li class="toctree-l1"><a class="reference internal" href="nn/module.html">Module</a><ul>
<li class="toctree-l2"><a class="reference internal" href="nn/module.html#mlx.nn.Module"><code class="docutils literal notranslate"><span class="pre">Module</span></code></a></li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.training.html">mlx.nn.Module.training</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.training.html#mlx.nn.Module.training"><code class="docutils literal notranslate"><span class="pre">Module.training</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.state.html">mlx.nn.Module.state</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.state.html#mlx.nn.Module.state"><code class="docutils literal notranslate"><span class="pre">Module.state</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.apply.html">mlx.nn.Module.apply</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.apply.html#mlx.nn.Module.apply"><code class="docutils literal notranslate"><span class="pre">Module.apply()</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.apply_to_modules.html">mlx.nn.Module.apply_to_modules</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.apply_to_modules.html#mlx.nn.Module.apply_to_modules"><code class="docutils literal notranslate"><span class="pre">Module.apply_to_modules()</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.children.html">mlx.nn.Module.children</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.children.html#mlx.nn.Module.children"><code class="docutils literal notranslate"><span class="pre">Module.children()</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.eval.html">mlx.nn.Module.eval</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.eval.html#mlx.nn.Module.eval"><code class="docutils literal notranslate"><span class="pre">Module.eval()</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.filter_and_map.html">mlx.nn.Module.filter_and_map</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.filter_and_map.html#mlx.nn.Module.filter_and_map"><code class="docutils literal notranslate"><span class="pre">Module.filter_and_map()</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.freeze.html">mlx.nn.Module.freeze</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.freeze.html#mlx.nn.Module.freeze"><code class="docutils literal notranslate"><span class="pre">Module.freeze()</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.leaf_modules.html">mlx.nn.Module.leaf_modules</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.leaf_modules.html#mlx.nn.Module.leaf_modules"><code class="docutils literal notranslate"><span class="pre">Module.leaf_modules()</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.load_weights.html">mlx.nn.Module.load_weights</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.load_weights.html#mlx.nn.Module.load_weights"><code class="docutils literal notranslate"><span class="pre">Module.load_weights()</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.modules.html">mlx.nn.Module.modules</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.modules.html#mlx.nn.Module.modules"><code class="docutils literal notranslate"><span class="pre">Module.modules()</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.named_modules.html">mlx.nn.Module.named_modules</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.named_modules.html#mlx.nn.Module.named_modules"><code class="docutils literal notranslate"><span class="pre">Module.named_modules()</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.parameters.html">mlx.nn.Module.parameters</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.parameters.html#mlx.nn.Module.parameters"><code class="docutils literal notranslate"><span class="pre">Module.parameters()</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.save_weights.html">mlx.nn.Module.save_weights</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.save_weights.html#mlx.nn.Module.save_weights"><code class="docutils literal notranslate"><span class="pre">Module.save_weights()</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.set_dtype.html">mlx.nn.Module.set_dtype</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.set_dtype.html#mlx.nn.Module.set_dtype"><code class="docutils literal notranslate"><span class="pre">Module.set_dtype()</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.train.html">mlx.nn.Module.train</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.train.html#mlx.nn.Module.train"><code class="docutils literal notranslate"><span class="pre">Module.train()</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.trainable_parameters.html">mlx.nn.Module.trainable_parameters</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.trainable_parameters.html#mlx.nn.Module.trainable_parameters"><code class="docutils literal notranslate"><span class="pre">Module.trainable_parameters()</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.unfreeze.html">mlx.nn.Module.unfreeze</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.unfreeze.html#mlx.nn.Module.unfreeze"><code class="docutils literal notranslate"><span class="pre">Module.unfreeze()</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.update.html">mlx.nn.Module.update</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.update.html#mlx.nn.Module.update"><code class="docutils literal notranslate"><span class="pre">Module.update()</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.update_modules.html">mlx.nn.Module.update_modules</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Module.update_modules.html#mlx.nn.Module.update_modules"><code class="docutils literal notranslate"><span class="pre">Module.update_modules()</span></code></a></li>
</ul>
</li>
</ul>
</li>
<li class="toctree-l1"><a class="reference internal" href="nn/layers.html">Layers</a><ul>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.ALiBi.html">mlx.nn.ALiBi</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.ALiBi.html#mlx.nn.ALiBi"><code class="docutils literal notranslate"><span class="pre">ALiBi</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.AvgPool1d.html">mlx.nn.AvgPool1d</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.AvgPool1d.html#mlx.nn.AvgPool1d"><code class="docutils literal notranslate"><span class="pre">AvgPool1d</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.AvgPool2d.html">mlx.nn.AvgPool2d</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.AvgPool2d.html#mlx.nn.AvgPool2d"><code class="docutils literal notranslate"><span class="pre">AvgPool2d</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.AvgPool3d.html">mlx.nn.AvgPool3d</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.AvgPool3d.html#mlx.nn.AvgPool3d"><code class="docutils literal notranslate"><span class="pre">AvgPool3d</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.BatchNorm.html">mlx.nn.BatchNorm</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.BatchNorm.html#mlx.nn.BatchNorm"><code class="docutils literal notranslate"><span class="pre">BatchNorm</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.CELU.html">mlx.nn.CELU</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.CELU.html#mlx.nn.CELU"><code class="docutils literal notranslate"><span class="pre">CELU</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Conv1d.html">mlx.nn.Conv1d</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Conv1d.html#mlx.nn.Conv1d"><code class="docutils literal notranslate"><span class="pre">Conv1d</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Conv2d.html">mlx.nn.Conv2d</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Conv2d.html#mlx.nn.Conv2d"><code class="docutils literal notranslate"><span class="pre">Conv2d</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Conv3d.html">mlx.nn.Conv3d</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Conv3d.html#mlx.nn.Conv3d"><code class="docutils literal notranslate"><span class="pre">Conv3d</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.ConvTranspose1d.html">mlx.nn.ConvTranspose1d</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.ConvTranspose1d.html#mlx.nn.ConvTranspose1d"><code class="docutils literal notranslate"><span class="pre">ConvTranspose1d</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.ConvTranspose2d.html">mlx.nn.ConvTranspose2d</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.ConvTranspose2d.html#mlx.nn.ConvTranspose2d"><code class="docutils literal notranslate"><span class="pre">ConvTranspose2d</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.ConvTranspose3d.html">mlx.nn.ConvTranspose3d</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.ConvTranspose3d.html#mlx.nn.ConvTranspose3d"><code class="docutils literal notranslate"><span class="pre">ConvTranspose3d</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Dropout.html">mlx.nn.Dropout</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Dropout.html#mlx.nn.Dropout"><code class="docutils literal notranslate"><span class="pre">Dropout</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Dropout2d.html">mlx.nn.Dropout2d</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Dropout2d.html#mlx.nn.Dropout2d"><code class="docutils literal notranslate"><span class="pre">Dropout2d</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Dropout3d.html">mlx.nn.Dropout3d</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Dropout3d.html#mlx.nn.Dropout3d"><code class="docutils literal notranslate"><span class="pre">Dropout3d</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Embedding.html">mlx.nn.Embedding</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Embedding.html#mlx.nn.Embedding"><code class="docutils literal notranslate"><span class="pre">Embedding</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.ELU.html">mlx.nn.ELU</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.ELU.html#mlx.nn.ELU"><code class="docutils literal notranslate"><span class="pre">ELU</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.GELU.html">mlx.nn.GELU</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.GELU.html#mlx.nn.GELU"><code class="docutils literal notranslate"><span class="pre">GELU</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.GLU.html">mlx.nn.GLU</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.GLU.html#mlx.nn.GLU"><code class="docutils literal notranslate"><span class="pre">GLU</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.GroupNorm.html">mlx.nn.GroupNorm</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.GroupNorm.html#mlx.nn.GroupNorm"><code class="docutils literal notranslate"><span class="pre">GroupNorm</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.GRU.html">mlx.nn.GRU</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.GRU.html#mlx.nn.GRU"><code class="docutils literal notranslate"><span class="pre">GRU</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.HardShrink.html">mlx.nn.HardShrink</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.HardShrink.html#mlx.nn.HardShrink"><code class="docutils literal notranslate"><span class="pre">HardShrink</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.HardTanh.html">mlx.nn.HardTanh</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.HardTanh.html#mlx.nn.HardTanh"><code class="docutils literal notranslate"><span class="pre">HardTanh</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Hardswish.html">mlx.nn.Hardswish</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Hardswish.html#mlx.nn.Hardswish"><code class="docutils literal notranslate"><span class="pre">Hardswish</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.InstanceNorm.html">mlx.nn.InstanceNorm</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.InstanceNorm.html#mlx.nn.InstanceNorm"><code class="docutils literal notranslate"><span class="pre">InstanceNorm</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.LayerNorm.html">mlx.nn.LayerNorm</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.LayerNorm.html#mlx.nn.LayerNorm"><code class="docutils literal notranslate"><span class="pre">LayerNorm</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.LeakyReLU.html">mlx.nn.LeakyReLU</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.LeakyReLU.html#mlx.nn.LeakyReLU"><code class="docutils literal notranslate"><span class="pre">LeakyReLU</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Linear.html">mlx.nn.Linear</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Linear.html#mlx.nn.Linear"><code class="docutils literal notranslate"><span class="pre">Linear</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.LogSigmoid.html">mlx.nn.LogSigmoid</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.LogSigmoid.html#mlx.nn.LogSigmoid"><code class="docutils literal notranslate"><span class="pre">LogSigmoid</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.LogSoftmax.html">mlx.nn.LogSoftmax</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.LogSoftmax.html#mlx.nn.LogSoftmax"><code class="docutils literal notranslate"><span class="pre">LogSoftmax</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.LSTM.html">mlx.nn.LSTM</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.LSTM.html#mlx.nn.LSTM"><code class="docutils literal notranslate"><span class="pre">LSTM</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.MaxPool1d.html">mlx.nn.MaxPool1d</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.MaxPool1d.html#mlx.nn.MaxPool1d"><code class="docutils literal notranslate"><span class="pre">MaxPool1d</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.MaxPool2d.html">mlx.nn.MaxPool2d</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.MaxPool2d.html#mlx.nn.MaxPool2d"><code class="docutils literal notranslate"><span class="pre">MaxPool2d</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.MaxPool3d.html">mlx.nn.MaxPool3d</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.MaxPool3d.html#mlx.nn.MaxPool3d"><code class="docutils literal notranslate"><span class="pre">MaxPool3d</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Mish.html">mlx.nn.Mish</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Mish.html#mlx.nn.Mish"><code class="docutils literal notranslate"><span class="pre">Mish</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.MultiHeadAttention.html">mlx.nn.MultiHeadAttention</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.MultiHeadAttention.html#mlx.nn.MultiHeadAttention"><code class="docutils literal notranslate"><span class="pre">MultiHeadAttention</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.PReLU.html">mlx.nn.PReLU</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.PReLU.html#mlx.nn.PReLU"><code class="docutils literal notranslate"><span class="pre">PReLU</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.QuantizedEmbedding.html">mlx.nn.QuantizedEmbedding</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.QuantizedEmbedding.html#mlx.nn.QuantizedEmbedding"><code class="docutils literal notranslate"><span class="pre">QuantizedEmbedding</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.QuantizedLinear.html">mlx.nn.QuantizedLinear</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.QuantizedLinear.html#mlx.nn.QuantizedLinear"><code class="docutils literal notranslate"><span class="pre">QuantizedLinear</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.RMSNorm.html">mlx.nn.RMSNorm</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.RMSNorm.html#mlx.nn.RMSNorm"><code class="docutils literal notranslate"><span class="pre">RMSNorm</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.ReLU.html">mlx.nn.ReLU</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.ReLU.html#mlx.nn.ReLU"><code class="docutils literal notranslate"><span class="pre">ReLU</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.ReLU6.html">mlx.nn.ReLU6</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.ReLU6.html#mlx.nn.ReLU6"><code class="docutils literal notranslate"><span class="pre">ReLU6</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.RNN.html">mlx.nn.RNN</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.RNN.html#mlx.nn.RNN"><code class="docutils literal notranslate"><span class="pre">RNN</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.RoPE.html">mlx.nn.RoPE</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.RoPE.html#mlx.nn.RoPE"><code class="docutils literal notranslate"><span class="pre">RoPE</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.SELU.html">mlx.nn.SELU</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.SELU.html#mlx.nn.SELU"><code class="docutils literal notranslate"><span class="pre">SELU</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Sequential.html">mlx.nn.Sequential</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Sequential.html#mlx.nn.Sequential"><code class="docutils literal notranslate"><span class="pre">Sequential</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Sigmoid.html">mlx.nn.Sigmoid</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Sigmoid.html#mlx.nn.Sigmoid"><code class="docutils literal notranslate"><span class="pre">Sigmoid</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.SiLU.html">mlx.nn.SiLU</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.SiLU.html#mlx.nn.SiLU"><code class="docutils literal notranslate"><span class="pre">SiLU</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.SinusoidalPositionalEncoding.html">mlx.nn.SinusoidalPositionalEncoding</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.SinusoidalPositionalEncoding.html#mlx.nn.SinusoidalPositionalEncoding"><code class="docutils literal notranslate"><span class="pre">SinusoidalPositionalEncoding</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Softmin.html">mlx.nn.Softmin</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Softmin.html#mlx.nn.Softmin"><code class="docutils literal notranslate"><span class="pre">Softmin</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Softshrink.html">mlx.nn.Softshrink</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Softshrink.html#mlx.nn.Softshrink"><code class="docutils literal notranslate"><span class="pre">Softshrink</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Softsign.html">mlx.nn.Softsign</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Softsign.html#mlx.nn.Softsign"><code class="docutils literal notranslate"><span class="pre">Softsign</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Softmax.html">mlx.nn.Softmax</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Softmax.html#mlx.nn.Softmax"><code class="docutils literal notranslate"><span class="pre">Softmax</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Softplus.html">mlx.nn.Softplus</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Softplus.html#mlx.nn.Softplus"><code class="docutils literal notranslate"><span class="pre">Softplus</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Step.html">mlx.nn.Step</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Step.html#mlx.nn.Step"><code class="docutils literal notranslate"><span class="pre">Step</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Tanh.html">mlx.nn.Tanh</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Tanh.html#mlx.nn.Tanh"><code class="docutils literal notranslate"><span class="pre">Tanh</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Transformer.html">mlx.nn.Transformer</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Transformer.html#mlx.nn.Transformer"><code class="docutils literal notranslate"><span class="pre">Transformer</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.Upsample.html">mlx.nn.Upsample</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.Upsample.html#mlx.nn.Upsample"><code class="docutils literal notranslate"><span class="pre">Upsample</span></code></a></li>
</ul>
</li>
</ul>
</li>
<li class="toctree-l1"><a class="reference internal" href="nn/functions.html">Functions</a><ul>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.elu.html">mlx.nn.elu</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.elu.html#mlx.nn.elu"><code class="docutils literal notranslate"><span class="pre">elu</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.celu.html">mlx.nn.celu</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.celu.html#mlx.nn.celu"><code class="docutils literal notranslate"><span class="pre">celu</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.gelu.html">mlx.nn.gelu</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.gelu.html#mlx.nn.gelu"><code class="docutils literal notranslate"><span class="pre">gelu</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.gelu_approx.html">mlx.nn.gelu_approx</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.gelu_approx.html#mlx.nn.gelu_approx"><code class="docutils literal notranslate"><span class="pre">gelu_approx</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.gelu_fast_approx.html">mlx.nn.gelu_fast_approx</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.gelu_fast_approx.html#mlx.nn.gelu_fast_approx"><code class="docutils literal notranslate"><span class="pre">gelu_fast_approx</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.glu.html">mlx.nn.glu</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.glu.html#mlx.nn.glu"><code class="docutils literal notranslate"><span class="pre">glu</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.hard_shrink.html">mlx.nn.hard_shrink</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.hard_shrink.html#mlx.nn.hard_shrink"><code class="docutils literal notranslate"><span class="pre">hard_shrink</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.hard_tanh.html">mlx.nn.hard_tanh</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.hard_tanh.html#mlx.nn.hard_tanh"><code class="docutils literal notranslate"><span class="pre">hard_tanh</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.hardswish.html">mlx.nn.hardswish</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.hardswish.html#mlx.nn.hardswish"><code class="docutils literal notranslate"><span class="pre">hardswish</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.leaky_relu.html">mlx.nn.leaky_relu</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.leaky_relu.html#mlx.nn.leaky_relu"><code class="docutils literal notranslate"><span class="pre">leaky_relu</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.log_sigmoid.html">mlx.nn.log_sigmoid</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.log_sigmoid.html#mlx.nn.log_sigmoid"><code class="docutils literal notranslate"><span class="pre">log_sigmoid</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.log_softmax.html">mlx.nn.log_softmax</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.log_softmax.html#mlx.nn.log_softmax"><code class="docutils literal notranslate"><span class="pre">log_softmax</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.mish.html">mlx.nn.mish</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.mish.html#mlx.nn.mish"><code class="docutils literal notranslate"><span class="pre">mish</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.prelu.html">mlx.nn.prelu</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.prelu.html#mlx.nn.prelu"><code class="docutils literal notranslate"><span class="pre">prelu</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.relu.html">mlx.nn.relu</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.relu.html#mlx.nn.relu"><code class="docutils literal notranslate"><span class="pre">relu</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.relu6.html">mlx.nn.relu6</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.relu6.html#mlx.nn.relu6"><code class="docutils literal notranslate"><span class="pre">relu6</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.selu.html">mlx.nn.selu</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.selu.html#mlx.nn.selu"><code class="docutils literal notranslate"><span class="pre">selu</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.sigmoid.html">mlx.nn.sigmoid</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.sigmoid.html#mlx.nn.sigmoid"><code class="docutils literal notranslate"><span class="pre">sigmoid</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.silu.html">mlx.nn.silu</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.silu.html#mlx.nn.silu"><code class="docutils literal notranslate"><span class="pre">silu</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.softmax.html">mlx.nn.softmax</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.softmax.html#mlx.nn.softmax"><code class="docutils literal notranslate"><span class="pre">softmax</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.softmin.html">mlx.nn.softmin</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.softmin.html#mlx.nn.softmin"><code class="docutils literal notranslate"><span class="pre">softmin</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.softplus.html">mlx.nn.softplus</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.softplus.html#mlx.nn.softplus"><code class="docutils literal notranslate"><span class="pre">softplus</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.softshrink.html">mlx.nn.softshrink</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.softshrink.html#mlx.nn.softshrink"><code class="docutils literal notranslate"><span class="pre">softshrink</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.step.html">mlx.nn.step</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.step.html#mlx.nn.step"><code class="docutils literal notranslate"><span class="pre">step</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.tanh.html">mlx.nn.tanh</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.tanh.html#mlx.nn.tanh"><code class="docutils literal notranslate"><span class="pre">tanh</span></code></a></li>
</ul>
</li>
</ul>
</li>
<li class="toctree-l1"><a class="reference internal" href="nn/losses.html">Loss Functions</a><ul>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.binary_cross_entropy.html">mlx.nn.losses.binary_cross_entropy</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.binary_cross_entropy.html#mlx.nn.losses.binary_cross_entropy"><code class="docutils literal notranslate"><span class="pre">binary_cross_entropy</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.cosine_similarity_loss.html">mlx.nn.losses.cosine_similarity_loss</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.cosine_similarity_loss.html#mlx.nn.losses.cosine_similarity_loss"><code class="docutils literal notranslate"><span class="pre">cosine_similarity_loss</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.cross_entropy.html">mlx.nn.losses.cross_entropy</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.cross_entropy.html#mlx.nn.losses.cross_entropy"><code class="docutils literal notranslate"><span class="pre">cross_entropy</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.gaussian_nll_loss.html">mlx.nn.losses.gaussian_nll_loss</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.gaussian_nll_loss.html#mlx.nn.losses.gaussian_nll_loss"><code class="docutils literal notranslate"><span class="pre">gaussian_nll_loss</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.hinge_loss.html">mlx.nn.losses.hinge_loss</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.hinge_loss.html#mlx.nn.losses.hinge_loss"><code class="docutils literal notranslate"><span class="pre">hinge_loss</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.huber_loss.html">mlx.nn.losses.huber_loss</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.huber_loss.html#mlx.nn.losses.huber_loss"><code class="docutils literal notranslate"><span class="pre">huber_loss</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.kl_div_loss.html">mlx.nn.losses.kl_div_loss</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.kl_div_loss.html#mlx.nn.losses.kl_div_loss"><code class="docutils literal notranslate"><span class="pre">kl_div_loss</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.l1_loss.html">mlx.nn.losses.l1_loss</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.l1_loss.html#mlx.nn.losses.l1_loss"><code class="docutils literal notranslate"><span class="pre">l1_loss</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.log_cosh_loss.html">mlx.nn.losses.log_cosh_loss</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.log_cosh_loss.html#mlx.nn.losses.log_cosh_loss"><code class="docutils literal notranslate"><span class="pre">log_cosh_loss</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.margin_ranking_loss.html">mlx.nn.losses.margin_ranking_loss</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.margin_ranking_loss.html#mlx.nn.losses.margin_ranking_loss"><code class="docutils literal notranslate"><span class="pre">margin_ranking_loss</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.mse_loss.html">mlx.nn.losses.mse_loss</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.mse_loss.html#mlx.nn.losses.mse_loss"><code class="docutils literal notranslate"><span class="pre">mse_loss</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.nll_loss.html">mlx.nn.losses.nll_loss</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.nll_loss.html#mlx.nn.losses.nll_loss"><code class="docutils literal notranslate"><span class="pre">nll_loss</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.smooth_l1_loss.html">mlx.nn.losses.smooth_l1_loss</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.smooth_l1_loss.html#mlx.nn.losses.smooth_l1_loss"><code class="docutils literal notranslate"><span class="pre">smooth_l1_loss</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.triplet_loss.html">mlx.nn.losses.triplet_loss</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary_functions/mlx.nn.losses.triplet_loss.html#mlx.nn.losses.triplet_loss"><code class="docutils literal notranslate"><span class="pre">triplet_loss</span></code></a></li>
</ul>
</li>
</ul>
</li>
<li class="toctree-l1"><a class="reference internal" href="nn/init.html">Initializers</a><ul>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.init.constant.html">mlx.nn.init.constant</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.init.constant.html#mlx.nn.init.constant"><code class="docutils literal notranslate"><span class="pre">constant()</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.init.normal.html">mlx.nn.init.normal</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.init.normal.html#mlx.nn.init.normal"><code class="docutils literal notranslate"><span class="pre">normal()</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.init.uniform.html">mlx.nn.init.uniform</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.init.uniform.html#mlx.nn.init.uniform"><code class="docutils literal notranslate"><span class="pre">uniform()</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.init.identity.html">mlx.nn.init.identity</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.init.identity.html#mlx.nn.init.identity"><code class="docutils literal notranslate"><span class="pre">identity()</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.init.glorot_normal.html">mlx.nn.init.glorot_normal</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.init.glorot_normal.html#mlx.nn.init.glorot_normal"><code class="docutils literal notranslate"><span class="pre">glorot_normal()</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.init.glorot_uniform.html">mlx.nn.init.glorot_uniform</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.init.glorot_uniform.html#mlx.nn.init.glorot_uniform"><code class="docutils literal notranslate"><span class="pre">glorot_uniform()</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.init.he_normal.html">mlx.nn.init.he_normal</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.init.he_normal.html#mlx.nn.init.he_normal"><code class="docutils literal notranslate"><span class="pre">he_normal()</span></code></a></li>
</ul>
</li>
<li class="toctree-l2"><a class="reference internal" href="nn/_autosummary/mlx.nn.init.he_uniform.html">mlx.nn.init.he_uniform</a><ul>
<li class="toctree-l3"><a class="reference internal" href="nn/_autosummary/mlx.nn.init.he_uniform.html#mlx.nn.init.he_uniform"><code class="docutils literal notranslate"><span class="pre">he_uniform()</span></code></a></li>
</ul>
</li>
</ul>
</li>
</ul>
</div>
</section>
</section>
</article>
<footer class="prev-next-footer d-print-none">
<div class="prev-next-area">
<a class="left-prev"
href="_autosummary/mlx.core.metal.stop_capture.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">mlx.core.metal.stop_capture</p>
</div>
</a>
<a class="right-next"
href="_autosummary/mlx.nn.value_and_grad.html"
title="next page">
<div class="prev-next-info">
<p class="prev-next-subtitle">next</p>
<p class="prev-next-title">mlx.nn.value_and_grad</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="#quick-start-with-neural-networks">Quick Start with Neural Networks</a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#the-module-class">The Module Class</a><ul class="visible nav section-nav flex-column">
<li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#parameters">Parameters</a></li>
<li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#updating-the-parameters">Updating the Parameters</a></li>
<li class="toc-h3 nav-item toc-entry"><a class="reference internal nav-link" href="#inspecting-modules">Inspecting Modules</a></li>
</ul>
</li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#value-and-grad">Value and Grad</a><ul class="visible nav section-nav flex-column">
</ul>
</li>
</ul>
</nav></div>
</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>
</div>
<div class="footer-item">
<p class="copyright">
© Copyright 2023, MLX Contributors.
<br/>
</p>
</div>
<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=dfe6caa3a7d634c4db9b"></script>
<script src="../_static/scripts/pydata-sphinx-theme.js?digest=dfe6caa3a7d634c4db9b"></script>
<footer class="bd-footer">
</footer>
</body>
</html>
Reference in New Issue View Git Blame Copy Permalink