Matmul.product_tree: def of bendtorch-core@0.1.0.0
Import: import bendtorch-core@0.1.0.0/src/tensor/matmul.bend as Matmul
Lines 598-619 of bendtorch-core@0.1.0.0/src/tensor/matmul.bend: https://jend.leeeet.dev/read.txt?path=bendtorch-core@0.1.0.0/src/tensor/matmul.bend&start_line=598&end_line=619

import Base
import bendtorch-core@0.1.0.0/src/tensor/parallel.bend as Parallel

# The cells base .. base + 2^depth on one thread: `Parallel.tree` with the
# product cell, but each subtree of eight cells in one row is a row block.
def product_tree(+layout: Layout, +depth: Nat, +count: Nat, +base: Nat, state: Array<F32> & Array<F32>) ->
  Array<F32> & (Array<F32> & Array<F32>):
  match depth:
    case 0n:
      Parallel.tree(~Layout, ~Array<F32>, ~Array<F32>, ~product_cell_at, layout, 0n, count, base, state)
    case 1n+d1:
      match d1:
        case 0n:
          Parallel.tree(~Layout, ~Array<F32>, ~Array<F32>, ~product_cell_at, layout, 1n, count, base, state)
        case 1n+d2:
          match d2:
            case 0n:
              Parallel.tree(~Layout, ~Array<F32>, ~Array<F32>, ~product_cell_at, layout, 2n, count, base, state)
            case 1n+d3:
              match d3:
                case 0n:
                  block_or_tree(layout, count, base, state, one_row(layout, count, base))
                case 1n+d4:
                  Parallel.node(~Array<F32>, ~Array<F32>, l => r => product_tree(layout, d1, count, Parallel.upper(d1, base), (l, r)),
                    product_tree(layout, d1, count, base, state))

Uses (12):
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Array
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=F32
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Nat
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Pair
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Sigma
  https://jend.leeeet.dev/definition.txt?path=bendtorch-core@0.1.0.0/src/tensor/matmul.bend&name=Layout
  https://jend.leeeet.dev/definition.txt?path=bendtorch-core@0.1.0.0/src/tensor/matmul.bend&name=block_or_tree
  https://jend.leeeet.dev/definition.txt?path=bendtorch-core@0.1.0.0/src/tensor/matmul.bend&name=one_row
  https://jend.leeeet.dev/definition.txt?path=bendtorch-core@0.1.0.0/src/tensor/matmul.bend&name=product_cell_at
  https://jend.leeeet.dev/definition.txt?path=bendtorch-core@0.1.0.0/src/tensor/parallel.bend&name=node
  https://jend.leeeet.dev/definition.txt?path=bendtorch-core@0.1.0.0/src/tensor/parallel.bend&name=tree
  https://jend.leeeet.dev/definition.txt?path=bendtorch-core@0.1.0.0/src/tensor/parallel.bend&name=upper