def plan_size

stelliferous@0.0.2.0 · tensor_band_pool.bend:25

import Base
import stelliferous@0.0.2.0/tensor_band.bend as Bands
import stelliferous@0.0.2.0/tensor_band_window.bend as BandWindow
import stelliferous@0.0.2.0/tensor_band_schedule.bend as BandSchedule

25def plan_size(layout: Bands.Plan,+window: BandWindow.Geometry) -> BandSchedule.Size:
26  match layout:
27    case Bands.Rows{first,rows}: window_size(BandWindow.window(window,first,rows))
28    case Bands.Split{left,right}: BandSchedule.merge(plan_size(left,window),plan_size(right,window))