mirror of
https://github.com/ml-explore/mlx.git
synced 2025-11-07 12:48:15 +08:00
rebase
This commit is contained in:
CircleCI Docs
265
docs/build/html/fast__primitives_8h_source.html
vendored
265
docs/build/html/fast__primitives_8h_source.html
vendored
@@ -353,142 +353,144 @@ $(function(){initNavTree('fast__primitives_8h_source.html',''); initResizable(tr
|
||||
</div>
|
||||
<div class="line"><a id="l00203" name="l00203"></a><span class="lineno"> 203</span> </div>
|
||||
<div class="foldopen" id="foldopen00204" data-start="{" data-end="};">
|
||||
<div class="line"><a id="l00204" name="l00204"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html"> 204</a></span><span class="keyword">class </span><a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#ab3f78d30e5bb3e76cfe701f2358e4748">ScaledDotProductAttention</a> : <span class="keyword">public</span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_custom.html#a4186fea23f7156c38960426821fca313">Custom</a> {</div>
|
||||
<div class="line"><a id="l00204" name="l00204"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html"> 204</a></span><span class="keyword">class </span><a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#a09c99b460cca606b2ebb22f90b3d13a2">ScaledDotProductAttention</a> : <span class="keyword">public</span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_custom.html#a4186fea23f7156c38960426821fca313">Custom</a> {</div>
|
||||
<div class="line"><a id="l00205" name="l00205"></a><span class="lineno"> 205</span> <span class="keyword">public</span>:</div>
|
||||
<div class="foldopen" id="foldopen00206" data-start="{" data-end="}">
|
||||
<div class="line"><a id="l00206" name="l00206"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#ab3f78d30e5bb3e76cfe701f2358e4748"> 206</a></span> <span class="keyword">explicit</span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#ab3f78d30e5bb3e76cfe701f2358e4748">ScaledDotProductAttention</a>(</div>
|
||||
<div class="line"><a id="l00206" name="l00206"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#a09c99b460cca606b2ebb22f90b3d13a2"> 206</a></span> <span class="keyword">explicit</span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#a09c99b460cca606b2ebb22f90b3d13a2">ScaledDotProductAttention</a>(</div>
|
||||
<div class="line"><a id="l00207" name="l00207"></a><span class="lineno"> 207</span> <a class="code hl_struct" href="structmlx_1_1core_1_1_stream.html">Stream</a> <a class="code hl_function" href="classmlx_1_1core_1_1_primitive.html#a46e6257397a662528f9f831842ac456a">stream</a>,</div>
|
||||
<div class="line"><a id="l00208" name="l00208"></a><span class="lineno"> 208</span> std::function<std::vector<array>(std::vector<array>)> fallback,</div>
|
||||
<div class="line"><a id="l00209" name="l00209"></a><span class="lineno"> 209</span> <span class="keyword">const</span> <span class="keywordtype">float</span> scale)</div>
|
||||
<div class="line"><a id="l00210" name="l00210"></a><span class="lineno"> 210</span> : <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_custom.html#a4186fea23f7156c38960426821fca313">Custom</a>(<a class="code hl_function" href="classmlx_1_1core_1_1_primitive.html#a46e6257397a662528f9f831842ac456a">stream</a>, fallback), scale_(scale) {}</div>
|
||||
<div class="line"><a id="l00209" name="l00209"></a><span class="lineno"> 209</span> <span class="keyword">const</span> <span class="keywordtype">float</span> scale,</div>
|
||||
<div class="line"><a id="l00210" name="l00210"></a><span class="lineno"> 210</span> <span class="keyword">const</span> <span class="keywordtype">bool</span> <a class="code hl_variable" href="steel__attention_8h.html#abfa50278ba59a90e0acb7e5d94500741">do_causal</a>)</div>
|
||||
<div class="line"><a id="l00211" name="l00211"></a><span class="lineno"> 211</span> : <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_custom.html#a4186fea23f7156c38960426821fca313">Custom</a>(<a class="code hl_function" href="classmlx_1_1core_1_1_primitive.html#a46e6257397a662528f9f831842ac456a">stream</a>, fallback), scale_(scale), do_causal_(<a class="code hl_variable" href="steel__attention_8h.html#abfa50278ba59a90e0acb7e5d94500741">do_causal</a>) {}</div>
|
||||
</div>
|
||||
<div class="line"><a id="l00211" name="l00211"></a><span class="lineno"> 211</span> </div>
|
||||
<div class="foldopen" id="foldopen00212" data-start="{" data-end="}">
|
||||
<div class="line"><a id="l00212" name="l00212"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#ae20851e002f7fcb6d4f97817596f6328"> 212</a></span> <span class="keywordtype">void</span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#ae20851e002f7fcb6d4f97817596f6328">eval_cpu</a>(<span class="keyword">const</span> std::vector<array>& inputs, std::vector<array>& outputs)<span class="keyword"></span></div>
|
||||
<div class="line"><a id="l00213" name="l00213"></a><span class="lineno"> 213</span><span class="keyword"> override </span>{</div>
|
||||
<div class="line"><a id="l00214" name="l00214"></a><span class="lineno"> 214</span> <span class="keywordflow">throw</span> std::runtime_error(<span class="stringliteral">"NYI"</span>);</div>
|
||||
<div class="line"><a id="l00215" name="l00215"></a><span class="lineno"> 215</span> }</div>
|
||||
<div class="line"><a id="l00212" name="l00212"></a><span class="lineno"> 212</span> </div>
|
||||
<div class="foldopen" id="foldopen00213" data-start="{" data-end="}">
|
||||
<div class="line"><a id="l00213" name="l00213"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#ae20851e002f7fcb6d4f97817596f6328"> 213</a></span> <span class="keywordtype">void</span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#ae20851e002f7fcb6d4f97817596f6328">eval_cpu</a>(<span class="keyword">const</span> std::vector<array>& inputs, std::vector<array>& outputs)<span class="keyword"></span></div>
|
||||
<div class="line"><a id="l00214" name="l00214"></a><span class="lineno"> 214</span><span class="keyword"> override </span>{</div>
|
||||
<div class="line"><a id="l00215" name="l00215"></a><span class="lineno"> 215</span> <span class="keywordflow">throw</span> std::runtime_error(<span class="stringliteral">"NYI"</span>);</div>
|
||||
<div class="line"><a id="l00216" name="l00216"></a><span class="lineno"> 216</span> }</div>
|
||||
</div>
|
||||
<div class="line"><a id="l00216" name="l00216"></a><span class="lineno"> 216</span> </div>
|
||||
<div class="foldopen" id="foldopen00217" data-start="{" data-end="}">
|
||||
<div class="line"><a id="l00217" name="l00217"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#a505f38ba93a3499895f5312e0112e73d"> 217</a></span> <span class="keywordtype">void</span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#a505f38ba93a3499895f5312e0112e73d">eval_gpu</a>(<span class="keyword">const</span> std::vector<array>& inputs, std::vector<array>& outputs)<span class="keyword"></span></div>
|
||||
<div class="line"><a id="l00218" name="l00218"></a><span class="lineno"> 218</span><span class="keyword"> override </span>{</div>
|
||||
<div class="line"><a id="l00219" name="l00219"></a><span class="lineno"> 219</span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#a505f38ba93a3499895f5312e0112e73d">eval_gpu</a>(inputs, outputs[0]);</div>
|
||||
<div class="line"><a id="l00220" name="l00220"></a><span class="lineno"> 220</span> }</div>
|
||||
<div class="line"><a id="l00217" name="l00217"></a><span class="lineno"> 217</span> </div>
|
||||
<div class="foldopen" id="foldopen00218" data-start="{" data-end="}">
|
||||
<div class="line"><a id="l00218" name="l00218"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#a505f38ba93a3499895f5312e0112e73d"> 218</a></span> <span class="keywordtype">void</span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#a505f38ba93a3499895f5312e0112e73d">eval_gpu</a>(<span class="keyword">const</span> std::vector<array>& inputs, std::vector<array>& outputs)<span class="keyword"></span></div>
|
||||
<div class="line"><a id="l00219" name="l00219"></a><span class="lineno"> 219</span><span class="keyword"> override </span>{</div>
|
||||
<div class="line"><a id="l00220" name="l00220"></a><span class="lineno"> 220</span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#a505f38ba93a3499895f5312e0112e73d">eval_gpu</a>(inputs, outputs[0]);</div>
|
||||
<div class="line"><a id="l00221" name="l00221"></a><span class="lineno"> 221</span> }</div>
|
||||
</div>
|
||||
<div class="line"><a id="l00221" name="l00221"></a><span class="lineno"> 221</span> </div>
|
||||
<div class="line"><a id="l00222" name="l00222"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#ad51666e69f670e286293aff96eb435a9"> 222</a></span> <span class="keywordtype">void</span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#ad51666e69f670e286293aff96eb435a9">eval_gpu</a>(<span class="keyword">const</span> std::vector<array>& inputs, <a class="code hl_class" href="classmlx_1_1core_1_1array.html">array</a>& out);</div>
|
||||
<div class="line"><a id="l00223" name="l00223"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#af08b1294f3f93505a96fdfa85b1edd62"> 223</a></span> <span class="keywordtype">bool</span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#af08b1294f3f93505a96fdfa85b1edd62">is_equivalent</a>(<span class="keyword">const</span> <a class="code hl_function" href="classmlx_1_1core_1_1_primitive.html#afc69f22ee1f6e8a9ecc2c3a8f43b8fdb">Primitive</a>& other) <span class="keyword">const override</span>;</div>
|
||||
<div class="line"><a id="l00224" name="l00224"></a><span class="lineno"> 224</span> </div>
|
||||
<div class="line"><a id="l00225" name="l00225"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#a6cc2092fa5b8e7585921b8e0f3ec3db7"> 225</a></span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#a6cc2092fa5b8e7585921b8e0f3ec3db7">DEFINE_PRINT</a>(<a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#ab3f78d30e5bb3e76cfe701f2358e4748">ScaledDotProductAttention</a>);</div>
|
||||
<div class="foldopen" id="foldopen00226" data-start="{" data-end="}">
|
||||
<div class="line"><a id="l00226" name="l00226"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#a64d2ce4b46b529a6a9ef068947bc623e"> 226</a></span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#a64d2ce4b46b529a6a9ef068947bc623e">DEFINE_INPUT_OUTPUT_SHAPE</a>()</div>
|
||||
<div class="line"><a id="l00227" name="l00227"></a><span class="lineno"> 227</span> auto state()<span class="keyword"> const </span>{</div>
|
||||
<div class="line"><a id="l00228" name="l00228"></a><span class="lineno"> 228</span> <span class="keywordflow">return</span> std::make_pair(<span class="keyword">nullptr</span>, scale_);</div>
|
||||
<div class="line"><a id="l00229" name="l00229"></a><span class="lineno"> 229</span> }</div>
|
||||
<div class="line"><a id="l00222" name="l00222"></a><span class="lineno"> 222</span> </div>
|
||||
<div class="line"><a id="l00223" name="l00223"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#ad51666e69f670e286293aff96eb435a9"> 223</a></span> <span class="keywordtype">void</span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#ad51666e69f670e286293aff96eb435a9">eval_gpu</a>(<span class="keyword">const</span> std::vector<array>& inputs, <a class="code hl_class" href="classmlx_1_1core_1_1array.html">array</a>& out);</div>
|
||||
<div class="line"><a id="l00224" name="l00224"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#af08b1294f3f93505a96fdfa85b1edd62"> 224</a></span> <span class="keywordtype">bool</span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#af08b1294f3f93505a96fdfa85b1edd62">is_equivalent</a>(<span class="keyword">const</span> <a class="code hl_function" href="classmlx_1_1core_1_1_primitive.html#afc69f22ee1f6e8a9ecc2c3a8f43b8fdb">Primitive</a>& other) <span class="keyword">const override</span>;</div>
|
||||
<div class="line"><a id="l00225" name="l00225"></a><span class="lineno"> 225</span> </div>
|
||||
<div class="line"><a id="l00226" name="l00226"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#a6cc2092fa5b8e7585921b8e0f3ec3db7"> 226</a></span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#a6cc2092fa5b8e7585921b8e0f3ec3db7">DEFINE_PRINT</a>(<a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#a09c99b460cca606b2ebb22f90b3d13a2">ScaledDotProductAttention</a>);</div>
|
||||
<div class="foldopen" id="foldopen00227" data-start="{" data-end="}">
|
||||
<div class="line"><a id="l00227" name="l00227"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#a64d2ce4b46b529a6a9ef068947bc623e"> 227</a></span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#a64d2ce4b46b529a6a9ef068947bc623e">DEFINE_INPUT_OUTPUT_SHAPE</a>()</div>
|
||||
<div class="line"><a id="l00228" name="l00228"></a><span class="lineno"> 228</span> auto state()<span class="keyword"> const </span>{</div>
|
||||
<div class="line"><a id="l00229" name="l00229"></a><span class="lineno"> 229</span> <span class="keywordflow">return</span> std::make_tuple(<span class="keyword">nullptr</span>, scale_, do_causal_);</div>
|
||||
<div class="line"><a id="l00230" name="l00230"></a><span class="lineno"> 230</span> }</div>
|
||||
</div>
|
||||
<div class="line"><a id="l00230" name="l00230"></a><span class="lineno"> 230</span> </div>
|
||||
<div class="line"><a id="l00231" name="l00231"></a><span class="lineno"> 231</span> <span class="keyword">private</span>:</div>
|
||||
<div class="line"><a id="l00232" name="l00232"></a><span class="lineno"> 232</span> std::function<std::vector<array>(std::vector<array>)> fallback_;</div>
|
||||
<div class="line"><a id="l00233" name="l00233"></a><span class="lineno"> 233</span> <span class="keywordtype">float</span> scale_;</div>
|
||||
<div class="line"><a id="l00234" name="l00234"></a><span class="lineno"> 234</span>};</div>
|
||||
<div class="line"><a id="l00231" name="l00231"></a><span class="lineno"> 231</span> </div>
|
||||
<div class="line"><a id="l00232" name="l00232"></a><span class="lineno"> 232</span> <span class="keyword">private</span>:</div>
|
||||
<div class="line"><a id="l00233" name="l00233"></a><span class="lineno"> 233</span> std::function<std::vector<array>(std::vector<array>)> fallback_;</div>
|
||||
<div class="line"><a id="l00234" name="l00234"></a><span class="lineno"> 234</span> <span class="keywordtype">float</span> scale_;</div>
|
||||
<div class="line"><a id="l00235" name="l00235"></a><span class="lineno"> 235</span> <span class="keywordtype">bool</span> do_causal_;</div>
|
||||
<div class="line"><a id="l00236" name="l00236"></a><span class="lineno"> 236</span>};</div>
|
||||
</div>
|
||||
<div class="line"><a id="l00235" name="l00235"></a><span class="lineno"> 235</span> </div>
|
||||
<div class="foldopen" id="foldopen00236" data-start="{" data-end="};">
|
||||
<div class="line"><a id="l00236" name="l00236"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html"> 236</a></span><span class="keyword">class </span><a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#a84d5fa9e8c3de407fbcc5f38d2ed1473">AffineQuantize</a> : <span class="keyword">public</span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_custom.html#a4186fea23f7156c38960426821fca313">Custom</a> {</div>
|
||||
<div class="line"><a id="l00237" name="l00237"></a><span class="lineno"> 237</span> <span class="keyword">public</span>:</div>
|
||||
<div class="foldopen" id="foldopen00238" data-start="{" data-end="}">
|
||||
<div class="line"><a id="l00238" name="l00238"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#a84d5fa9e8c3de407fbcc5f38d2ed1473"> 238</a></span> <span class="keyword">explicit</span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#a84d5fa9e8c3de407fbcc5f38d2ed1473">AffineQuantize</a>(</div>
|
||||
<div class="line"><a id="l00239" name="l00239"></a><span class="lineno"> 239</span> <a class="code hl_struct" href="structmlx_1_1core_1_1_stream.html">Stream</a> <a class="code hl_function" href="classmlx_1_1core_1_1_primitive.html#a46e6257397a662528f9f831842ac456a">stream</a>,</div>
|
||||
<div class="line"><a id="l00240" name="l00240"></a><span class="lineno"> 240</span> std::function<std::vector<array>(std::vector<array>)> fallback,</div>
|
||||
<div class="line"><a id="l00241" name="l00241"></a><span class="lineno"> 241</span> <span class="keywordtype">int</span> group_size,</div>
|
||||
<div class="line"><a id="l00242" name="l00242"></a><span class="lineno"> 242</span> <span class="keywordtype">int</span> bits,</div>
|
||||
<div class="line"><a id="l00243" name="l00243"></a><span class="lineno"> 243</span> <span class="keywordtype">bool</span> <a class="code hl_function" href="group__ops.html#gabff758a5c1ce32ad7e8b78aba0164077">dequantize</a>)</div>
|
||||
<div class="line"><a id="l00244" name="l00244"></a><span class="lineno"> 244</span> : <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_custom.html#a4186fea23f7156c38960426821fca313">Custom</a>(<a class="code hl_function" href="classmlx_1_1core_1_1_primitive.html#a46e6257397a662528f9f831842ac456a">stream</a>, fallback),</div>
|
||||
<div class="line"><a id="l00245" name="l00245"></a><span class="lineno"> 245</span> group_size_(group_size),</div>
|
||||
<div class="line"><a id="l00246" name="l00246"></a><span class="lineno"> 246</span> bits_(bits),</div>
|
||||
<div class="line"><a id="l00247" name="l00247"></a><span class="lineno"> 247</span> dequantize_(<a class="code hl_function" href="group__ops.html#gabff758a5c1ce32ad7e8b78aba0164077">dequantize</a>) {}</div>
|
||||
<div class="line"><a id="l00237" name="l00237"></a><span class="lineno"> 237</span> </div>
|
||||
<div class="foldopen" id="foldopen00238" data-start="{" data-end="};">
|
||||
<div class="line"><a id="l00238" name="l00238"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html"> 238</a></span><span class="keyword">class </span><a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#a84d5fa9e8c3de407fbcc5f38d2ed1473">AffineQuantize</a> : <span class="keyword">public</span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_custom.html#a4186fea23f7156c38960426821fca313">Custom</a> {</div>
|
||||
<div class="line"><a id="l00239" name="l00239"></a><span class="lineno"> 239</span> <span class="keyword">public</span>:</div>
|
||||
<div class="foldopen" id="foldopen00240" data-start="{" data-end="}">
|
||||
<div class="line"><a id="l00240" name="l00240"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#a84d5fa9e8c3de407fbcc5f38d2ed1473"> 240</a></span> <span class="keyword">explicit</span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#a84d5fa9e8c3de407fbcc5f38d2ed1473">AffineQuantize</a>(</div>
|
||||
<div class="line"><a id="l00241" name="l00241"></a><span class="lineno"> 241</span> <a class="code hl_struct" href="structmlx_1_1core_1_1_stream.html">Stream</a> <a class="code hl_function" href="classmlx_1_1core_1_1_primitive.html#a46e6257397a662528f9f831842ac456a">stream</a>,</div>
|
||||
<div class="line"><a id="l00242" name="l00242"></a><span class="lineno"> 242</span> std::function<std::vector<array>(std::vector<array>)> fallback,</div>
|
||||
<div class="line"><a id="l00243" name="l00243"></a><span class="lineno"> 243</span> <span class="keywordtype">int</span> group_size,</div>
|
||||
<div class="line"><a id="l00244" name="l00244"></a><span class="lineno"> 244</span> <span class="keywordtype">int</span> bits,</div>
|
||||
<div class="line"><a id="l00245" name="l00245"></a><span class="lineno"> 245</span> <span class="keywordtype">bool</span> <a class="code hl_function" href="group__ops.html#gabff758a5c1ce32ad7e8b78aba0164077">dequantize</a>)</div>
|
||||
<div class="line"><a id="l00246" name="l00246"></a><span class="lineno"> 246</span> : <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_custom.html#a4186fea23f7156c38960426821fca313">Custom</a>(<a class="code hl_function" href="classmlx_1_1core_1_1_primitive.html#a46e6257397a662528f9f831842ac456a">stream</a>, fallback),</div>
|
||||
<div class="line"><a id="l00247" name="l00247"></a><span class="lineno"> 247</span> group_size_(group_size),</div>
|
||||
<div class="line"><a id="l00248" name="l00248"></a><span class="lineno"> 248</span> bits_(bits),</div>
|
||||
<div class="line"><a id="l00249" name="l00249"></a><span class="lineno"> 249</span> dequantize_(<a class="code hl_function" href="group__ops.html#gabff758a5c1ce32ad7e8b78aba0164077">dequantize</a>) {}</div>
|
||||
</div>
|
||||
<div class="line"><a id="l00248" name="l00248"></a><span class="lineno"> 248</span> </div>
|
||||
<div class="line"><a id="l00249" name="l00249"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#a3b5d628628d245b38911118d4a0ff9fd"> 249</a></span> <span class="keywordtype">void</span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#a3b5d628628d245b38911118d4a0ff9fd">eval_cpu</a>(<span class="keyword">const</span> std::vector<array>& inputs, std::vector<array>& outputs)</div>
|
||||
<div class="line"><a id="l00250" name="l00250"></a><span class="lineno"> 250</span> <span class="keyword">override</span>;</div>
|
||||
<div class="line"><a id="l00251" name="l00251"></a><span class="lineno"> 251</span> </div>
|
||||
<div class="line"><a id="l00252" name="l00252"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#a63812b2abaf26ad7e7fa4c9e82db1628"> 252</a></span> <span class="keywordtype">void</span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#a63812b2abaf26ad7e7fa4c9e82db1628">eval_gpu</a>(<span class="keyword">const</span> std::vector<array>& inputs, std::vector<array>& outputs)</div>
|
||||
<div class="line"><a id="l00253" name="l00253"></a><span class="lineno"> 253</span> <span class="keyword">override</span>;</div>
|
||||
<div class="line"><a id="l00254" name="l00254"></a><span class="lineno"> 254</span> </div>
|
||||
<div class="line"><a id="l00255" name="l00255"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#a4b8f1b1f633002c8ca6fa8f0ef4dd587"> 255</a></span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#a4b8f1b1f633002c8ca6fa8f0ef4dd587">DEFINE_PRINT</a>(<a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#a84d5fa9e8c3de407fbcc5f38d2ed1473">AffineQuantize</a>);</div>
|
||||
<div class="line"><a id="l00250" name="l00250"></a><span class="lineno"> 250</span> </div>
|
||||
<div class="line"><a id="l00251" name="l00251"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#a3b5d628628d245b38911118d4a0ff9fd"> 251</a></span> <span class="keywordtype">void</span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#a3b5d628628d245b38911118d4a0ff9fd">eval_cpu</a>(<span class="keyword">const</span> std::vector<array>& inputs, std::vector<array>& outputs)</div>
|
||||
<div class="line"><a id="l00252" name="l00252"></a><span class="lineno"> 252</span> <span class="keyword">override</span>;</div>
|
||||
<div class="line"><a id="l00253" name="l00253"></a><span class="lineno"> 253</span> </div>
|
||||
<div class="line"><a id="l00254" name="l00254"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#a63812b2abaf26ad7e7fa4c9e82db1628"> 254</a></span> <span class="keywordtype">void</span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#a63812b2abaf26ad7e7fa4c9e82db1628">eval_gpu</a>(<span class="keyword">const</span> std::vector<array>& inputs, std::vector<array>& outputs)</div>
|
||||
<div class="line"><a id="l00255" name="l00255"></a><span class="lineno"> 255</span> <span class="keyword">override</span>;</div>
|
||||
<div class="line"><a id="l00256" name="l00256"></a><span class="lineno"> 256</span> </div>
|
||||
<div class="line"><a id="l00257" name="l00257"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#a5936175e5923aec272d6f718785f57a1"> 257</a></span> <span class="keywordtype">bool</span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#a5936175e5923aec272d6f718785f57a1">is_equivalent</a>(<span class="keyword">const</span> <a class="code hl_function" href="classmlx_1_1core_1_1_primitive.html#afc69f22ee1f6e8a9ecc2c3a8f43b8fdb">Primitive</a>& other) <span class="keyword">const override</span>;</div>
|
||||
<div class="line"><a id="l00258" name="l00258"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#a50934862ccdb16a3dcce6626c5727080"> 258</a></span> std::vector<Shape> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#a50934862ccdb16a3dcce6626c5727080">output_shapes</a>(<span class="keyword">const</span> std::vector<array>& inputs) <span class="keyword">override</span>;</div>
|
||||
<div class="foldopen" id="foldopen00259" data-start="{" data-end="}">
|
||||
<div class="line"><a id="l00259" name="l00259"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#aa5a03284c6f5639d684dd34d86050cf9"> 259</a></span> <span class="keyword">auto</span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#aa5a03284c6f5639d684dd34d86050cf9">state</a>()<span class="keyword"> const </span>{</div>
|
||||
<div class="line"><a id="l00260" name="l00260"></a><span class="lineno"> 260</span> <span class="keywordflow">return</span> std::make_tuple(<span class="keyword">nullptr</span>, group_size_, bits_, dequantize_);</div>
|
||||
<div class="line"><a id="l00261" name="l00261"></a><span class="lineno"> 261</span> }</div>
|
||||
<div class="line"><a id="l00257" name="l00257"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#a4b8f1b1f633002c8ca6fa8f0ef4dd587"> 257</a></span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#a4b8f1b1f633002c8ca6fa8f0ef4dd587">DEFINE_PRINT</a>(<a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#a84d5fa9e8c3de407fbcc5f38d2ed1473">AffineQuantize</a>);</div>
|
||||
<div class="line"><a id="l00258" name="l00258"></a><span class="lineno"> 258</span> </div>
|
||||
<div class="line"><a id="l00259" name="l00259"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#a5936175e5923aec272d6f718785f57a1"> 259</a></span> <span class="keywordtype">bool</span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#a5936175e5923aec272d6f718785f57a1">is_equivalent</a>(<span class="keyword">const</span> <a class="code hl_function" href="classmlx_1_1core_1_1_primitive.html#afc69f22ee1f6e8a9ecc2c3a8f43b8fdb">Primitive</a>& other) <span class="keyword">const override</span>;</div>
|
||||
<div class="line"><a id="l00260" name="l00260"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#a50934862ccdb16a3dcce6626c5727080"> 260</a></span> std::vector<Shape> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#a50934862ccdb16a3dcce6626c5727080">output_shapes</a>(<span class="keyword">const</span> std::vector<array>& inputs) <span class="keyword">override</span>;</div>
|
||||
<div class="foldopen" id="foldopen00261" data-start="{" data-end="}">
|
||||
<div class="line"><a id="l00261" name="l00261"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#aa5a03284c6f5639d684dd34d86050cf9"> 261</a></span> <span class="keyword">auto</span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#aa5a03284c6f5639d684dd34d86050cf9">state</a>()<span class="keyword"> const </span>{</div>
|
||||
<div class="line"><a id="l00262" name="l00262"></a><span class="lineno"> 262</span> <span class="keywordflow">return</span> std::make_tuple(<span class="keyword">nullptr</span>, group_size_, bits_, dequantize_);</div>
|
||||
<div class="line"><a id="l00263" name="l00263"></a><span class="lineno"> 263</span> }</div>
|
||||
</div>
|
||||
<div class="line"><a id="l00262" name="l00262"></a><span class="lineno"> 262</span> </div>
|
||||
<div class="line"><a id="l00263" name="l00263"></a><span class="lineno"> 263</span> <span class="keyword">private</span>:</div>
|
||||
<div class="line"><a id="l00264" name="l00264"></a><span class="lineno"> 264</span> std::function<std::vector<array>(std::vector<array>)> fallback_;</div>
|
||||
<div class="line"><a id="l00265" name="l00265"></a><span class="lineno"> 265</span> <span class="keywordtype">int</span> group_size_;</div>
|
||||
<div class="line"><a id="l00266" name="l00266"></a><span class="lineno"> 266</span> <span class="keywordtype">int</span> bits_;</div>
|
||||
<div class="line"><a id="l00267" name="l00267"></a><span class="lineno"> 267</span> <span class="keywordtype">bool</span> dequantize_;</div>
|
||||
<div class="line"><a id="l00268" name="l00268"></a><span class="lineno"> 268</span>};</div>
|
||||
<div class="line"><a id="l00264" name="l00264"></a><span class="lineno"> 264</span> </div>
|
||||
<div class="line"><a id="l00265" name="l00265"></a><span class="lineno"> 265</span> <span class="keyword">private</span>:</div>
|
||||
<div class="line"><a id="l00266" name="l00266"></a><span class="lineno"> 266</span> std::function<std::vector<array>(std::vector<array>)> fallback_;</div>
|
||||
<div class="line"><a id="l00267" name="l00267"></a><span class="lineno"> 267</span> <span class="keywordtype">int</span> group_size_;</div>
|
||||
<div class="line"><a id="l00268" name="l00268"></a><span class="lineno"> 268</span> <span class="keywordtype">int</span> bits_;</div>
|
||||
<div class="line"><a id="l00269" name="l00269"></a><span class="lineno"> 269</span> <span class="keywordtype">bool</span> dequantize_;</div>
|
||||
<div class="line"><a id="l00270" name="l00270"></a><span class="lineno"> 270</span>};</div>
|
||||
</div>
|
||||
<div class="line"><a id="l00269" name="l00269"></a><span class="lineno"> 269</span> </div>
|
||||
<div class="foldopen" id="foldopen00270" data-start="{" data-end="};">
|
||||
<div class="line"><a id="l00270" name="l00270"></a><span class="lineno"><a class="line" href="structmlx_1_1core_1_1fast_1_1_custom_kernel_shape_info.html"> 270</a></span><span class="keyword">struct </span><a class="code hl_struct" href="structmlx_1_1core_1_1fast_1_1_custom_kernel_shape_info.html">CustomKernelShapeInfo</a> {</div>
|
||||
<div class="line"><a id="l00271" name="l00271"></a><span class="lineno"><a class="line" href="structmlx_1_1core_1_1fast_1_1_custom_kernel_shape_info.html#a63db720fe0c2abc4b71e22a58a015f8a"> 271</a></span> <span class="keywordtype">bool</span> <a class="code hl_variable" href="structmlx_1_1core_1_1fast_1_1_custom_kernel_shape_info.html#a63db720fe0c2abc4b71e22a58a015f8a">shape</a> = <span class="keyword">false</span>;</div>
|
||||
<div class="line"><a id="l00272" name="l00272"></a><span class="lineno"><a class="line" href="structmlx_1_1core_1_1fast_1_1_custom_kernel_shape_info.html#a63954de7da62942ec69afcaaa19d46f2"> 272</a></span> <span class="keywordtype">bool</span> <a class="code hl_variable" href="structmlx_1_1core_1_1fast_1_1_custom_kernel_shape_info.html#a63954de7da62942ec69afcaaa19d46f2">strides</a> = <span class="keyword">false</span>;</div>
|
||||
<div class="line"><a id="l00273" name="l00273"></a><span class="lineno"><a class="line" href="structmlx_1_1core_1_1fast_1_1_custom_kernel_shape_info.html#ae605df33f449872e3da9777d97008051"> 273</a></span> <span class="keywordtype">bool</span> <a class="code hl_variable" href="structmlx_1_1core_1_1fast_1_1_custom_kernel_shape_info.html#ae605df33f449872e3da9777d97008051">ndim</a> = <span class="keyword">false</span>;</div>
|
||||
<div class="line"><a id="l00274" name="l00274"></a><span class="lineno"> 274</span>};</div>
|
||||
<div class="line"><a id="l00271" name="l00271"></a><span class="lineno"> 271</span> </div>
|
||||
<div class="foldopen" id="foldopen00272" data-start="{" data-end="};">
|
||||
<div class="line"><a id="l00272" name="l00272"></a><span class="lineno"><a class="line" href="structmlx_1_1core_1_1fast_1_1_custom_kernel_shape_info.html"> 272</a></span><span class="keyword">struct </span><a class="code hl_struct" href="structmlx_1_1core_1_1fast_1_1_custom_kernel_shape_info.html">CustomKernelShapeInfo</a> {</div>
|
||||
<div class="line"><a id="l00273" name="l00273"></a><span class="lineno"><a class="line" href="structmlx_1_1core_1_1fast_1_1_custom_kernel_shape_info.html#a63db720fe0c2abc4b71e22a58a015f8a"> 273</a></span> <span class="keywordtype">bool</span> <a class="code hl_variable" href="structmlx_1_1core_1_1fast_1_1_custom_kernel_shape_info.html#a63db720fe0c2abc4b71e22a58a015f8a">shape</a> = <span class="keyword">false</span>;</div>
|
||||
<div class="line"><a id="l00274" name="l00274"></a><span class="lineno"><a class="line" href="structmlx_1_1core_1_1fast_1_1_custom_kernel_shape_info.html#a63954de7da62942ec69afcaaa19d46f2"> 274</a></span> <span class="keywordtype">bool</span> <a class="code hl_variable" href="structmlx_1_1core_1_1fast_1_1_custom_kernel_shape_info.html#a63954de7da62942ec69afcaaa19d46f2">strides</a> = <span class="keyword">false</span>;</div>
|
||||
<div class="line"><a id="l00275" name="l00275"></a><span class="lineno"><a class="line" href="structmlx_1_1core_1_1fast_1_1_custom_kernel_shape_info.html#ae605df33f449872e3da9777d97008051"> 275</a></span> <span class="keywordtype">bool</span> <a class="code hl_variable" href="structmlx_1_1core_1_1fast_1_1_custom_kernel_shape_info.html#ae605df33f449872e3da9777d97008051">ndim</a> = <span class="keyword">false</span>;</div>
|
||||
<div class="line"><a id="l00276" name="l00276"></a><span class="lineno"> 276</span>};</div>
|
||||
</div>
|
||||
<div class="line"><a id="l00275" name="l00275"></a><span class="lineno"> 275</span> </div>
|
||||
<div class="foldopen" id="foldopen00276" data-start="{" data-end="};">
|
||||
<div class="line"><a id="l00276" name="l00276"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_custom_kernel.html"> 276</a></span><span class="keyword">class </span><a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_custom_kernel.html#a954893e07f0d36715b4e1e414b6f2153">CustomKernel</a> : <span class="keyword">public</span> <a class="code hl_function" href="classmlx_1_1core_1_1_primitive.html#afc69f22ee1f6e8a9ecc2c3a8f43b8fdb">Primitive</a> {</div>
|
||||
<div class="line"><a id="l00277" name="l00277"></a><span class="lineno"> 277</span> <span class="keyword">public</span>:</div>
|
||||
<div class="foldopen" id="foldopen00278" data-start="{" data-end="}">
|
||||
<div class="line"><a id="l00278" name="l00278"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_custom_kernel.html#a954893e07f0d36715b4e1e414b6f2153"> 278</a></span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_custom_kernel.html#a954893e07f0d36715b4e1e414b6f2153">CustomKernel</a>(</div>
|
||||
<div class="line"><a id="l00279" name="l00279"></a><span class="lineno"> 279</span> <a class="code hl_struct" href="structmlx_1_1core_1_1_stream.html">Stream</a> <a class="code hl_function" href="classmlx_1_1core_1_1_primitive.html#a46e6257397a662528f9f831842ac456a">stream</a>,</div>
|
||||
<div class="line"><a id="l00280" name="l00280"></a><span class="lineno"> 280</span> std::string name,</div>
|
||||
<div class="line"><a id="l00281" name="l00281"></a><span class="lineno"> 281</span> std::string source,</div>
|
||||
<div class="line"><a id="l00282" name="l00282"></a><span class="lineno"> 282</span> std::tuple<int, int, int> grid,</div>
|
||||
<div class="line"><a id="l00283" name="l00283"></a><span class="lineno"> 283</span> std::tuple<int, int, int> threadgroup,</div>
|
||||
<div class="line"><a id="l00284" name="l00284"></a><span class="lineno"> 284</span> std::vector<CustomKernelShapeInfo> shape_infos,</div>
|
||||
<div class="line"><a id="l00285" name="l00285"></a><span class="lineno"> 285</span> <span class="keywordtype">bool</span> ensure_row_contiguous,</div>
|
||||
<div class="line"><a id="l00286" name="l00286"></a><span class="lineno"> 286</span> std::optional<float> init_value)</div>
|
||||
<div class="line"><a id="l00287" name="l00287"></a><span class="lineno"> 287</span> : <a class="code hl_function" href="classmlx_1_1core_1_1_primitive.html#afc69f22ee1f6e8a9ecc2c3a8f43b8fdb">Primitive</a>(<a class="code hl_function" href="classmlx_1_1core_1_1_primitive.html#a46e6257397a662528f9f831842ac456a">stream</a>),</div>
|
||||
<div class="line"><a id="l00288" name="l00288"></a><span class="lineno"> 288</span> source_(<a class="code hl_function" href="group__ops.html#ga2a466024f8061febc0a64be557644cb0">std</a>::move(source)),</div>
|
||||
<div class="line"><a id="l00289" name="l00289"></a><span class="lineno"> 289</span> name_(<a class="code hl_function" href="group__ops.html#ga2a466024f8061febc0a64be557644cb0">std</a>::move(name)),</div>
|
||||
<div class="line"><a id="l00290" name="l00290"></a><span class="lineno"> 290</span> grid_(grid),</div>
|
||||
<div class="line"><a id="l00291" name="l00291"></a><span class="lineno"> 291</span> threadgroup_(threadgroup),</div>
|
||||
<div class="line"><a id="l00292" name="l00292"></a><span class="lineno"> 292</span> shape_infos_(<a class="code hl_function" href="group__ops.html#ga2a466024f8061febc0a64be557644cb0">std</a>::move(shape_infos)),</div>
|
||||
<div class="line"><a id="l00293" name="l00293"></a><span class="lineno"> 293</span> ensure_row_contiguous_(ensure_row_contiguous),</div>
|
||||
<div class="line"><a id="l00294" name="l00294"></a><span class="lineno"> 294</span> init_value_(init_value) {}</div>
|
||||
<div class="line"><a id="l00277" name="l00277"></a><span class="lineno"> 277</span> </div>
|
||||
<div class="foldopen" id="foldopen00278" data-start="{" data-end="};">
|
||||
<div class="line"><a id="l00278" name="l00278"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_custom_kernel.html"> 278</a></span><span class="keyword">class </span><a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_custom_kernel.html#a954893e07f0d36715b4e1e414b6f2153">CustomKernel</a> : <span class="keyword">public</span> <a class="code hl_function" href="classmlx_1_1core_1_1_primitive.html#afc69f22ee1f6e8a9ecc2c3a8f43b8fdb">Primitive</a> {</div>
|
||||
<div class="line"><a id="l00279" name="l00279"></a><span class="lineno"> 279</span> <span class="keyword">public</span>:</div>
|
||||
<div class="foldopen" id="foldopen00280" data-start="{" data-end="}">
|
||||
<div class="line"><a id="l00280" name="l00280"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_custom_kernel.html#a954893e07f0d36715b4e1e414b6f2153"> 280</a></span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_custom_kernel.html#a954893e07f0d36715b4e1e414b6f2153">CustomKernel</a>(</div>
|
||||
<div class="line"><a id="l00281" name="l00281"></a><span class="lineno"> 281</span> <a class="code hl_struct" href="structmlx_1_1core_1_1_stream.html">Stream</a> <a class="code hl_function" href="classmlx_1_1core_1_1_primitive.html#a46e6257397a662528f9f831842ac456a">stream</a>,</div>
|
||||
<div class="line"><a id="l00282" name="l00282"></a><span class="lineno"> 282</span> std::string name,</div>
|
||||
<div class="line"><a id="l00283" name="l00283"></a><span class="lineno"> 283</span> std::string source,</div>
|
||||
<div class="line"><a id="l00284" name="l00284"></a><span class="lineno"> 284</span> std::tuple<int, int, int> grid,</div>
|
||||
<div class="line"><a id="l00285" name="l00285"></a><span class="lineno"> 285</span> std::tuple<int, int, int> threadgroup,</div>
|
||||
<div class="line"><a id="l00286" name="l00286"></a><span class="lineno"> 286</span> std::vector<CustomKernelShapeInfo> shape_infos,</div>
|
||||
<div class="line"><a id="l00287" name="l00287"></a><span class="lineno"> 287</span> <span class="keywordtype">bool</span> ensure_row_contiguous,</div>
|
||||
<div class="line"><a id="l00288" name="l00288"></a><span class="lineno"> 288</span> std::optional<float> init_value)</div>
|
||||
<div class="line"><a id="l00289" name="l00289"></a><span class="lineno"> 289</span> : <a class="code hl_function" href="classmlx_1_1core_1_1_primitive.html#afc69f22ee1f6e8a9ecc2c3a8f43b8fdb">Primitive</a>(<a class="code hl_function" href="classmlx_1_1core_1_1_primitive.html#a46e6257397a662528f9f831842ac456a">stream</a>),</div>
|
||||
<div class="line"><a id="l00290" name="l00290"></a><span class="lineno"> 290</span> source_(<a class="code hl_function" href="group__ops.html#ga2a466024f8061febc0a64be557644cb0">std</a>::move(source)),</div>
|
||||
<div class="line"><a id="l00291" name="l00291"></a><span class="lineno"> 291</span> name_(<a class="code hl_function" href="group__ops.html#ga2a466024f8061febc0a64be557644cb0">std</a>::move(name)),</div>
|
||||
<div class="line"><a id="l00292" name="l00292"></a><span class="lineno"> 292</span> grid_(grid),</div>
|
||||
<div class="line"><a id="l00293" name="l00293"></a><span class="lineno"> 293</span> threadgroup_(threadgroup),</div>
|
||||
<div class="line"><a id="l00294" name="l00294"></a><span class="lineno"> 294</span> shape_infos_(<a class="code hl_function" href="group__ops.html#ga2a466024f8061febc0a64be557644cb0">std</a>::move(shape_infos)),</div>
|
||||
<div class="line"><a id="l00295" name="l00295"></a><span class="lineno"> 295</span> ensure_row_contiguous_(ensure_row_contiguous),</div>
|
||||
<div class="line"><a id="l00296" name="l00296"></a><span class="lineno"> 296</span> init_value_(init_value) {}</div>
|
||||
</div>
|
||||
<div class="line"><a id="l00295" name="l00295"></a><span class="lineno"> 295</span> </div>
|
||||
<div class="foldopen" id="foldopen00296" data-start="{" data-end="}">
|
||||
<div class="line"><a id="l00296" name="l00296"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_custom_kernel.html#a4ad1b7a9919753c759093f3e21a15bad"> 296</a></span> <span class="keywordtype">void</span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_custom_kernel.html#a4ad1b7a9919753c759093f3e21a15bad">eval_cpu</a>(<span class="keyword">const</span> std::vector<array>& inputs, std::vector<array>& outputs)<span class="keyword"></span></div>
|
||||
<div class="line"><a id="l00297" name="l00297"></a><span class="lineno"> 297</span><span class="keyword"> override </span>{</div>
|
||||
<div class="line"><a id="l00298" name="l00298"></a><span class="lineno"> 298</span> <span class="keywordflow">throw</span> std::runtime_error(<span class="stringliteral">"Custom Metal kernels only run on GPU."</span>);</div>
|
||||
<div class="line"><a id="l00299" name="l00299"></a><span class="lineno"> 299</span> }</div>
|
||||
<div class="line"><a id="l00297" name="l00297"></a><span class="lineno"> 297</span> </div>
|
||||
<div class="foldopen" id="foldopen00298" data-start="{" data-end="}">
|
||||
<div class="line"><a id="l00298" name="l00298"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_custom_kernel.html#a4ad1b7a9919753c759093f3e21a15bad"> 298</a></span> <span class="keywordtype">void</span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_custom_kernel.html#a4ad1b7a9919753c759093f3e21a15bad">eval_cpu</a>(<span class="keyword">const</span> std::vector<array>& inputs, std::vector<array>& outputs)<span class="keyword"></span></div>
|
||||
<div class="line"><a id="l00299" name="l00299"></a><span class="lineno"> 299</span><span class="keyword"> override </span>{</div>
|
||||
<div class="line"><a id="l00300" name="l00300"></a><span class="lineno"> 300</span> <span class="keywordflow">throw</span> std::runtime_error(<span class="stringliteral">"Custom Metal kernels only run on GPU."</span>);</div>
|
||||
<div class="line"><a id="l00301" name="l00301"></a><span class="lineno"> 301</span> }</div>
|
||||
</div>
|
||||
<div class="line"><a id="l00300" name="l00300"></a><span class="lineno"> 300</span> </div>
|
||||
<div class="line"><a id="l00301" name="l00301"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_custom_kernel.html#a2ed2a16b23053f8195068386a99fd6db"> 301</a></span> <span class="keywordtype">void</span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_custom_kernel.html#a2ed2a16b23053f8195068386a99fd6db">eval_gpu</a>(<span class="keyword">const</span> std::vector<array>& inputs, std::vector<array>& outputs)</div>
|
||||
<div class="line"><a id="l00302" name="l00302"></a><span class="lineno"> 302</span> <span class="keyword">override</span>;</div>
|
||||
<div class="line"><a id="l00303" name="l00303"></a><span class="lineno"> 303</span> </div>
|
||||
<div class="line"><a id="l00304" name="l00304"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_custom_kernel.html#a116ecf31c8672c94e5ea06c1d43e9534"> 304</a></span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_custom_kernel.html#a116ecf31c8672c94e5ea06c1d43e9534">DEFINE_PRINT</a>(<a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_custom_kernel.html#a954893e07f0d36715b4e1e414b6f2153">CustomKernel</a>);</div>
|
||||
<div class="line"><a id="l00302" name="l00302"></a><span class="lineno"> 302</span> </div>
|
||||
<div class="line"><a id="l00303" name="l00303"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_custom_kernel.html#a2ed2a16b23053f8195068386a99fd6db"> 303</a></span> <span class="keywordtype">void</span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_custom_kernel.html#a2ed2a16b23053f8195068386a99fd6db">eval_gpu</a>(<span class="keyword">const</span> std::vector<array>& inputs, std::vector<array>& outputs)</div>
|
||||
<div class="line"><a id="l00304" name="l00304"></a><span class="lineno"> 304</span> <span class="keyword">override</span>;</div>
|
||||
<div class="line"><a id="l00305" name="l00305"></a><span class="lineno"> 305</span> </div>
|
||||
<div class="line"><a id="l00306" name="l00306"></a><span class="lineno"> 306</span> <span class="keyword">private</span>:</div>
|
||||
<div class="line"><a id="l00307" name="l00307"></a><span class="lineno"> 307</span> std::string source_;</div>
|
||||
<div class="line"><a id="l00308" name="l00308"></a><span class="lineno"> 308</span> std::string name_;</div>
|
||||
<div class="line"><a id="l00309" name="l00309"></a><span class="lineno"> 309</span> std::tuple<int, int, int> grid_;</div>
|
||||
<div class="line"><a id="l00310" name="l00310"></a><span class="lineno"> 310</span> std::tuple<int, int, int> threadgroup_;</div>
|
||||
<div class="line"><a id="l00311" name="l00311"></a><span class="lineno"> 311</span> std::vector<CustomKernelShapeInfo> shape_infos_;</div>
|
||||
<div class="line"><a id="l00312" name="l00312"></a><span class="lineno"> 312</span> <span class="keywordtype">bool</span> ensure_row_contiguous_;</div>
|
||||
<div class="line"><a id="l00313" name="l00313"></a><span class="lineno"> 313</span> std::optional<float> init_value_;</div>
|
||||
<div class="line"><a id="l00314" name="l00314"></a><span class="lineno"> 314</span>};</div>
|
||||
<div class="line"><a id="l00306" name="l00306"></a><span class="lineno"><a class="line" href="classmlx_1_1core_1_1fast_1_1_custom_kernel.html#a116ecf31c8672c94e5ea06c1d43e9534"> 306</a></span> <a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_custom_kernel.html#a116ecf31c8672c94e5ea06c1d43e9534">DEFINE_PRINT</a>(<a class="code hl_function" href="classmlx_1_1core_1_1fast_1_1_custom_kernel.html#a954893e07f0d36715b4e1e414b6f2153">CustomKernel</a>);</div>
|
||||
<div class="line"><a id="l00307" name="l00307"></a><span class="lineno"> 307</span> </div>
|
||||
<div class="line"><a id="l00308" name="l00308"></a><span class="lineno"> 308</span> <span class="keyword">private</span>:</div>
|
||||
<div class="line"><a id="l00309" name="l00309"></a><span class="lineno"> 309</span> std::string source_;</div>
|
||||
<div class="line"><a id="l00310" name="l00310"></a><span class="lineno"> 310</span> std::string name_;</div>
|
||||
<div class="line"><a id="l00311" name="l00311"></a><span class="lineno"> 311</span> std::tuple<int, int, int> grid_;</div>
|
||||
<div class="line"><a id="l00312" name="l00312"></a><span class="lineno"> 312</span> std::tuple<int, int, int> threadgroup_;</div>
|
||||
<div class="line"><a id="l00313" name="l00313"></a><span class="lineno"> 313</span> std::vector<CustomKernelShapeInfo> shape_infos_;</div>
|
||||
<div class="line"><a id="l00314" name="l00314"></a><span class="lineno"> 314</span> <span class="keywordtype">bool</span> ensure_row_contiguous_;</div>
|
||||
<div class="line"><a id="l00315" name="l00315"></a><span class="lineno"> 315</span> std::optional<float> init_value_;</div>
|
||||
<div class="line"><a id="l00316" name="l00316"></a><span class="lineno"> 316</span>};</div>
|
||||
</div>
|
||||
<div class="line"><a id="l00315" name="l00315"></a><span class="lineno"> 315</span> </div>
|
||||
<div class="line"><a id="l00316" name="l00316"></a><span class="lineno"> 316</span>} <span class="comment">// namespace mlx::core::fast</span></div>
|
||||
<div class="line"><a id="l00317" name="l00317"></a><span class="lineno"> 317</span> </div>
|
||||
<div class="line"><a id="l00318" name="l00318"></a><span class="lineno"> 318</span>} <span class="comment">// namespace mlx::core::fast</span></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1_primitive_html_a46e6257397a662528f9f831842ac456a"><div class="ttname"><a href="classmlx_1_1core_1_1_primitive.html#a46e6257397a662528f9f831842ac456a">mlx::core::Primitive::stream</a></div><div class="ttdeci">const Stream & stream()</div><div class="ttdoc">The stream the primitive will run on.</div><div class="ttdef"><b>Definition</b> primitives.h:58</div></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1_primitive_html_a6140a502af4c2bbbc776ab26e9afebcd"><div class="ttname"><a href="classmlx_1_1core_1_1_primitive.html#a6140a502af4c2bbbc776ab26e9afebcd">mlx::core::Primitive::is_equivalent</a></div><div class="ttdeci">virtual bool is_equivalent(const Primitive &other) const</div><div class="ttdoc">Equivalence check defaults to false unless overridden by the primitive.</div><div class="ttdef"><b>Definition</b> primitives.h:107</div></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1_primitive_html_afc69f22ee1f6e8a9ecc2c3a8f43b8fdb"><div class="ttname"><a href="classmlx_1_1core_1_1_primitive.html#afc69f22ee1f6e8a9ecc2c3a8f43b8fdb">mlx::core::Primitive::Primitive</a></div><div class="ttdeci">Primitive(Stream stream)</div><div class="ttdef"><b>Definition</b> primitives.h:50</div></div>
|
||||
@@ -498,16 +500,16 @@ $(function(){initNavTree('fast__primitives_8h_source.html',''); initResizable(tr
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1fast_1_1_affine_quantize_html_a50934862ccdb16a3dcce6626c5727080"><div class="ttname"><a href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#a50934862ccdb16a3dcce6626c5727080">mlx::core::fast::AffineQuantize::output_shapes</a></div><div class="ttdeci">std::vector< Shape > output_shapes(const std::vector< array > &inputs) override</div><div class="ttdoc">Get the output shapes of the primitive.</div></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1fast_1_1_affine_quantize_html_a5936175e5923aec272d6f718785f57a1"><div class="ttname"><a href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#a5936175e5923aec272d6f718785f57a1">mlx::core::fast::AffineQuantize::is_equivalent</a></div><div class="ttdeci">bool is_equivalent(const Primitive &other) const override</div><div class="ttdoc">Equivalence check defaults to false unless overridden by the primitive.</div></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1fast_1_1_affine_quantize_html_a63812b2abaf26ad7e7fa4c9e82db1628"><div class="ttname"><a href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#a63812b2abaf26ad7e7fa4c9e82db1628">mlx::core::fast::AffineQuantize::eval_gpu</a></div><div class="ttdeci">void eval_gpu(const std::vector< array > &inputs, std::vector< array > &outputs) override</div></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1fast_1_1_affine_quantize_html_a84d5fa9e8c3de407fbcc5f38d2ed1473"><div class="ttname"><a href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#a84d5fa9e8c3de407fbcc5f38d2ed1473">mlx::core::fast::AffineQuantize::AffineQuantize</a></div><div class="ttdeci">AffineQuantize(Stream stream, std::function< std::vector< array >(std::vector< array >)> fallback, int group_size, int bits, bool dequantize)</div><div class="ttdef"><b>Definition</b> fast_primitives.h:238</div></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1fast_1_1_affine_quantize_html_aa5a03284c6f5639d684dd34d86050cf9"><div class="ttname"><a href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#aa5a03284c6f5639d684dd34d86050cf9">mlx::core::fast::AffineQuantize::state</a></div><div class="ttdeci">auto state() const</div><div class="ttdef"><b>Definition</b> fast_primitives.h:259</div></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1fast_1_1_affine_quantize_html_a84d5fa9e8c3de407fbcc5f38d2ed1473"><div class="ttname"><a href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#a84d5fa9e8c3de407fbcc5f38d2ed1473">mlx::core::fast::AffineQuantize::AffineQuantize</a></div><div class="ttdeci">AffineQuantize(Stream stream, std::function< std::vector< array >(std::vector< array >)> fallback, int group_size, int bits, bool dequantize)</div><div class="ttdef"><b>Definition</b> fast_primitives.h:240</div></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1fast_1_1_affine_quantize_html_aa5a03284c6f5639d684dd34d86050cf9"><div class="ttname"><a href="classmlx_1_1core_1_1fast_1_1_affine_quantize.html#aa5a03284c6f5639d684dd34d86050cf9">mlx::core::fast::AffineQuantize::state</a></div><div class="ttdeci">auto state() const</div><div class="ttdef"><b>Definition</b> fast_primitives.h:261</div></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1fast_1_1_custom_html_a4186fea23f7156c38960426821fca313"><div class="ttname"><a href="classmlx_1_1core_1_1fast_1_1_custom.html#a4186fea23f7156c38960426821fca313">mlx::core::fast::Custom::Custom</a></div><div class="ttdeci">Custom(Stream stream, std::function< std::vector< array >(std::vector< array >)> fallback)</div><div class="ttdef"><b>Definition</b> fast_primitives.h:14</div></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1fast_1_1_custom_html_a74be4bcd0382f7f6400bf73fd5569c91"><div class="ttname"><a href="classmlx_1_1core_1_1fast_1_1_custom.html#a74be4bcd0382f7f6400bf73fd5569c91">mlx::core::fast::Custom::vjp</a></div><div class="ttdeci">virtual std::vector< array > vjp(const std::vector< array > &primals, const std::vector< array > &cotangents, const std::vector< int > &argnums, const std::vector< array > &outputs) override</div><div class="ttdoc">The vector-Jacobian product.</div></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1fast_1_1_custom_html_a7f4c3a4c48c6807faa36fb31e39dad8d"><div class="ttname"><a href="classmlx_1_1core_1_1fast_1_1_custom.html#a7f4c3a4c48c6807faa36fb31e39dad8d">mlx::core::fast::Custom::vmap</a></div><div class="ttdeci">virtual std::pair< std::vector< array >, std::vector< int > > vmap(const std::vector< array > &inputs, const std::vector< int > &axes) override</div><div class="ttdoc">The primitive must know how to vectorize itself across the given axes.</div></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1fast_1_1_custom_html_ac77b28702654df8e7d882a49357a9584"><div class="ttname"><a href="classmlx_1_1core_1_1fast_1_1_custom.html#ac77b28702654df8e7d882a49357a9584">mlx::core::fast::Custom::jvp</a></div><div class="ttdeci">virtual std::vector< array > jvp(const std::vector< array > &primals, const std::vector< array > &tangents, const std::vector< int > &argnums) override</div><div class="ttdoc">The Jacobian-vector product.</div></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1fast_1_1_custom_kernel_html_a116ecf31c8672c94e5ea06c1d43e9534"><div class="ttname"><a href="classmlx_1_1core_1_1fast_1_1_custom_kernel.html#a116ecf31c8672c94e5ea06c1d43e9534">mlx::core::fast::CustomKernel::DEFINE_PRINT</a></div><div class="ttdeci">DEFINE_PRINT(CustomKernel)</div></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1fast_1_1_custom_kernel_html_a2ed2a16b23053f8195068386a99fd6db"><div class="ttname"><a href="classmlx_1_1core_1_1fast_1_1_custom_kernel.html#a2ed2a16b23053f8195068386a99fd6db">mlx::core::fast::CustomKernel::eval_gpu</a></div><div class="ttdeci">void eval_gpu(const std::vector< array > &inputs, std::vector< array > &outputs) override</div></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1fast_1_1_custom_kernel_html_a4ad1b7a9919753c759093f3e21a15bad"><div class="ttname"><a href="classmlx_1_1core_1_1fast_1_1_custom_kernel.html#a4ad1b7a9919753c759093f3e21a15bad">mlx::core::fast::CustomKernel::eval_cpu</a></div><div class="ttdeci">void eval_cpu(const std::vector< array > &inputs, std::vector< array > &outputs) override</div><div class="ttdoc">A primitive must know how to evaluate itself on the CPU/GPU for the given inputs and populate the out...</div><div class="ttdef"><b>Definition</b> fast_primitives.h:296</div></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1fast_1_1_custom_kernel_html_a954893e07f0d36715b4e1e414b6f2153"><div class="ttname"><a href="classmlx_1_1core_1_1fast_1_1_custom_kernel.html#a954893e07f0d36715b4e1e414b6f2153">mlx::core::fast::CustomKernel::CustomKernel</a></div><div class="ttdeci">CustomKernel(Stream stream, std::string name, std::string source, std::tuple< int, int, int > grid, std::tuple< int, int, int > threadgroup, std::vector< CustomKernelShapeInfo > shape_infos, bool ensure_row_contiguous, std::optional< float > init_value)</div><div class="ttdef"><b>Definition</b> fast_primitives.h:278</div></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1fast_1_1_custom_kernel_html_a4ad1b7a9919753c759093f3e21a15bad"><div class="ttname"><a href="classmlx_1_1core_1_1fast_1_1_custom_kernel.html#a4ad1b7a9919753c759093f3e21a15bad">mlx::core::fast::CustomKernel::eval_cpu</a></div><div class="ttdeci">void eval_cpu(const std::vector< array > &inputs, std::vector< array > &outputs) override</div><div class="ttdoc">A primitive must know how to evaluate itself on the CPU/GPU for the given inputs and populate the out...</div><div class="ttdef"><b>Definition</b> fast_primitives.h:298</div></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1fast_1_1_custom_kernel_html_a954893e07f0d36715b4e1e414b6f2153"><div class="ttname"><a href="classmlx_1_1core_1_1fast_1_1_custom_kernel.html#a954893e07f0d36715b4e1e414b6f2153">mlx::core::fast::CustomKernel::CustomKernel</a></div><div class="ttdeci">CustomKernel(Stream stream, std::string name, std::string source, std::tuple< int, int, int > grid, std::tuple< int, int, int > threadgroup, std::vector< CustomKernelShapeInfo > shape_infos, bool ensure_row_contiguous, std::optional< float > init_value)</div><div class="ttdef"><b>Definition</b> fast_primitives.h:280</div></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1fast_1_1_layer_norm_html_a467fcf02b3ddf1d8b6d476b244ae3568"><div class="ttname"><a href="classmlx_1_1core_1_1fast_1_1_layer_norm.html#a467fcf02b3ddf1d8b6d476b244ae3568">mlx::core::fast::LayerNorm::DEFINE_PRINT</a></div><div class="ttdeci">DEFINE_PRINT(LayerNorm) bool is_equivalent(const Primitive &other) const override</div></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1fast_1_1_layer_norm_html_a5ac38d50e62850589bf51ee313303153"><div class="ttname"><a href="classmlx_1_1core_1_1fast_1_1_layer_norm.html#a5ac38d50e62850589bf51ee313303153">mlx::core::fast::LayerNorm::LayerNorm</a></div><div class="ttdeci">LayerNorm(Stream stream, std::function< std::vector< array >(std::vector< array >)> fallback, float eps)</div><div class="ttdef"><b>Definition</b> fast_primitives.h:100</div></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1fast_1_1_layer_norm_html_a5d7a4c1c9ee84e327d1c371733108c05"><div class="ttname"><a href="classmlx_1_1core_1_1fast_1_1_layer_norm.html#a5d7a4c1c9ee84e327d1c371733108c05">mlx::core::fast::LayerNorm::eval_cpu</a></div><div class="ttdeci">void eval_cpu(const std::vector< array > &inputs, std::vector< array > &outputs) override</div><div class="ttdoc">A primitive must know how to evaluate itself on the CPU/GPU for the given inputs and populate the out...</div><div class="ttdef"><b>Definition</b> fast_primitives.h:106</div></div>
|
||||
@@ -536,22 +538,23 @@ $(function(){initNavTree('fast__primitives_8h_source.html',''); initResizable(tr
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1fast_1_1_ro_p_e_html_a60b399d7f38c0f5f50342a6b97f0eb1a"><div class="ttname"><a href="classmlx_1_1core_1_1fast_1_1_ro_p_e.html#a60b399d7f38c0f5f50342a6b97f0eb1a">mlx::core::fast::RoPE::RoPE</a></div><div class="ttdeci">RoPE(Stream stream, std::function< std::vector< array >(std::vector< array >)> fallback, int dims, bool traditional, float base, float scale, bool forward)</div><div class="ttdef"><b>Definition</b> fast_primitives.h:159</div></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1fast_1_1_ro_p_e_html_a913b6b00fc518b25ac3947e4e15790f2"><div class="ttname"><a href="classmlx_1_1core_1_1fast_1_1_ro_p_e.html#a913b6b00fc518b25ac3947e4e15790f2">mlx::core::fast::RoPE::eval_gpu</a></div><div class="ttdeci">void eval_gpu(const std::vector< array > &inputs, std::vector< array > &outputs) override</div></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1fast_1_1_ro_p_e_html_ad999105414badd66c8fd9e069454a533"><div class="ttname"><a href="classmlx_1_1core_1_1fast_1_1_ro_p_e.html#ad999105414badd66c8fd9e069454a533">mlx::core::fast::RoPE::vjp</a></div><div class="ttdeci">std::vector< array > vjp(const std::vector< array > &primals, const std::vector< array > &cotangents, const std::vector< int > &argnums, const std::vector< array > &outputs) override</div><div class="ttdoc">The vector-Jacobian product.</div></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention_html_a505f38ba93a3499895f5312e0112e73d"><div class="ttname"><a href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#a505f38ba93a3499895f5312e0112e73d">mlx::core::fast::ScaledDotProductAttention::eval_gpu</a></div><div class="ttdeci">void eval_gpu(const std::vector< array > &inputs, std::vector< array > &outputs) override</div><div class="ttdef"><b>Definition</b> fast_primitives.h:217</div></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention_html_a64d2ce4b46b529a6a9ef068947bc623e"><div class="ttname"><a href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#a64d2ce4b46b529a6a9ef068947bc623e">mlx::core::fast::ScaledDotProductAttention::DEFINE_INPUT_OUTPUT_SHAPE</a></div><div class="ttdeci">DEFINE_INPUT_OUTPUT_SHAPE() auto state() const</div><div class="ttdef"><b>Definition</b> fast_primitives.h:226</div></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention_html_a09c99b460cca606b2ebb22f90b3d13a2"><div class="ttname"><a href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#a09c99b460cca606b2ebb22f90b3d13a2">mlx::core::fast::ScaledDotProductAttention::ScaledDotProductAttention</a></div><div class="ttdeci">ScaledDotProductAttention(Stream stream, std::function< std::vector< array >(std::vector< array >)> fallback, const float scale, const bool do_causal)</div><div class="ttdef"><b>Definition</b> fast_primitives.h:206</div></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention_html_a505f38ba93a3499895f5312e0112e73d"><div class="ttname"><a href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#a505f38ba93a3499895f5312e0112e73d">mlx::core::fast::ScaledDotProductAttention::eval_gpu</a></div><div class="ttdeci">void eval_gpu(const std::vector< array > &inputs, std::vector< array > &outputs) override</div><div class="ttdef"><b>Definition</b> fast_primitives.h:218</div></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention_html_a64d2ce4b46b529a6a9ef068947bc623e"><div class="ttname"><a href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#a64d2ce4b46b529a6a9ef068947bc623e">mlx::core::fast::ScaledDotProductAttention::DEFINE_INPUT_OUTPUT_SHAPE</a></div><div class="ttdeci">DEFINE_INPUT_OUTPUT_SHAPE() auto state() const</div><div class="ttdef"><b>Definition</b> fast_primitives.h:227</div></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention_html_a6cc2092fa5b8e7585921b8e0f3ec3db7"><div class="ttname"><a href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#a6cc2092fa5b8e7585921b8e0f3ec3db7">mlx::core::fast::ScaledDotProductAttention::DEFINE_PRINT</a></div><div class="ttdeci">DEFINE_PRINT(ScaledDotProductAttention)</div></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention_html_ab3f78d30e5bb3e76cfe701f2358e4748"><div class="ttname"><a href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#ab3f78d30e5bb3e76cfe701f2358e4748">mlx::core::fast::ScaledDotProductAttention::ScaledDotProductAttention</a></div><div class="ttdeci">ScaledDotProductAttention(Stream stream, std::function< std::vector< array >(std::vector< array >)> fallback, const float scale)</div><div class="ttdef"><b>Definition</b> fast_primitives.h:206</div></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention_html_ad51666e69f670e286293aff96eb435a9"><div class="ttname"><a href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#ad51666e69f670e286293aff96eb435a9">mlx::core::fast::ScaledDotProductAttention::eval_gpu</a></div><div class="ttdeci">void eval_gpu(const std::vector< array > &inputs, array &out)</div></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention_html_ae20851e002f7fcb6d4f97817596f6328"><div class="ttname"><a href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#ae20851e002f7fcb6d4f97817596f6328">mlx::core::fast::ScaledDotProductAttention::eval_cpu</a></div><div class="ttdeci">void eval_cpu(const std::vector< array > &inputs, std::vector< array > &outputs) override</div><div class="ttdoc">A primitive must know how to evaluate itself on the CPU/GPU for the given inputs and populate the out...</div><div class="ttdef"><b>Definition</b> fast_primitives.h:212</div></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention_html_ae20851e002f7fcb6d4f97817596f6328"><div class="ttname"><a href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#ae20851e002f7fcb6d4f97817596f6328">mlx::core::fast::ScaledDotProductAttention::eval_cpu</a></div><div class="ttdeci">void eval_cpu(const std::vector< array > &inputs, std::vector< array > &outputs) override</div><div class="ttdoc">A primitive must know how to evaluate itself on the CPU/GPU for the given inputs and populate the out...</div><div class="ttdef"><b>Definition</b> fast_primitives.h:213</div></div>
|
||||
<div class="ttc" id="aclassmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention_html_af08b1294f3f93505a96fdfa85b1edd62"><div class="ttname"><a href="classmlx_1_1core_1_1fast_1_1_scaled_dot_product_attention.html#af08b1294f3f93505a96fdfa85b1edd62">mlx::core::fast::ScaledDotProductAttention::is_equivalent</a></div><div class="ttdeci">bool is_equivalent(const Primitive &other) const override</div><div class="ttdoc">Equivalence check defaults to false unless overridden by the primitive.</div></div>
|
||||
<div class="ttc" id="agroup__ops_html_ga2a466024f8061febc0a64be557644cb0"><div class="ttname"><a href="group__ops.html#ga2a466024f8061febc0a64be557644cb0">mlx::core::std</a></div><div class="ttdeci">array std(const array &a, bool keepdims, int ddof=0, StreamOrDevice s={})</div><div class="ttdoc">Computes the standard deviation of the elements of an array.</div></div>
|
||||
<div class="ttc" id="agroup__ops_html_gabff758a5c1ce32ad7e8b78aba0164077"><div class="ttname"><a href="group__ops.html#gabff758a5c1ce32ad7e8b78aba0164077">mlx::core::dequantize</a></div><div class="ttdeci">array dequantize(const array &w, const array &scales, const array &biases, int group_size=64, int bits=4, StreamOrDevice s={})</div><div class="ttdoc">Dequantize a matrix produced by quantize()</div></div>
|
||||
<div class="ttc" id="anamespacemlx_1_1core_1_1fast_html"><div class="ttname"><a href="namespacemlx_1_1core_1_1fast.html">mlx::core::fast</a></div><div class="ttdef"><b>Definition</b> fast.h:9</div></div>
|
||||
<div class="ttc" id="anamespacemlx_1_1core_1_1fast_html"><div class="ttname"><a href="namespacemlx_1_1core_1_1fast.html">mlx::core::fast</a></div><div class="ttdef"><b>Definition</b> fast.h:10</div></div>
|
||||
<div class="ttc" id="aprimitives_8h_html"><div class="ttname"><a href="primitives_8h.html">primitives.h</a></div></div>
|
||||
<div class="ttc" id="asteel__attention_8h_html_abfa50278ba59a90e0acb7e5d94500741"><div class="ttname"><a href="steel__attention_8h.html#abfa50278ba59a90e0acb7e5d94500741">do_causal</a></div><div class="ttdeci">constant bool do_causal</div><div class="ttdef"><b>Definition</b> steel_attention.h:13</div></div>
|
||||
<div class="ttc" id="astructmlx_1_1core_1_1_stream_html"><div class="ttname"><a href="structmlx_1_1core_1_1_stream.html">mlx::core::Stream</a></div><div class="ttdef"><b>Definition</b> stream.h:9</div></div>
|
||||
<div class="ttc" id="astructmlx_1_1core_1_1fast_1_1_custom_kernel_shape_info_html"><div class="ttname"><a href="structmlx_1_1core_1_1fast_1_1_custom_kernel_shape_info.html">mlx::core::fast::CustomKernelShapeInfo</a></div><div class="ttdef"><b>Definition</b> fast_primitives.h:270</div></div>
|
||||
<div class="ttc" id="astructmlx_1_1core_1_1fast_1_1_custom_kernel_shape_info_html_a63954de7da62942ec69afcaaa19d46f2"><div class="ttname"><a href="structmlx_1_1core_1_1fast_1_1_custom_kernel_shape_info.html#a63954de7da62942ec69afcaaa19d46f2">mlx::core::fast::CustomKernelShapeInfo::strides</a></div><div class="ttdeci">bool strides</div><div class="ttdef"><b>Definition</b> fast_primitives.h:272</div></div>
|
||||
<div class="ttc" id="astructmlx_1_1core_1_1fast_1_1_custom_kernel_shape_info_html_a63db720fe0c2abc4b71e22a58a015f8a"><div class="ttname"><a href="structmlx_1_1core_1_1fast_1_1_custom_kernel_shape_info.html#a63db720fe0c2abc4b71e22a58a015f8a">mlx::core::fast::CustomKernelShapeInfo::shape</a></div><div class="ttdeci">bool shape</div><div class="ttdef"><b>Definition</b> fast_primitives.h:271</div></div>
|
||||
<div class="ttc" id="astructmlx_1_1core_1_1fast_1_1_custom_kernel_shape_info_html_ae605df33f449872e3da9777d97008051"><div class="ttname"><a href="structmlx_1_1core_1_1fast_1_1_custom_kernel_shape_info.html#ae605df33f449872e3da9777d97008051">mlx::core::fast::CustomKernelShapeInfo::ndim</a></div><div class="ttdeci">bool ndim</div><div class="ttdef"><b>Definition</b> fast_primitives.h:273</div></div>
|
||||
<div class="ttc" id="astructmlx_1_1core_1_1fast_1_1_custom_kernel_shape_info_html"><div class="ttname"><a href="structmlx_1_1core_1_1fast_1_1_custom_kernel_shape_info.html">mlx::core::fast::CustomKernelShapeInfo</a></div><div class="ttdef"><b>Definition</b> fast_primitives.h:272</div></div>
|
||||
<div class="ttc" id="astructmlx_1_1core_1_1fast_1_1_custom_kernel_shape_info_html_a63954de7da62942ec69afcaaa19d46f2"><div class="ttname"><a href="structmlx_1_1core_1_1fast_1_1_custom_kernel_shape_info.html#a63954de7da62942ec69afcaaa19d46f2">mlx::core::fast::CustomKernelShapeInfo::strides</a></div><div class="ttdeci">bool strides</div><div class="ttdef"><b>Definition</b> fast_primitives.h:274</div></div>
|
||||
<div class="ttc" id="astructmlx_1_1core_1_1fast_1_1_custom_kernel_shape_info_html_a63db720fe0c2abc4b71e22a58a015f8a"><div class="ttname"><a href="structmlx_1_1core_1_1fast_1_1_custom_kernel_shape_info.html#a63db720fe0c2abc4b71e22a58a015f8a">mlx::core::fast::CustomKernelShapeInfo::shape</a></div><div class="ttdeci">bool shape</div><div class="ttdef"><b>Definition</b> fast_primitives.h:273</div></div>
|
||||
<div class="ttc" id="astructmlx_1_1core_1_1fast_1_1_custom_kernel_shape_info_html_ae605df33f449872e3da9777d97008051"><div class="ttname"><a href="structmlx_1_1core_1_1fast_1_1_custom_kernel_shape_info.html#ae605df33f449872e3da9777d97008051">mlx::core::fast::CustomKernelShapeInfo::ndim</a></div><div class="ttdeci">bool ndim</div><div class="ttdef"><b>Definition</b> fast_primitives.h:275</div></div>
|
||||
</div><!-- fragment --></div><!-- contents -->
|
||||
</div><!-- doc-content -->
|
||||
<!-- start footer part -->
|
||||
|
||||
Reference in New Issue
Block a user