TensorBandSchedule.allocation_elements: def of stelliferous@0.0.2.0
Import: import stelliferous@0.0.2.0/tensor_band_schedule.bend as TensorBandSchedule
Lines 40-45 of stelliferous@0.0.2.0/tensor_band_schedule.bend: https://jend.leeeet.dev/read.txt?path=stelliferous@0.0.2.0/tensor_band_schedule.bend&start_line=40&end_line=45

import Base
import stelliferous@0.0.2.0/tensor_band.bend as Bands
import stelliferous@0.0.2.0/storage_buffer.bend as Storage

# For equal-shaped operands the logical read/write volume of repartition is
# identical. The target owners' rounded capacities decide the remaining cost.
def allocation_elements(layout: Bands.Plan,+geometry: Bands.Geometry) -> Nat:
  match layout:
    case Bands.Rows{first,rows}: U32.to_nat(U32.shln(1,Storage.depth_for(Bands.local_count(geometry,rows))))
    case Bands.Split{left,right}: Nat.add(allocation_elements(left,geometry),allocation_elements(right,geometry))

Uses (8):
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Nat
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Nat.add
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=U32.shln
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=U32.to_nat
  https://jend.leeeet.dev/definition.txt?path=stelliferous@0.0.2.0/storage_buffer.bend&name=depth_for
  https://jend.leeeet.dev/definition.txt?path=stelliferous@0.0.2.0/tensor_band.bend&name=Geometry
  https://jend.leeeet.dev/definition.txt?path=stelliferous@0.0.2.0/tensor_band.bend&name=Plan
  https://jend.leeeet.dev/definition.txt?path=stelliferous@0.0.2.0/tensor_band.bend&name=local_count