[{"data":1,"prerenderedAt":38},["ShallowReactive",2],{"chapter:kernels\u002Fpatterns\u002Freductions.json":3},{"project":4,"route":5,"title":6,"titleHtml":6,"navTitle":6,"part":7,"sourcePath":8,"editUrl":9,"html":10,"toc":11,"hasMermaid":31,"prev":32,"next":35},"kernels","\u002Fkernels\u002Fpatterns\u002Freductions","Reductions and Scans","Real Patterns","patterns\u002Freductions.md","https:\u002F\u002Fgithub.com\u002Fteenygrad\u002Fteenygrad\u002Fedit\u002Fmain\u002Fbooks\u002Fkernels\u002Fsrc\u002Fpatterns\u002Freductions.md","\u003Cp>Chapter 10 used two reductions to build a softmax. This chapter is the rest of\nthe family: what is available, how to write one the library does not have, and\nthe difference between a reduction and a scan.\u003C\u002Fp>\n\u003Ch2 id=\"the-two-shapes\">The two shapes\u003C\u002Fh2>\n\u003Cp>A \u003Cstrong>reduction\u003C\u002Fstrong> turns many values into one. Sum, maximum, count.\u003C\u002Fp>\n\u003Cpre class=\"code-panel\" data-lang=\"text\">\u003Ccode>[3, 1, 4, 1, 5]  --sum-->  14\n\u003C\u002Fcode>\u003C\u002Fpre>\n\u003Cp>A \u003Cstrong>scan\u003C\u002Fstrong> turns many values into the same many values, each holding the\nreduction of everything up to it. Also called a prefix operation.\u003C\u002Fp>\n\u003Cpre class=\"code-panel\" data-lang=\"text\">\u003Ccode>[3, 1, 4, 1, 5]  --cumsum-->  [3, 4, 8, 9, 14]\n\u003C\u002Fcode>\u003C\u002Fpre>\n\u003Cp>Reductions are cheap and common. Scans are less common and more expensive,\nbecause every output depends on every earlier input, but they are how you\nimplement anything involving running totals — offsets into a variable-length\nbuffer, sampling from a distribution, sorting.\u003C\u002Fp>\n\u003Ch2 id=\"what-is-built-in\">What is built in\u003C\u002Fh2>\n\u003Cp>Reductions, all taking an \u003Ccode>axis\u003C\u002Fcode> and \u003Ccode>keep_dims\u003C\u002Fcode>:\u003C\u002Fp>\n\u003Ctable>\n\u003Cthead>\n\u003Ctr>\n\u003Cth>Method\u003C\u002Fth>\n\u003Cth>Result\u003C\u002Fth>\n\u003C\u002Ftr>\n\u003C\u002Fthead>\n\u003Ctbody>\n\u003Ctr>\n\u003Ctd>\u003Ccode>sum\u003C\u002Fcode>, \u003Ccode>max\u003C\u002Fcode>, \u003Ccode>min\u003C\u002Fcode>\u003C\u002Ftd>\n\u003Ctd>The obvious\u003C\u002Ftd>\n\u003C\u002Ftr>\n\u003Ctr>\n\u003Ctd>\u003Ccode>max_with_indices\u003C\u002Fcode>, \u003Ccode>min_with_indices\u003C\u002Fcode>\u003C\u002Ftd>\n\u003Ctd>The value \u003Cem>and\u003C\u002Fem> where it was\u003C\u002Ftd>\n\u003C\u002Ftr>\n\u003Ctr>\n\u003Ctd>\u003Ccode>argmax\u003C\u002Fcode>, \u003Ccode>argmin\u003C\u002Fcode>\u003C\u002Ftd>\n\u003Ctd>Just where it was\u003C\u002Ftd>\n\u003C\u002Ftr>\n\u003Ctr>\n\u003Ctd>\u003Ccode>xor_sum\u003C\u002Fcode>\u003C\u002Ftd>\n\u003Ctd>XOR-fold, integers only\u003C\u002Ftd>\n\u003C\u002Ftr>\n\u003C\u002Ftbody>\n\u003C\u002Ftable>\n\u003Cp>Scans and friends:\u003C\u002Fp>\n\u003Ctable>\n\u003Cthead>\n\u003Ctr>\n\u003Cth>Method\u003C\u002Fth>\n\u003Cth>Result\u003C\u002Fth>\n\u003C\u002Ftr>\n\u003C\u002Fthead>\n\u003Ctbody>\n\u003Ctr>\n\u003Ctd>\u003Ccode>cumsum\u003C\u002Fcode>, \u003Ccode>cumprod\u003C\u002Fcode>\u003C\u002Ftd>\n\u003Ctd>Running total \u002F product along an axis\u003C\u002Ftd>\n\u003C\u002Ftr>\n\u003Ctr>\n\u003Ctd>\u003Ccode>sort\u003C\u002Fcode>\u003C\u002Ftd>\n\u003Ctd>Sorted along a dimension\u003C\u002Ftd>\n\u003C\u002Ftr>\n\u003Ctr>\n\u003Ctd>\u003Ccode>histogram\u003C\u002Fcode>\u003C\u002Ftd>\n\u003Ctd>Counts into \u003Ccode>num_bins\u003C\u002Fcode> bins of width 1\u003C\u002Ftd>\n\u003C\u002Ftr>\n\u003C\u002Ftbody>\n\u003C\u002Ftable>\n\u003Cp>Two conventions apply throughout, and both were introduced in Chapter 10:\u003C\u002Fp>\n\u003Cul>\n\u003Cli>\u003Cstrong>\u003Ccode>axis\u003C\u002Fcode>\u003C\u002Fstrong> — \u003Ccode>Some(n)\u003C\u002Fcode> reduces dimension \u003Ccode>n\u003C\u002Fcode>, \u003Ccode>None\u003C\u002Fcode> reduces everything.\u003C\u002Fli>\n\u003Cli>\u003Cstrong>\u003Ccode>keep_dims\u003C\u002Fcode>\u003C\u002Fstrong> — \u003Ccode>true\u003C\u002Fcode> leaves a length-1 dimension so the result can\nbroadcast back against the input. This is almost always what you want inside\na kernel.\u003C\u002Fli>\n\u003C\u002Ful>\n\u003Cp>The \u003Ccode>*_with_indices\u003C\u002Fcode> and \u003Ccode>arg*\u003C\u002Fcode> variants also take \u003Ccode>tie_break_left\u003C\u002Fcode>. With \u003Ccode>true\u003C\u002Fcode>\nthe leftmost of equal values wins. It matters more than it sounds: if your\nkernel and your reference implementation break ties differently, a test on data\nwith duplicates fails for a reason that looks like a real bug.\u003C\u002Fp>\n\u003Ch2 id=\"a-worked-one\">A worked one\u003C\u002Fh2>\n\u003Cp>The library’s sum-reduction kernel:\u003C\u002Fp>\n\u003Cpre data-lang=\"rust\" class=\"shiki teeny-datasheet\" style=\"background-color:#16181a;color:#e6e8e3\" tabindex=\"0\">\u003Ccode>\u003Cspan class=\"line\">\u003Cspan style=\"color:#8A9088\">#[\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">kernel\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">]\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#FF5F9E\">pub\u003C\u002Fspan>\u003Cspan style=\"color:#FF5F9E\"> fn\u003C\u002Fspan>\u003Cspan style=\"color:#7FB6D9\"> reduce_sum_forward\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> Triton\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> D\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> Num\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003Cspan style=\"color:#FF5F9E\"> const\u003C\u002Fspan>\u003Cspan style=\"color:#B79AD4\"> BLOCK_INNER\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> i32\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>(\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#E6E8E3\">    x_ptr\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">Pointer\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">D\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>,\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#E6E8E3\">    y_ptr\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">Pointer\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">D\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>,\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#E6E8E3\">    n_inner\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> i32\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#E6E8E3\">    n_outer\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> i32\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#8A9088\">)\u003C\u002Fspan>\u003Cspan style=\"color:#FF5F9E\"> where\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#6FBF98\">    T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">I32Tensor\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> types\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">Tensor\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">i32\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> 1\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>,\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#6FBF98\">    T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">I32Tensor\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> Comparison\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">i32\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> BoolTensor\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"> =\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">BoolTensor\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>,\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#6FBF98\">    T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">Pointer\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">D\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>:\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> AddOffsets\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">i32\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> 1\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">I32Tensor\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> Output\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"> =\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">Tensor\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">Pointer\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">D\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>>>,\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#8A9088\">{\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#FF5F9E\">    let\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> row \u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">=\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#7FB6D9\">program_id\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">(\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">Axis\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">X\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">);\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#FF5F9E\">    if\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> row \u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>=\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> n_outer \u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">{\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#FF5F9E\">        return\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">;\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#8A9088\">    }\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#FF5F9E\">    let\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> col_offsets \u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">=\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#7FB6D9\">arange\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">(\u003C\u002Fspan>\u003Cspan style=\"color:#B79AD4\">0\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003Cspan style=\"color:#B79AD4\"> BLOCK_INNER\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">);\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#FF5F9E\">    let\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> offsets \u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">=\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> col_offsets \u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">+\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> row \u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">*\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> n_inner\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">;\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#FF5F9E\">    let\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> mask \u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">=\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> col_offsets\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">.\u003C\u002Fspan>\u003Cspan style=\"color:#7FB6D9\">lt\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">(\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\">n_inner\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">);\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#FF5F9E\">    let\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> x \u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">=\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#7FB6D9\">load\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">(\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#E6E8E3\">        x_ptr\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">.\u003C\u002Fspan>\u003Cspan style=\"color:#7FB6D9\">add_offsets\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">(\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\">offsets\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">),\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#6FBF98\">        Some\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">(\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\">mask\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">),\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#6FBF98\">        Some\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">(\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#7FB6D9\">zeros\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::<\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">D\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>(&[\u003C\u002Fspan>\u003Cspan style=\"color:#B79AD4\">BLOCK_INNER\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">])),\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#8A9088\">        &[],\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#6FBF98\">        None\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#6FBF98\">        None\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#6FBF98\">        None\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#B79AD4\">        false\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#8A9088\">    );\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#FF5F9E\">    let\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> sum \u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">=\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#7FB6D9\">sum\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">(\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\">x\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> Some\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">(\u003C\u002Fspan>\u003Cspan style=\"color:#B79AD4\">0\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">),\u003C\u002Fspan>\u003Cspan style=\"color:#B79AD4\"> true\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">);\u003C\u002Fspan>\u003Cspan style=\"color:#7F877D;font-style:italic\"> \u002F\u002F [1] or scalar\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#FF5F9E\">    let\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> row_offsets \u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">=\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#7FB6D9\">arange\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">(\u003C\u002Fspan>\u003Cspan style=\"color:#B79AD4\">0\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003Cspan style=\"color:#B79AD4\"> 1\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">)\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"> +\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> row\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">;\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#6FBF98\">    T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#7FB6D9\">store\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">(\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\">y_ptr\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">.\u003C\u002Fspan>\u003Cspan style=\"color:#7FB6D9\">add_offsets\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">(\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\">row_offsets\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">),\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> sum\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> None\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"> &[],\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> None\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> None\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">);\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#8A9088\">}\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003C\u002Fspan>\u003C\u002Fcode>\u003C\u002Fpre>\n\u003Cp>The pattern is Chapter 10’s, without the numerical-stability step: one program\nper output, load the slice being reduced, mask it, reduce, store one value.\u003C\u002Fp>\n\u003Cp>Note the masked-lane fill. It has to be the identity for the operation — zero\nfor a sum — or the masked lanes contribute garbage. Chapter 10 has the table of\nidentities; this is the kernel where getting it wrong is easiest, because the\nresult is a single number that looks plausible.\u003C\u002Fp>\n\u003Ch2 id=\"writing-your-own\">Writing your own\u003C\u002Fh2>\n\u003Cp>When the operation you need is not in the list, \u003Ccode>T::reduce\u003C\u002Fcode> takes a combine\nfunction:\u003C\u002Fp>\n\u003Cpre data-lang=\"rust\" class=\"shiki teeny-datasheet\" style=\"background-color:#16181a;color:#e6e8e3\" tabindex=\"0\">\u003Ccode>\u003Cspan class=\"line\">\u003Cspan style=\"color:#FF5F9E\">fn\u003C\u002Fspan>\u003Cspan style=\"color:#7FB6D9\"> combine_max_abs\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> Triton\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> D\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> Float\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>(\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\">a\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">Tensor\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">D\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>,\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> b\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">Tensor\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">D\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>)\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"> ->\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">Tensor\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">D\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"> {\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#6FBF98\">    T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#7FB6D9\">maximum\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">(\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#7FB6D9\">abs\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">(\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\">a\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">),\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#7FB6D9\">abs\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">(\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\">b\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">))\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#8A9088\">}\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003Cspan style=\"color:#FF5F9E\">let\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> result \u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">=\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#7FB6D9\">reduce\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">(\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\">x\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003Cspan style=\"color:#B79AD4\"> 0\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> combine_max_abs\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::<\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">T\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> D\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>,\u003C\u002Fspan>\u003Cspan style=\"color:#B79AD4\"> true\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">);\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003C\u002Fspan>\u003C\u002Fcode>\u003C\u002Fpre>\n\u003Cp>and \u003Ccode>T::associative_scan\u003C\u002Fcode> is the same idea for a prefix operation, plus a\n\u003Ccode>reverse\u003C\u002Fcode> flag.\u003C\u002Fp>\n\u003Cp>Two requirements, and the second is a genuine gotcha.\u003C\u002Fp>\n\u003Cp>\u003Cstrong>The function must be associative.\u003C\u002Fstrong> The compiler builds a tree and combines\npairs in an unspecified order, so \u003Ccode>f(f(a, b), c)\u003C\u002Fcode> and \u003Ccode>f(a, f(b, c))\u003C\u002Fcode> must agree.\nMaximum is associative. Subtraction is not. Floating-point addition is not\n\u003Cem>exactly\u003C\u002Fem> associative, which is why a GPU sum and a CPU sum can differ in the\nlast bits — expected, and not a bug.\u003C\u002Fp>\n\u003Cp>\u003Cstrong>It must be a \u003Ccode>fn\u003C\u002Fcode> pointer, not a closure.\u003C\u002Fstrong>\u003C\u002Fp>\n\u003Cpre data-lang=\"rust\" class=\"shiki teeny-datasheet\" style=\"background-color:#16181a;color:#e6e8e3\" tabindex=\"0\">\u003Ccode>\u003Cspan class=\"line\">\u003Cspan style=\"color:#FF5F9E\">fn\u003C\u002Fspan>\u003Cspan style=\"color:#7FB6D9\"> reduce\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">D\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> O\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>(\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\">x\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"> ...,\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> axis\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> i32\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">,\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> combine_fn\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#FF5F9E\"> fn\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">(\u003C\u002Fspan>\u003Cspan style=\"color:#FF5F9E\">Self\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">Tensor\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">O\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>,\u003C\u002Fspan>\u003Cspan style=\"color:#FF5F9E\"> Self\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">Tensor\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">O\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>)\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"> ->\u003C\u002Fspan>\u003Cspan style=\"color:#FF5F9E\"> Self\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">::\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">Tensor\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"><\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\">O\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">>,\u003C\u002Fspan>\u003Cspan style=\"color:#E6E8E3\"> keep_dims\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">:\u003C\u002Fspan>\u003Cspan style=\"color:#6FBF98\"> bool\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\">)\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"> ->\u003C\u002Fspan>\u003Cspan style=\"color:#8A9088\"> ...;\u003C\u002Fspan>\u003C\u002Fspan>\n\u003Cspan class=\"line\">\u003C\u002Fspan>\u003C\u002Fcode>\u003C\u002Fpre>\n\u003Cp>A closure that captures anything is rejected. This follows directly from\nChapter 3: the kernel body is compiled from captured source text, so the combine\nfunction has to be a statically-known name that can be written out. A closure’s\ncaptured environment cannot be.\u003C\u002Fp>\n\u003Cp>Python Triton has the same restriction — the combine function needs\n\u003Ccode>@triton.jit\u003C\u002Fcode> — but the Rust error message does not mention kernels at all. It\nis a generic closure-coercion complaint, and it is worth recognising:\u003C\u002Fp>\n\u003Cpre class=\"code-panel\" data-lang=\"text\">\u003Ccode>expected fn pointer `fn(...) -> ...`\nfound closure `[closure@src\u002F...]`\n\u003C\u002Fcode>\u003C\u002Fpre>\n\u003Ch2 id=\"cost\">Cost\u003C\u002Fh2>\n\u003Cp>A reduction over \u003Ccode>n\u003C\u002Fcode> lanes takes \u003Ccode>log2(n)\u003C\u002Fcode> steps, not \u003Ccode>n\u003C\u002Fcode>. Halving the working\nset each round is what makes it cheap, and it is why powers of two matter for\nblock sizes — Chapter 6’s second rule.\u003C\u002Fp>\n\u003Cp>A scan is more expensive: the standard algorithm makes two passes over the tree,\nso roughly twice the work of a reduction. Still \u003Ccode>O(log n)\u003C\u002Fcode> depth, but do not\nreach for \u003Ccode>cumsum\u003C\u002Fcode> where \u003Ccode>sum\u003C\u002Fcode> would do.\u003C\u002Fp>\n\u003Ch2 id=\"reducing-across-programs\">Reducing across programs\u003C\u002Fh2>\n\u003Cp>Everything here reduces \u003Cem>within\u003C\u002Fem> one program. Getting a single number out of a\nwhole tensor that does not fit in one block is a different problem, and there are\ntwo answers.\u003C\u002Fp>\n\u003Cp>The first is two kernels: one produces a partial result per program, the second\nreduces those partials. Predictable, deterministic, and needs a scratch buffer.\u003C\u002Fp>\n\u003Cp>The second is atomics — every program folds its partial into one location in\nmemory. That is the next chapter, along with why it is not the default.\u003C\u002Fp>\n",[12,16,19,22,25,28],{"id":13,"text":14,"level":15},"the-two-shapes","The two shapes",2,{"id":17,"text":18,"level":15},"what-is-built-in","What is built in",{"id":20,"text":21,"level":15},"a-worked-one","A worked one",{"id":23,"text":24,"level":15},"writing-your-own","Writing your own",{"id":26,"text":27,"level":15},"cost","Cost",{"id":29,"text":30,"level":15},"reducing-across-programs","Reducing across programs",false,{"title":33,"titleHtml":33,"route":34},"Fusing an Epilogue","\u002Fkernels\u002Fpatterns\u002Ffusion",{"title":36,"titleHtml":36,"route":37},"Atomics","\u002Fkernels\u002Fpatterns\u002Fatomics",1786271829981]