docs update

This commit is contained in:
Awni Hannun
2024-06-06 20:28:06 -07:00
committed by CircleCI Docs
parent d7a78fbe2b
commit 85f70be0e6
2046 changed files with 323202 additions and 2319 deletions

View File

@@ -8,7 +8,7 @@
<meta charset="utf-8" />
<meta name="viewport" content="width=device-width, initial-scale=1.0" /><meta name="generator" content="Docutils 0.18.1: http://docutils.sourceforge.net/" />
<title>Operations &#8212; MLX 0.14.0 documentation</title>
<title>Operations &#8212; MLX 0.15.0 documentation</title>
@@ -36,7 +36,7 @@
<link rel="preload" as="script" href="../_static/scripts/pydata-sphinx-theme.js?digest=5b4479735964841361fd" />
<script src="../_static/vendor/fontawesome/6.1.2/js/all.min.js?digest=5b4479735964841361fd"></script>
<script src="../_static/documentation_options.js?v=8eac8a5b"></script>
<script src="../_static/documentation_options.js?v=7276bfb7"></script>
<script src="../_static/doctools.js?v=888ff710"></script>
<script src="../_static/sphinx_highlight.js?v=dc90522c"></script>
<script src="../_static/scripts/sphinx-book-theme.js?v=efea14e4"></script>
@@ -131,8 +131,8 @@
<img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.14.0 documentation - Home"/>
<script>document.write(`<img src="../_static/mlx_logo_dark.png" class="logo__image only-dark" alt="MLX 0.14.0 documentation - Home"/>`);</script>
<img src="../_static/mlx_logo.png" class="logo__image only-light" alt="MLX 0.15.0 documentation - Home"/>
<script>document.write(`<img src="../_static/mlx_logo_dark.png" class="logo__image only-dark" alt="MLX 0.15.0 documentation - Home"/>`);</script>
</a></div>
@@ -163,6 +163,7 @@
<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>
</ul>
<p aria-level="2" class="caption" role="heading"><span class="caption-text">Examples</span></p>
@@ -392,6 +393,7 @@
<li class="toctree-l2"><a class="reference internal" href="../python/_autosummary/mlx.core.tril.html">mlx.core.tril</a></li>
<li class="toctree-l2"><a class="reference internal" href="../python/_autosummary/mlx.core.triu.html">mlx.core.triu</a></li>
<li class="toctree-l2"><a class="reference internal" href="../python/_autosummary/mlx.core.var.html">mlx.core.var</a></li>
<li class="toctree-l2"><a class="reference internal" href="../python/_autosummary/mlx.core.view.html">mlx.core.view</a></li>
<li class="toctree-l2"><a class="reference internal" href="../python/_autosummary/mlx.core.where.html">mlx.core.where</a></li>
<li class="toctree-l2"><a class="reference internal" href="../python/_autosummary/mlx.core.zeros.html">mlx.core.zeros</a></li>
<li class="toctree-l2"><a class="reference internal" href="../python/_autosummary/mlx.core.zeros_like.html">mlx.core.zeros_like</a></li>
@@ -506,10 +508,15 @@
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary/mlx.nn.Dropout3d.html">mlx.nn.Dropout3d</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary/mlx.nn.Embedding.html">mlx.nn.Embedding</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary/mlx.nn.GELU.html">mlx.nn.GELU</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary/mlx.nn.GLU.html">mlx.nn.GLU</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary/mlx.nn.GroupNorm.html">mlx.nn.GroupNorm</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary/mlx.nn.GRU.html">mlx.nn.GRU</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary/mlx.nn.HardShrink.html">mlx.nn.HardShrink</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary/mlx.nn.HardTanh.html">mlx.nn.HardTanh</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary/mlx.nn.Hardswish.html">mlx.nn.Hardswish</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary/mlx.nn.InstanceNorm.html">mlx.nn.InstanceNorm</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary/mlx.nn.LayerNorm.html">mlx.nn.LayerNorm</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary/mlx.nn.LeakyReLU.html">mlx.nn.LeakyReLU</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary/mlx.nn.Linear.html">mlx.nn.Linear</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary/mlx.nn.LSTM.html">mlx.nn.LSTM</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary/mlx.nn.MaxPool1d.html">mlx.nn.MaxPool1d</a></li>
@@ -521,14 +528,20 @@
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary/mlx.nn.QuantizedLinear.html">mlx.nn.QuantizedLinear</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary/mlx.nn.RMSNorm.html">mlx.nn.RMSNorm</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary/mlx.nn.ReLU.html">mlx.nn.ReLU</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary/mlx.nn.ReLU6.html">mlx.nn.ReLU6</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary/mlx.nn.RNN.html">mlx.nn.RNN</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary/mlx.nn.RoPE.html">mlx.nn.RoPE</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary/mlx.nn.SELU.html">mlx.nn.SELU</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary/mlx.nn.Sequential.html">mlx.nn.Sequential</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary/mlx.nn.SiLU.html">mlx.nn.SiLU</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary/mlx.nn.SinusoidalPositionalEncoding.html">mlx.nn.SinusoidalPositionalEncoding</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary/mlx.nn.Softmin.html">mlx.nn.Softmin</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary/mlx.nn.Softshrink.html">mlx.nn.Softshrink</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary/mlx.nn.Softsign.html">mlx.nn.Softsign</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary/mlx.nn.Softmax.html">mlx.nn.Softmax</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary/mlx.nn.Softplus.html">mlx.nn.Softplus</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary/mlx.nn.Step.html">mlx.nn.Step</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary/mlx.nn.Tanh.html">mlx.nn.Tanh</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary/mlx.nn.Transformer.html">mlx.nn.Transformer</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary/mlx.nn.Upsample.html">mlx.nn.Upsample</a></li>
</ul>
@@ -539,6 +552,8 @@
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary_functions/mlx.nn.gelu_approx.html">mlx.nn.gelu_approx</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary_functions/mlx.nn.gelu_fast_approx.html">mlx.nn.gelu_fast_approx</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary_functions/mlx.nn.glu.html">mlx.nn.glu</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary_functions/mlx.nn.hard_shrink.html">mlx.nn.hard_shrink</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary_functions/mlx.nn.hard_tanh.html">mlx.nn.hard_tanh</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary_functions/mlx.nn.hardswish.html">mlx.nn.hardswish</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary_functions/mlx.nn.leaky_relu.html">mlx.nn.leaky_relu</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary_functions/mlx.nn.log_sigmoid.html">mlx.nn.log_sigmoid</a></li>
@@ -551,6 +566,7 @@
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary_functions/mlx.nn.sigmoid.html">mlx.nn.sigmoid</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary_functions/mlx.nn.silu.html">mlx.nn.silu</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary_functions/mlx.nn.softmax.html">mlx.nn.softmax</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary_functions/mlx.nn.softmin.html">mlx.nn.softmin</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary_functions/mlx.nn.softplus.html">mlx.nn.softplus</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary_functions/mlx.nn.softshrink.html">mlx.nn.softshrink</a></li>
<li class="toctree-l3"><a class="reference internal" href="../python/nn/_autosummary_functions/mlx.nn.step.html">mlx.nn.step</a></li>
@@ -618,7 +634,15 @@
<li class="toctree-l2"><a class="reference internal" href="../python/_autosummary/mlx.optimizers.clip_grad_norm.html">mlx.optimizers.clip_grad_norm</a></li>
</ul>
</li>
<li class="toctree-l1 has-children"><a class="reference internal" href="../python/tree_utils.html">Tree Utils</a><input class="toctree-checkbox" id="toctree-checkbox-21" name="toctree-checkbox-21" type="checkbox"/><label class="toctree-toggle" for="toctree-checkbox-21"><i class="fa-solid fa-chevron-down"></i></label><ul>
<li class="toctree-l1 has-children"><a class="reference internal" href="../python/distributed.html">Distributed Communication</a><input class="toctree-checkbox" id="toctree-checkbox-21" name="toctree-checkbox-21" type="checkbox"/><label class="toctree-toggle" for="toctree-checkbox-21"><i class="fa-solid fa-chevron-down"></i></label><ul>
<li class="toctree-l2"><a class="reference internal" href="../python/_autosummary/mlx.core.distributed.Group.html">mlx.core.distributed.Group</a></li>
<li class="toctree-l2"><a class="reference internal" href="../python/_autosummary/mlx.core.distributed.is_available.html">mlx.core.distributed.is_available</a></li>
<li class="toctree-l2"><a class="reference internal" href="../python/_autosummary/mlx.core.distributed.init.html">mlx.core.distributed.init</a></li>
<li class="toctree-l2"><a class="reference internal" href="../python/_autosummary/mlx.core.distributed.all_sum.html">mlx.core.distributed.all_sum</a></li>
<li class="toctree-l2"><a class="reference internal" href="../python/_autosummary/mlx.core.distributed.all_gather.html">mlx.core.distributed.all_gather</a></li>
</ul>
</li>
<li class="toctree-l1 has-children"><a class="reference internal" href="../python/tree_utils.html">Tree Utils</a><input class="toctree-checkbox" id="toctree-checkbox-22" name="toctree-checkbox-22" type="checkbox"/><label class="toctree-toggle" for="toctree-checkbox-22"><i class="fa-solid fa-chevron-down"></i></label><ul>
<li class="toctree-l2"><a class="reference internal" href="../python/_autosummary/mlx.utils.tree_flatten.html">mlx.utils.tree_flatten</a></li>
<li class="toctree-l2"><a class="reference internal" href="../python/_autosummary/mlx.utils.tree_unflatten.html">mlx.utils.tree_unflatten</a></li>
<li class="toctree-l2"><a class="reference internal" href="../python/_autosummary/mlx.utils.tree_map.html">mlx.utils.tree_map</a></li>
@@ -1056,16 +1080,19 @@ document.write(`
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv416quantized_matmulRK5arrayRK5arrayRK5arrayRK5arraybii14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">quantized_matmul()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv48quantizeRK5arrayii14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">quantize()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv410dequantizeRK5arrayRK5arrayRK5arrayii14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">dequantize()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv416block_sparse_qmmRK5arrayRK5arrayRK5arrayRK5arrayNSt8optionalI5arrayEENSt8optionalI5arrayEEbii14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">block_sparse_qmm()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv410gather_qmmRK5arrayRK5arrayRK5arrayRK5arrayNSt8optionalI5arrayEENSt8optionalI5arrayEEbii14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">gather_qmm()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv49tensordotRK5arrayRK5arrayKi14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">tensordot()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv49tensordotRK5arrayRK5arrayRKNSt6vectorIiEERKNSt6vectorIiEE14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">tensordot()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv45outerRK5arrayRK5array14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">outer()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv45innerRK5arrayRK5array14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">inner()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv45addmm5array5array5arrayRKfRKf14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">addmm()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv415block_masked_mm5array5arrayiNSt8optionalI5arrayEENSt8optionalI5arrayEENSt8optionalI5arrayEE14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">block_masked_mm()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv415block_sparse_mm5array5arrayNSt8optionalI5arrayEENSt8optionalI5arrayEE14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">block_sparse_mm()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv49gather_mm5array5arrayNSt8optionalI5arrayEENSt8optionalI5arrayEE14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">gather_mm()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv48diagonalRK5arrayiii14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">diagonal()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv44diagRK5arrayi14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">diag()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv45traceRK5arrayiii5Dtype14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">trace()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv45traceRK5arrayiii14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">trace()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv45traceRK5array14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">trace()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv47dependsRKNSt6vectorI5arrayEERKNSt6vectorI5arrayEE"><code class="docutils literal notranslate"><span class="pre">depends()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv410atleast_1dRK5array14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">atleast_1d()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv410atleast_1dRKNSt6vectorI5arrayEE14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">atleast_1d()</span></code></a></li>
@@ -1085,6 +1112,7 @@ document.write(`
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv4lsRK5arrayRK5array"><code class="docutils literal notranslate"><span class="pre">operator&lt;&lt;()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv411right_shiftRK5arrayRK5array14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">right_shift()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv4rsRK5arrayRK5array"><code class="docutils literal notranslate"><span class="pre">operator&gt;&gt;()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv44viewRK5arrayRK5Dtype14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">view()</span></code></a></li>
</ul>
</nav>
</div>
@@ -2397,7 +2425,67 @@ document.write(`
<dl class="cpp function">
<dt class="sig sig-object cpp" id="_CPPv47scatterRK5arrayRKNSt6vectorI5arrayEERK5arrayRKNSt6vectorIiEE14StreamOrDevice">
<span id="_CPPv37scatterRK5arrayRKNSt6vectorI5arrayEERK5arrayRKNSt6vectorIiEE14StreamOrDevice"></span><span id="_CPPv27scatterRK5arrayRKNSt6vectorI5arrayEERK5arrayRKNSt6vectorIiEE14StreamOrDevice"></span><span id="scatter__arrayCR.std::vector:array:CR.arrayCR.std::vector:i:CR.StreamOrDevice"></span><span class="target" id="group__ops_1gad438be8f90bae9d37c6853b8f4225d61"></span><span class="n"><span class="pre">array</span></span><span class="w"> </span><span class="sig-name descname"><span class="n"><span class="pre">scatter</span></span></span><span class="sig-paren">(</span><span class="k"><span class="pre">const</span></span><span class="w"> </span><span class="n"><span class="pre">array</span></span><span class="w"> </span><span class="p"><span class="pre">&amp;</span></span><span class="n sig-param"><span class="pre">a</span></span>, <span class="k"><span class="pre">const</span></span><span class="w"> </span><a class="reference internal" href="#_CPPv4StRK5arraybi14StreamOrDevice" title="std"><span class="n"><span class="pre">std</span></span></a><span class="p"><span class="pre">::</span></span><span class="n"><span class="pre">vector</span></span><span class="p"><span class="pre">&lt;</span></span><span class="n"><span class="pre">array</span></span><span class="p"><span class="pre">&gt;</span></span><span class="w"> </span><span class="p"><span class="pre">&amp;</span></span><span class="n sig-param"><span class="pre">indices</span></span>, <span class="k"><span class="pre">const</span></span><span class="w"> </span><span class="n"><span class="pre">array</span></span><span class="w"> </span><span class="p"><span class="pre">&amp;</span></span><span class="n sig-param"><span class="pre">updates</span></span>, <span class="k"><span class="pre">const</span></span><span class="w"> </span><a class="reference internal" href="#_CPPv4StRK5arraybi14StreamOrDevice" title="std"><span class="n"><span class="pre">std</span></span></a><span class="p"><span class="pre">::</span></span><span class="n"><span class="pre">vector</span></span><span class="p"><span class="pre">&lt;</span></span><span class="kt"><span class="pre">int</span></span><span class="p"><span class="pre">&gt;</span></span><span class="w"> </span><span class="p"><span class="pre">&amp;</span></span><span class="n sig-param"><span class="pre">axes</span></span>, <span class="n"><span class="pre">StreamOrDevice</span></span><span class="w"> </span><span class="n sig-param"><span class="pre">s</span></span><span class="w"> </span><span class="p"><span class="pre">=</span></span><span class="w"> </span><span class="p"><span class="pre">{</span></span><span class="p"><span class="pre">}</span></span><span class="sig-paren">)</span><a class="headerlink" href="#_CPPv47scatterRK5arrayRKNSt6vectorI5arrayEERK5arrayRKNSt6vectorIiEE14StreamOrDevice" title="Link to this definition">#</a><br /></dt>
<dd><p>Scatter updates to given linear indices. </p>
<dd><p>Scatter updates to the given indices. </p>
<p>The parameters <code class="docutils literal notranslate"><span class="pre">indices</span></code> and <code class="docutils literal notranslate"><span class="pre">axes</span></code> determine the locations of <code class="docutils literal notranslate"><span class="pre">a</span></code> that are updated with the values in <code class="docutils literal notranslate"><span class="pre">updates</span></code>. Assuming 1-d <code class="docutils literal notranslate"><span class="pre">indices</span></code> for simplicity, <code class="docutils literal notranslate"><span class="pre">indices[i]</span></code> are the indices on axis <code class="docutils literal notranslate"><span class="pre">axes[i]</span></code> to which the values in <code class="docutils literal notranslate"><span class="pre">updates</span></code> will be applied. Note each array in <code class="docutils literal notranslate"><span class="pre">indices</span></code> is assigned to a corresponding axis and hence <code class="docutils literal notranslate"><span class="pre">indices.size()</span> <span class="pre">==</span> <span class="pre">axes.size()</span></code>. If an index/axis pair is not provided then indices along that axis are assumed to be zero.</p>
<p>Note the rank of <code class="docutils literal notranslate"><span class="pre">updates</span></code> must be equal to the sum of the rank of the broadcasted <code class="docutils literal notranslate"><span class="pre">indices</span></code> and the rank of <code class="docutils literal notranslate"><span class="pre">a</span></code>. In other words, assuming the arrays in <code class="docutils literal notranslate"><span class="pre">indices</span></code> have the same shape, <code class="docutils literal notranslate"><span class="pre">updates.ndim()</span> <span class="pre">==</span> <span class="pre">indices[0].ndim()</span> <span class="pre">+</span> <span class="pre">a.ndim()</span></code>. The leading dimensions of <code class="docutils literal notranslate"><span class="pre">updates</span></code> correspond to the indices, and the remaining <code class="docutils literal notranslate"><span class="pre">a.ndim()</span></code> dimensions are the values that will be applied to the given location in <code class="docutils literal notranslate"><span class="pre">a</span></code>.</p>
<p>For example:</p>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="n">auto</span> <span class="ow">in</span> <span class="o">=</span> <span class="n">zeros</span><span class="p">({</span><span class="mi">4</span><span class="p">,</span> <span class="mi">4</span><span class="p">},</span> <span class="n">float32</span><span class="p">);</span>
<span class="n">auto</span> <span class="n">indices</span> <span class="o">=</span> <span class="n">array</span><span class="p">({</span><span class="mi">2</span><span class="p">});</span>
<span class="n">auto</span> <span class="n">updates</span> <span class="o">=</span> <span class="n">reshape</span><span class="p">(</span><span class="n">arange</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="n">float32</span><span class="p">),</span> <span class="p">{</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">});</span>
<span class="n">std</span><span class="p">::</span><span class="n">vector</span><span class="o">&lt;</span><span class="nb">int</span><span class="o">&gt;</span> <span class="n">axes</span><span class="p">{</span><span class="mi">0</span><span class="p">};</span>
<span class="n">auto</span> <span class="n">out</span> <span class="o">=</span> <span class="n">scatter</span><span class="p">(</span><span class="ow">in</span><span class="p">,</span> <span class="p">{</span><span class="n">indices</span><span class="p">},</span> <span class="n">updates</span><span class="p">,</span> <span class="n">axes</span><span class="p">);</span>
</pre></div>
</div>
<p>will produce:</p>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="n">array</span><span class="p">([[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span>
<span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span>
<span class="p">[</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span>
<span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">]],</span> <span class="n">dtype</span><span class="o">=</span><span class="n">float32</span><span class="p">)</span>
</pre></div>
</div>
<p>This scatters the two-element row vector <code class="docutils literal notranslate"><span class="pre">[1,</span> <span class="pre">2]</span></code> starting at the <code class="docutils literal notranslate"><span class="pre">(2,</span> <span class="pre">0)</span></code> position of <code class="docutils literal notranslate"><span class="pre">a</span></code>.</p>
<p>Adding another element to <code class="docutils literal notranslate"><span class="pre">indices</span></code> will scatter into another location of <code class="docutils literal notranslate"><span class="pre">a</span></code>. We also have to add an another update for the new index:</p>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="n">auto</span> <span class="ow">in</span> <span class="o">=</span> <span class="n">zeros</span><span class="p">({</span><span class="mi">4</span><span class="p">,</span> <span class="mi">4</span><span class="p">},</span> <span class="n">float32</span><span class="p">);</span>
<span class="n">auto</span> <span class="n">indices</span> <span class="o">=</span> <span class="n">array</span><span class="p">({</span><span class="mi">2</span><span class="p">,</span> <span class="mi">0</span><span class="p">});</span>
<span class="n">auto</span> <span class="n">updates</span> <span class="o">=</span> <span class="n">reshape</span><span class="p">(</span><span class="n">arange</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">5</span><span class="p">,</span> <span class="n">float32</span><span class="p">),</span> <span class="p">{</span><span class="mi">2</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">});</span>
<span class="n">std</span><span class="p">::</span><span class="n">vector</span><span class="o">&lt;</span><span class="nb">int</span><span class="o">&gt;</span> <span class="n">axes</span><span class="p">{</span><span class="mi">0</span><span class="p">};</span>
<span class="n">auto</span> <span class="n">out</span> <span class="o">=</span> <span class="n">scatter</span><span class="p">(</span><span class="ow">in</span><span class="p">,</span> <span class="p">{</span><span class="n">indices</span><span class="p">},</span> <span class="n">updates</span><span class="p">,</span> <span class="n">axes</span><span class="p">):</span>
</pre></div>
</div>
<p>will produce:</p>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="n">array</span><span class="p">([[</span><span class="mi">3</span><span class="p">,</span> <span class="mi">4</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span>
<span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span>
<span class="p">[</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span>
<span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">]],</span> <span class="n">dtype</span><span class="o">=</span><span class="n">float32</span><span class="p">)</span>
</pre></div>
</div>
<p>To control the scatter location on an additional axis, add another index array to <code class="docutils literal notranslate"><span class="pre">indices</span></code> and another axis to <code class="docutils literal notranslate"><span class="pre">axes</span></code>:</p>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="n">auto</span> <span class="ow">in</span> <span class="o">=</span> <span class="n">zeros</span><span class="p">({</span><span class="mi">4</span><span class="p">,</span> <span class="mi">4</span><span class="p">},</span> <span class="n">float32</span><span class="p">);</span>
<span class="n">auto</span> <span class="n">indices</span> <span class="o">=</span> <span class="n">std</span><span class="p">::</span><span class="n">vector</span><span class="p">{</span><span class="n">array</span><span class="p">({</span><span class="mi">2</span><span class="p">,</span> <span class="mi">0</span><span class="p">}),</span> <span class="n">array</span><span class="p">({</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">})};</span>
<span class="n">auto</span> <span class="n">updates</span> <span class="o">=</span> <span class="n">reshape</span><span class="p">(</span><span class="n">arange</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">5</span><span class="p">,</span> <span class="n">float32</span><span class="p">),</span> <span class="p">{</span><span class="mi">2</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">});</span>
<span class="n">std</span><span class="p">::</span><span class="n">vector</span><span class="o">&lt;</span><span class="nb">int</span><span class="o">&gt;</span> <span class="n">axes</span><span class="p">{</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">};</span>
<span class="n">auto</span> <span class="n">out</span> <span class="o">=</span> <span class="n">scatter</span><span class="p">(</span><span class="ow">in</span><span class="p">,</span> <span class="n">indices</span><span class="p">,</span> <span class="n">updates</span><span class="p">,</span> <span class="n">axes</span><span class="p">);</span>
</pre></div>
</div>
<p>will produce:</p>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="n">array</span><span class="p">([[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">4</span><span class="p">],</span>
<span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span>
<span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span>
<span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">]],</span> <span class="n">dtype</span><span class="o">=</span><span class="n">float32</span><span class="p">)</span>
</pre></div>
</div>
<p>Items in indices are broadcasted together. This means:</p>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="n">auto</span> <span class="n">indices</span> <span class="o">=</span> <span class="n">std</span><span class="p">::</span><span class="n">vector</span><span class="p">{</span><span class="n">array</span><span class="p">({</span><span class="mi">2</span><span class="p">,</span> <span class="mi">0</span><span class="p">}),</span> <span class="n">array</span><span class="p">({</span><span class="mi">1</span><span class="p">})};</span>
</pre></div>
</div>
<p>is equivalent to:</p>
<div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="n">auto</span> <span class="n">indices</span> <span class="o">=</span> <span class="n">std</span><span class="p">::</span><span class="n">vector</span><span class="p">{</span><span class="n">array</span><span class="p">({</span><span class="mi">2</span><span class="p">,</span> <span class="mi">0</span><span class="p">}),</span> <span class="n">array</span><span class="p">({</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">})};</span>
</pre></div>
</div>
<p>Note, <code class="docutils literal notranslate"><span class="pre">scatter</span></code> does not perform bounds checking on the indices and updates. Out-of-bounds accesses on <code class="docutils literal notranslate"><span class="pre">a</span></code> are undefined and typically result in unintended or invalid memory writes. </p>
</dd></dl>
<dl class="cpp function">
@@ -2558,8 +2646,8 @@ document.write(`
</dd></dl>
<dl class="cpp function">
<dt class="sig sig-object cpp" id="_CPPv416block_sparse_qmmRK5arrayRK5arrayRK5arrayRK5arrayNSt8optionalI5arrayEENSt8optionalI5arrayEEbii14StreamOrDevice">
<span id="_CPPv316block_sparse_qmmRK5arrayRK5arrayRK5arrayRK5arrayNSt8optionalI5arrayEENSt8optionalI5arrayEEbii14StreamOrDevice"></span><span id="_CPPv216block_sparse_qmmRK5arrayRK5arrayRK5arrayRK5arrayNSt8optionalI5arrayEENSt8optionalI5arrayEEbii14StreamOrDevice"></span><span id="block_sparse_qmm__arrayCR.arrayCR.arrayCR.arrayCR.std::optional:array:.std::optional:array:.b.i.i.StreamOrDevice"></span><span class="target" id="group__ops_1ga9aaa4381942a6d9a31d2ab8af8436e78"></span><span class="n"><span class="pre">array</span></span><span class="w"> </span><span class="sig-name descname"><span class="n"><span class="pre">block_sparse_qmm</span></span></span><span class="sig-paren">(</span><span class="k"><span class="pre">const</span></span><span class="w"> </span><span class="n"><span class="pre">array</span></span><span class="w"> </span><span class="p"><span class="pre">&amp;</span></span><span class="n sig-param"><span class="pre">x</span></span>, <span class="k"><span class="pre">const</span></span><span class="w"> </span><span class="n"><span class="pre">array</span></span><span class="w"> </span><span class="p"><span class="pre">&amp;</span></span><span class="n sig-param"><span class="pre">w</span></span>, <span class="k"><span class="pre">const</span></span><span class="w"> </span><span class="n"><span class="pre">array</span></span><span class="w"> </span><span class="p"><span class="pre">&amp;</span></span><span class="n sig-param"><span class="pre">scales</span></span>, <span class="k"><span class="pre">const</span></span><span class="w"> </span><span class="n"><span class="pre">array</span></span><span class="w"> </span><span class="p"><span class="pre">&amp;</span></span><span class="n sig-param"><span class="pre">biases</span></span>, <a class="reference internal" href="#_CPPv4StRK5arraybi14StreamOrDevice" title="std"><span class="n"><span class="pre">std</span></span></a><span class="p"><span class="pre">::</span></span><span class="n"><span class="pre">optional</span></span><span class="p"><span class="pre">&lt;</span></span><span class="n"><span class="pre">array</span></span><span class="p"><span class="pre">&gt;</span></span><span class="w"> </span><span class="n sig-param"><span class="pre">lhs_indices</span></span><span class="w"> </span><span class="p"><span class="pre">=</span></span><span class="w"> </span><a class="reference internal" href="#_CPPv4StRK5arraybi14StreamOrDevice" title="std"><span class="n"><span class="pre">std</span></span></a><span class="p"><span class="pre">::</span></span><span class="n"><span class="pre">nullopt</span></span>, <a class="reference internal" href="#_CPPv4StRK5arraybi14StreamOrDevice" title="std"><span class="n"><span class="pre">std</span></span></a><span class="p"><span class="pre">::</span></span><span class="n"><span class="pre">optional</span></span><span class="p"><span class="pre">&lt;</span></span><span class="n"><span class="pre">array</span></span><span class="p"><span class="pre">&gt;</span></span><span class="w"> </span><span class="n sig-param"><span class="pre">rhs_indices</span></span><span class="w"> </span><span class="p"><span class="pre">=</span></span><span class="w"> </span><a class="reference internal" href="#_CPPv4StRK5arraybi14StreamOrDevice" title="std"><span class="n"><span class="pre">std</span></span></a><span class="p"><span class="pre">::</span></span><span class="n"><span class="pre">nullopt</span></span>, <span class="kt"><span class="pre">bool</span></span><span class="w"> </span><span class="n sig-param"><span class="pre">transpose</span></span><span class="w"> </span><span class="p"><span class="pre">=</span></span><span class="w"> </span><span class="k"><span class="pre">true</span></span>, <span class="kt"><span class="pre">int</span></span><span class="w"> </span><span class="n sig-param"><span class="pre">group_size</span></span><span class="w"> </span><span class="p"><span class="pre">=</span></span><span class="w"> </span><span class="m"><span class="pre">64</span></span>, <span class="kt"><span class="pre">int</span></span><span class="w"> </span><span class="n sig-param"><span class="pre">bits</span></span><span class="w"> </span><span class="p"><span class="pre">=</span></span><span class="w"> </span><span class="m"><span class="pre">4</span></span>, <span class="n"><span class="pre">StreamOrDevice</span></span><span class="w"> </span><span class="n sig-param"><span class="pre">s</span></span><span class="w"> </span><span class="p"><span class="pre">=</span></span><span class="w"> </span><span class="p"><span class="pre">{</span></span><span class="p"><span class="pre">}</span></span><span class="sig-paren">)</span><a class="headerlink" href="#_CPPv416block_sparse_qmmRK5arrayRK5arrayRK5arrayRK5arrayNSt8optionalI5arrayEENSt8optionalI5arrayEEbii14StreamOrDevice" title="Link to this definition">#</a><br /></dt>
<dt class="sig sig-object cpp" id="_CPPv410gather_qmmRK5arrayRK5arrayRK5arrayRK5arrayNSt8optionalI5arrayEENSt8optionalI5arrayEEbii14StreamOrDevice">
<span id="_CPPv310gather_qmmRK5arrayRK5arrayRK5arrayRK5arrayNSt8optionalI5arrayEENSt8optionalI5arrayEEbii14StreamOrDevice"></span><span id="_CPPv210gather_qmmRK5arrayRK5arrayRK5arrayRK5arrayNSt8optionalI5arrayEENSt8optionalI5arrayEEbii14StreamOrDevice"></span><span id="gather_qmm__arrayCR.arrayCR.arrayCR.arrayCR.std::optional:array:.std::optional:array:.b.i.i.StreamOrDevice"></span><span class="target" id="group__ops_1ga368a0dc0e5dfb76922e7aa55a95f12f0"></span><span class="n"><span class="pre">array</span></span><span class="w"> </span><span class="sig-name descname"><span class="n"><span class="pre">gather_qmm</span></span></span><span class="sig-paren">(</span><span class="k"><span class="pre">const</span></span><span class="w"> </span><span class="n"><span class="pre">array</span></span><span class="w"> </span><span class="p"><span class="pre">&amp;</span></span><span class="n sig-param"><span class="pre">x</span></span>, <span class="k"><span class="pre">const</span></span><span class="w"> </span><span class="n"><span class="pre">array</span></span><span class="w"> </span><span class="p"><span class="pre">&amp;</span></span><span class="n sig-param"><span class="pre">w</span></span>, <span class="k"><span class="pre">const</span></span><span class="w"> </span><span class="n"><span class="pre">array</span></span><span class="w"> </span><span class="p"><span class="pre">&amp;</span></span><span class="n sig-param"><span class="pre">scales</span></span>, <span class="k"><span class="pre">const</span></span><span class="w"> </span><span class="n"><span class="pre">array</span></span><span class="w"> </span><span class="p"><span class="pre">&amp;</span></span><span class="n sig-param"><span class="pre">biases</span></span>, <a class="reference internal" href="#_CPPv4StRK5arraybi14StreamOrDevice" title="std"><span class="n"><span class="pre">std</span></span></a><span class="p"><span class="pre">::</span></span><span class="n"><span class="pre">optional</span></span><span class="p"><span class="pre">&lt;</span></span><span class="n"><span class="pre">array</span></span><span class="p"><span class="pre">&gt;</span></span><span class="w"> </span><span class="n sig-param"><span class="pre">lhs_indices</span></span><span class="w"> </span><span class="p"><span class="pre">=</span></span><span class="w"> </span><a class="reference internal" href="#_CPPv4StRK5arraybi14StreamOrDevice" title="std"><span class="n"><span class="pre">std</span></span></a><span class="p"><span class="pre">::</span></span><span class="n"><span class="pre">nullopt</span></span>, <a class="reference internal" href="#_CPPv4StRK5arraybi14StreamOrDevice" title="std"><span class="n"><span class="pre">std</span></span></a><span class="p"><span class="pre">::</span></span><span class="n"><span class="pre">optional</span></span><span class="p"><span class="pre">&lt;</span></span><span class="n"><span class="pre">array</span></span><span class="p"><span class="pre">&gt;</span></span><span class="w"> </span><span class="n sig-param"><span class="pre">rhs_indices</span></span><span class="w"> </span><span class="p"><span class="pre">=</span></span><span class="w"> </span><a class="reference internal" href="#_CPPv4StRK5arraybi14StreamOrDevice" title="std"><span class="n"><span class="pre">std</span></span></a><span class="p"><span class="pre">::</span></span><span class="n"><span class="pre">nullopt</span></span>, <span class="kt"><span class="pre">bool</span></span><span class="w"> </span><span class="n sig-param"><span class="pre">transpose</span></span><span class="w"> </span><span class="p"><span class="pre">=</span></span><span class="w"> </span><span class="k"><span class="pre">true</span></span>, <span class="kt"><span class="pre">int</span></span><span class="w"> </span><span class="n sig-param"><span class="pre">group_size</span></span><span class="w"> </span><span class="p"><span class="pre">=</span></span><span class="w"> </span><span class="m"><span class="pre">64</span></span>, <span class="kt"><span class="pre">int</span></span><span class="w"> </span><span class="n sig-param"><span class="pre">bits</span></span><span class="w"> </span><span class="p"><span class="pre">=</span></span><span class="w"> </span><span class="m"><span class="pre">4</span></span>, <span class="n"><span class="pre">StreamOrDevice</span></span><span class="w"> </span><span class="n sig-param"><span class="pre">s</span></span><span class="w"> </span><span class="p"><span class="pre">=</span></span><span class="w"> </span><span class="p"><span class="pre">{</span></span><span class="p"><span class="pre">}</span></span><span class="sig-paren">)</span><a class="headerlink" href="#_CPPv410gather_qmmRK5arrayRK5arrayRK5arrayRK5arrayNSt8optionalI5arrayEENSt8optionalI5arrayEEbii14StreamOrDevice" title="Link to this definition">#</a><br /></dt>
<dd><p>Compute matrix products with matrix-level gather. </p>
</dd></dl>
@@ -2599,8 +2687,8 @@ document.write(`
</dd></dl>
<dl class="cpp function">
<dt class="sig sig-object cpp" id="_CPPv415block_sparse_mm5array5arrayNSt8optionalI5arrayEENSt8optionalI5arrayEE14StreamOrDevice">
<span id="_CPPv315block_sparse_mm5array5arrayNSt8optionalI5arrayEENSt8optionalI5arrayEE14StreamOrDevice"></span><span id="_CPPv215block_sparse_mm5array5arrayNSt8optionalI5arrayEENSt8optionalI5arrayEE14StreamOrDevice"></span><span id="block_sparse_mm__array.array.std::optional:array:.std::optional:array:.StreamOrDevice"></span><span class="target" id="group__ops_1gaf5d82380cd204e7c33324cebcd1461ce"></span><span class="n"><span class="pre">array</span></span><span class="w"> </span><span class="sig-name descname"><span class="n"><span class="pre">block_sparse_mm</span></span></span><span class="sig-paren">(</span><span class="n"><span class="pre">array</span></span><span class="w"> </span><span class="n sig-param"><span class="pre">a</span></span>, <span class="n"><span class="pre">array</span></span><span class="w"> </span><span class="n sig-param"><span class="pre">b</span></span>, <a class="reference internal" href="#_CPPv4StRK5arraybi14StreamOrDevice" title="std"><span class="n"><span class="pre">std</span></span></a><span class="p"><span class="pre">::</span></span><span class="n"><span class="pre">optional</span></span><span class="p"><span class="pre">&lt;</span></span><span class="n"><span class="pre">array</span></span><span class="p"><span class="pre">&gt;</span></span><span class="w"> </span><span class="n sig-param"><span class="pre">lhs_indices</span></span><span class="w"> </span><span class="p"><span class="pre">=</span></span><span class="w"> </span><a class="reference internal" href="#_CPPv4StRK5arraybi14StreamOrDevice" title="std"><span class="n"><span class="pre">std</span></span></a><span class="p"><span class="pre">::</span></span><span class="n"><span class="pre">nullopt</span></span>, <a class="reference internal" href="#_CPPv4StRK5arraybi14StreamOrDevice" title="std"><span class="n"><span class="pre">std</span></span></a><span class="p"><span class="pre">::</span></span><span class="n"><span class="pre">optional</span></span><span class="p"><span class="pre">&lt;</span></span><span class="n"><span class="pre">array</span></span><span class="p"><span class="pre">&gt;</span></span><span class="w"> </span><span class="n sig-param"><span class="pre">rhs_indices</span></span><span class="w"> </span><span class="p"><span class="pre">=</span></span><span class="w"> </span><a class="reference internal" href="#_CPPv4StRK5arraybi14StreamOrDevice" title="std"><span class="n"><span class="pre">std</span></span></a><span class="p"><span class="pre">::</span></span><span class="n"><span class="pre">nullopt</span></span>, <span class="n"><span class="pre">StreamOrDevice</span></span><span class="w"> </span><span class="n sig-param"><span class="pre">s</span></span><span class="w"> </span><span class="p"><span class="pre">=</span></span><span class="w"> </span><span class="p"><span class="pre">{</span></span><span class="p"><span class="pre">}</span></span><span class="sig-paren">)</span><a class="headerlink" href="#_CPPv415block_sparse_mm5array5arrayNSt8optionalI5arrayEENSt8optionalI5arrayEE14StreamOrDevice" title="Link to this definition">#</a><br /></dt>
<dt class="sig sig-object cpp" id="_CPPv49gather_mm5array5arrayNSt8optionalI5arrayEENSt8optionalI5arrayEE14StreamOrDevice">
<span id="_CPPv39gather_mm5array5arrayNSt8optionalI5arrayEENSt8optionalI5arrayEE14StreamOrDevice"></span><span id="_CPPv29gather_mm5array5arrayNSt8optionalI5arrayEENSt8optionalI5arrayEE14StreamOrDevice"></span><span id="gather_mm__array.array.std::optional:array:.std::optional:array:.StreamOrDevice"></span><span class="target" id="group__ops_1ga8d50480266d258cac40ff51bcb0fc6a7"></span><span class="n"><span class="pre">array</span></span><span class="w"> </span><span class="sig-name descname"><span class="n"><span class="pre">gather_mm</span></span></span><span class="sig-paren">(</span><span class="n"><span class="pre">array</span></span><span class="w"> </span><span class="n sig-param"><span class="pre">a</span></span>, <span class="n"><span class="pre">array</span></span><span class="w"> </span><span class="n sig-param"><span class="pre">b</span></span>, <a class="reference internal" href="#_CPPv4StRK5arraybi14StreamOrDevice" title="std"><span class="n"><span class="pre">std</span></span></a><span class="p"><span class="pre">::</span></span><span class="n"><span class="pre">optional</span></span><span class="p"><span class="pre">&lt;</span></span><span class="n"><span class="pre">array</span></span><span class="p"><span class="pre">&gt;</span></span><span class="w"> </span><span class="n sig-param"><span class="pre">lhs_indices</span></span><span class="w"> </span><span class="p"><span class="pre">=</span></span><span class="w"> </span><a class="reference internal" href="#_CPPv4StRK5arraybi14StreamOrDevice" title="std"><span class="n"><span class="pre">std</span></span></a><span class="p"><span class="pre">::</span></span><span class="n"><span class="pre">nullopt</span></span>, <a class="reference internal" href="#_CPPv4StRK5arraybi14StreamOrDevice" title="std"><span class="n"><span class="pre">std</span></span></a><span class="p"><span class="pre">::</span></span><span class="n"><span class="pre">optional</span></span><span class="p"><span class="pre">&lt;</span></span><span class="n"><span class="pre">array</span></span><span class="p"><span class="pre">&gt;</span></span><span class="w"> </span><span class="n sig-param"><span class="pre">rhs_indices</span></span><span class="w"> </span><span class="p"><span class="pre">=</span></span><span class="w"> </span><a class="reference internal" href="#_CPPv4StRK5arraybi14StreamOrDevice" title="std"><span class="n"><span class="pre">std</span></span></a><span class="p"><span class="pre">::</span></span><span class="n"><span class="pre">nullopt</span></span>, <span class="n"><span class="pre">StreamOrDevice</span></span><span class="w"> </span><span class="n sig-param"><span class="pre">s</span></span><span class="w"> </span><span class="p"><span class="pre">=</span></span><span class="w"> </span><span class="p"><span class="pre">{</span></span><span class="p"><span class="pre">}</span></span><span class="sig-paren">)</span><a class="headerlink" href="#_CPPv49gather_mm5array5arrayNSt8optionalI5arrayEENSt8optionalI5arrayEE14StreamOrDevice" title="Link to this definition">#</a><br /></dt>
<dd><p>Compute matrix product with matrix-level gather. </p>
</dd></dl>
@@ -2616,6 +2704,22 @@ document.write(`
<dd><p>Extract diagonal from a 2d array or create a diagonal matrix. </p>
</dd></dl>
<dl class="cpp function">
<dt class="sig sig-object cpp" id="_CPPv45traceRK5arrayiii5Dtype14StreamOrDevice">
<span id="_CPPv35traceRK5arrayiii5Dtype14StreamOrDevice"></span><span id="_CPPv25traceRK5arrayiii5Dtype14StreamOrDevice"></span><span id="trace__arrayCR.i.i.i.Dtype.StreamOrDevice"></span><span class="target" id="group__ops_1gabf786129c7660ed8d5acb5499bc6fefd"></span><span class="n"><span class="pre">array</span></span><span class="w"> </span><span class="sig-name descname"><span class="n"><span class="pre">trace</span></span></span><span class="sig-paren">(</span><span class="k"><span class="pre">const</span></span><span class="w"> </span><span class="n"><span class="pre">array</span></span><span class="w"> </span><span class="p"><span class="pre">&amp;</span></span><span class="n sig-param"><span class="pre">a</span></span>, <span class="kt"><span class="pre">int</span></span><span class="w"> </span><span class="n sig-param"><span class="pre">offset</span></span>, <span class="kt"><span class="pre">int</span></span><span class="w"> </span><span class="n sig-param"><span class="pre">axis1</span></span>, <span class="kt"><span class="pre">int</span></span><span class="w"> </span><span class="n sig-param"><span class="pre">axis2</span></span>, <span class="n"><span class="pre">Dtype</span></span><span class="w"> </span><span class="n sig-param"><span class="pre">dtype</span></span>, <span class="n"><span class="pre">StreamOrDevice</span></span><span class="w"> </span><span class="n sig-param"><span class="pre">s</span></span><span class="w"> </span><span class="p"><span class="pre">=</span></span><span class="w"> </span><span class="p"><span class="pre">{</span></span><span class="p"><span class="pre">}</span></span><span class="sig-paren">)</span><a class="headerlink" href="#_CPPv45traceRK5arrayiii5Dtype14StreamOrDevice" title="Link to this definition">#</a><br /></dt>
<dd><p>Return the sum along a specified diagonal in the given array. </p>
</dd></dl>
<dl class="cpp function">
<dt class="sig sig-object cpp" id="_CPPv45traceRK5arrayiii14StreamOrDevice">
<span id="_CPPv35traceRK5arrayiii14StreamOrDevice"></span><span id="_CPPv25traceRK5arrayiii14StreamOrDevice"></span><span id="trace__arrayCR.i.i.i.StreamOrDevice"></span><span class="target" id="group__ops_1ga5ed43c2dbf7d6cbddbaa2fd682deaafd"></span><span class="n"><span class="pre">array</span></span><span class="w"> </span><span class="sig-name descname"><span class="n"><span class="pre">trace</span></span></span><span class="sig-paren">(</span><span class="k"><span class="pre">const</span></span><span class="w"> </span><span class="n"><span class="pre">array</span></span><span class="w"> </span><span class="p"><span class="pre">&amp;</span></span><span class="n sig-param"><span class="pre">a</span></span>, <span class="kt"><span class="pre">int</span></span><span class="w"> </span><span class="n sig-param"><span class="pre">offset</span></span>, <span class="kt"><span class="pre">int</span></span><span class="w"> </span><span class="n sig-param"><span class="pre">axis1</span></span>, <span class="kt"><span class="pre">int</span></span><span class="w"> </span><span class="n sig-param"><span class="pre">axis2</span></span>, <span class="n"><span class="pre">StreamOrDevice</span></span><span class="w"> </span><span class="n sig-param"><span class="pre">s</span></span><span class="w"> </span><span class="p"><span class="pre">=</span></span><span class="w"> </span><span class="p"><span class="pre">{</span></span><span class="p"><span class="pre">}</span></span><span class="sig-paren">)</span><a class="headerlink" href="#_CPPv45traceRK5arrayiii14StreamOrDevice" title="Link to this definition">#</a><br /></dt>
<dd></dd></dl>
<dl class="cpp function">
<dt class="sig sig-object cpp" id="_CPPv45traceRK5array14StreamOrDevice">
<span id="_CPPv35traceRK5array14StreamOrDevice"></span><span id="_CPPv25traceRK5array14StreamOrDevice"></span><span id="trace__arrayCR.StreamOrDevice"></span><span class="target" id="group__ops_1gaf25c00108feaafaa6350a4434cb0062e"></span><span class="n"><span class="pre">array</span></span><span class="w"> </span><span class="sig-name descname"><span class="n"><span class="pre">trace</span></span></span><span class="sig-paren">(</span><span class="k"><span class="pre">const</span></span><span class="w"> </span><span class="n"><span class="pre">array</span></span><span class="w"> </span><span class="p"><span class="pre">&amp;</span></span><span class="n sig-param"><span class="pre">a</span></span>, <span class="n"><span class="pre">StreamOrDevice</span></span><span class="w"> </span><span class="n sig-param"><span class="pre">s</span></span><span class="w"> </span><span class="p"><span class="pre">=</span></span><span class="w"> </span><span class="p"><span class="pre">{</span></span><span class="p"><span class="pre">}</span></span><span class="sig-paren">)</span><a class="headerlink" href="#_CPPv45traceRK5array14StreamOrDevice" title="Link to this definition">#</a><br /></dt>
<dd></dd></dl>
<dl class="cpp function">
<dt class="sig sig-object cpp" id="_CPPv47dependsRKNSt6vectorI5arrayEERKNSt6vectorI5arrayEE">
<span id="_CPPv37dependsRKNSt6vectorI5arrayEERKNSt6vectorI5arrayEE"></span><span id="_CPPv27dependsRKNSt6vectorI5arrayEERKNSt6vectorI5arrayEE"></span><span id="depends__std::vector:array:CR.std::vector:array:CR"></span><span class="target" id="group__ops_1gac4a51a68fbe1725436b026d2fbb95759"></span><a class="reference internal" href="#_CPPv4StRK5arraybi14StreamOrDevice" title="std"><span class="n"><span class="pre">std</span></span></a><span class="p"><span class="pre">::</span></span><span class="n"><span class="pre">vector</span></span><span class="p"><span class="pre">&lt;</span></span><span class="n"><span class="pre">array</span></span><span class="p"><span class="pre">&gt;</span></span><span class="w"> </span><span class="sig-name descname"><span class="n"><span class="pre">depends</span></span></span><span class="sig-paren">(</span><span class="k"><span class="pre">const</span></span><span class="w"> </span><a class="reference internal" href="#_CPPv4StRK5arraybi14StreamOrDevice" title="std"><span class="n"><span class="pre">std</span></span></a><span class="p"><span class="pre">::</span></span><span class="n"><span class="pre">vector</span></span><span class="p"><span class="pre">&lt;</span></span><span class="n"><span class="pre">array</span></span><span class="p"><span class="pre">&gt;</span></span><span class="w"> </span><span class="p"><span class="pre">&amp;</span></span><span class="n sig-param"><span class="pre">inputs</span></span>, <span class="k"><span class="pre">const</span></span><span class="w"> </span><a class="reference internal" href="#_CPPv4StRK5arraybi14StreamOrDevice" title="std"><span class="n"><span class="pre">std</span></span></a><span class="p"><span class="pre">::</span></span><span class="n"><span class="pre">vector</span></span><span class="p"><span class="pre">&lt;</span></span><span class="n"><span class="pre">array</span></span><span class="p"><span class="pre">&gt;</span></span><span class="w"> </span><span class="p"><span class="pre">&amp;</span></span><span class="n sig-param"><span class="pre">dependencies</span></span><span class="sig-paren">)</span><a class="headerlink" href="#_CPPv47dependsRKNSt6vectorI5arrayEERKNSt6vectorI5arrayEE" title="Link to this definition">#</a><br /></dt>
@@ -2721,6 +2825,11 @@ document.write(`
<span id="_CPPv3rsRK5arrayRK5array"></span><span id="_CPPv2rsRK5arrayRK5array"></span><span id="rshift-operator__arrayCR.arrayCR"></span><span class="target" id="group__ops_1ga498b61f7e8f056ae00297fa0dc17303a"></span><span class="n"><span class="pre">array</span></span><span class="w"> </span><span class="sig-name descname"><span class="k"><span class="pre">operator</span></span><span class="o"><span class="pre">&gt;&gt;</span></span></span><span class="sig-paren">(</span><span class="k"><span class="pre">const</span></span><span class="w"> </span><span class="n"><span class="pre">array</span></span><span class="w"> </span><span class="p"><span class="pre">&amp;</span></span><span class="n sig-param"><span class="pre">a</span></span>, <span class="k"><span class="pre">const</span></span><span class="w"> </span><span class="n"><span class="pre">array</span></span><span class="w"> </span><span class="p"><span class="pre">&amp;</span></span><span class="n sig-param"><span class="pre">b</span></span><span class="sig-paren">)</span><a class="headerlink" href="#_CPPv4rsRK5arrayRK5array" title="Link to this definition">#</a><br /></dt>
<dd></dd></dl>
<dl class="cpp function">
<dt class="sig sig-object cpp" id="_CPPv44viewRK5arrayRK5Dtype14StreamOrDevice">
<span id="_CPPv34viewRK5arrayRK5Dtype14StreamOrDevice"></span><span id="_CPPv24viewRK5arrayRK5Dtype14StreamOrDevice"></span><span id="view__arrayCR.DtypeCR.StreamOrDevice"></span><span class="target" id="group__ops_1ga3602aa91b7b124a0b41ec1b2137a1b02"></span><span class="n"><span class="pre">array</span></span><span class="w"> </span><span class="sig-name descname"><span class="n"><span class="pre">view</span></span></span><span class="sig-paren">(</span><span class="k"><span class="pre">const</span></span><span class="w"> </span><span class="n"><span class="pre">array</span></span><span class="w"> </span><span class="p"><span class="pre">&amp;</span></span><span class="n sig-param"><span class="pre">a</span></span>, <span class="k"><span class="pre">const</span></span><span class="w"> </span><span class="n"><span class="pre">Dtype</span></span><span class="w"> </span><span class="p"><span class="pre">&amp;</span></span><span class="n sig-param"><span class="pre">dtype</span></span>, <span class="n"><span class="pre">StreamOrDevice</span></span><span class="w"> </span><span class="n sig-param"><span class="pre">s</span></span><span class="w"> </span><span class="p"><span class="pre">=</span></span><span class="w"> </span><span class="p"><span class="pre">{</span></span><span class="p"><span class="pre">}</span></span><span class="sig-paren">)</span><a class="headerlink" href="#_CPPv44viewRK5arrayRK5Dtype14StreamOrDevice" title="Link to this definition">#</a><br /></dt>
<dd></dd></dl>
</section>
@@ -3025,16 +3134,19 @@ document.write(`
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv416quantized_matmulRK5arrayRK5arrayRK5arrayRK5arraybii14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">quantized_matmul()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv48quantizeRK5arrayii14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">quantize()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv410dequantizeRK5arrayRK5arrayRK5arrayii14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">dequantize()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv416block_sparse_qmmRK5arrayRK5arrayRK5arrayRK5arrayNSt8optionalI5arrayEENSt8optionalI5arrayEEbii14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">block_sparse_qmm()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv410gather_qmmRK5arrayRK5arrayRK5arrayRK5arrayNSt8optionalI5arrayEENSt8optionalI5arrayEEbii14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">gather_qmm()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv49tensordotRK5arrayRK5arrayKi14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">tensordot()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv49tensordotRK5arrayRK5arrayRKNSt6vectorIiEERKNSt6vectorIiEE14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">tensordot()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv45outerRK5arrayRK5array14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">outer()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv45innerRK5arrayRK5array14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">inner()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv45addmm5array5array5arrayRKfRKf14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">addmm()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv415block_masked_mm5array5arrayiNSt8optionalI5arrayEENSt8optionalI5arrayEENSt8optionalI5arrayEE14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">block_masked_mm()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv415block_sparse_mm5array5arrayNSt8optionalI5arrayEENSt8optionalI5arrayEE14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">block_sparse_mm()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv49gather_mm5array5arrayNSt8optionalI5arrayEENSt8optionalI5arrayEE14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">gather_mm()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv48diagonalRK5arrayiii14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">diagonal()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv44diagRK5arrayi14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">diag()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv45traceRK5arrayiii5Dtype14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">trace()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv45traceRK5arrayiii14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">trace()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv45traceRK5array14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">trace()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv47dependsRKNSt6vectorI5arrayEERKNSt6vectorI5arrayEE"><code class="docutils literal notranslate"><span class="pre">depends()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv410atleast_1dRK5array14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">atleast_1d()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv410atleast_1dRKNSt6vectorI5arrayEE14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">atleast_1d()</span></code></a></li>
@@ -3054,6 +3166,7 @@ document.write(`
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv4lsRK5arrayRK5array"><code class="docutils literal notranslate"><span class="pre">operator&lt;&lt;()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv411right_shiftRK5arrayRK5array14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">right_shift()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv4rsRK5arrayRK5array"><code class="docutils literal notranslate"><span class="pre">operator&gt;&gt;()</span></code></a></li>
<li class="toc-h2 nav-item toc-entry"><a class="reference internal nav-link" href="#_CPPv44viewRK5arrayRK5Dtype14StreamOrDevice"><code class="docutils literal notranslate"><span class="pre">view()</span></code></a></li>
</ul>
</nav></div>