tensor_select.bend checks

raw source on the hub · import stelliferous@0.0.2.0/tensor_select.bend as Tensor_select

Three-input selection: shared shape/band helpers and the common lane loop.

9 imports

import Base
import ./tensor.bend as Tensors
import ./tensor_shape.bend as Shapes
import ./tensor_layout.bend as Layout
import ./tensor_lanes.bend as Lanes
import ./tensor_band.bend as Bands
import ./tensor_band_schedule.bend as Schedule
import ./storage_buffer.bend as Storage
import ./traversal_partition.bend as Partition

Types

type BandPlan source · line 23 · raw

Data

BandPlan

@geometry:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor_band.Geometry -> @layout:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor_band.Plan -> BandPlan

Definitions

def dense source · line 12 · raw

@+shape:List<&2, U32> -> @+count:U32 -> @+workers:U32 -> @mask:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32> -> @yes:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32> -> @no:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32> -> 0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32>

def expanded source · line 18 · raw

@same:Bool -> @shape:List<&2, U32> -> @tensor:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32> -> 0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32>

def planned source · line 26 · raw

@shape:List<&2, U32> -> @fill:F32 -> @status:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Status -> @workers:U32 -> @+geometry:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor_band.Geometry -> @owners:Pair(0x9be3b13bf249759bc81e1958bcd1a4c0/tensor_band.Bands(F32), 0x9be3b13bf249759bc81e1958bcd1a4c0/tensor_band.Plan) -> Pair(0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32>, Maybe<&2, BandPlan>)

def plan source · line 31 · raw

@tensor:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32> -> Pair(0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32>, Maybe<&2, BandPlan>)

def common source · line 36 · raw

@left:Maybe<&2, BandPlan> -> @right:Maybe<&2, BandPlan> -> Maybe<&2, BandPlan>

def selected_leaf source · line 42 · raw

@geometry:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor_band.Geometry -> @first:U32 -> @rows:U32 -> @input:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor_band.Paired<0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32>, 0x9be3b13bf249759bc81e1958bcd1a4c0/tensor_band.Paired<0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32>, 0x9be3b13bf249759bc81e1958bcd1a4c0/storage_buffer.Buffer<F32>>> -> @resource:Unit -> 0x9be3b13bf249759bc81e1958bcd1a4c0/tensor_band.Bands(F32)

def selected_bands source · line 51 · raw

@+geometry:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor_band.Geometry -> @+layout:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor_band.Plan -> @+parallel:Bool -> @mask:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor_band.Bands(F32) -> @yes:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor_band.Bands(F32) -> @no:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor_band.Bands(F32) -> 0x9be3b13bf249759bc81e1958bcd1a4c0/tensor_band.Bands(F32)

def dispatch source · line 58 · raw

@layout:Maybe<&2, BandPlan> -> @shape:List<&2, U32> -> @count:U32 -> @fill:F32 -> @workers:U32 -> @mask:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32> -> @yes:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32> -> @no:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32> -> 0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32>

def matching source · line 68 · raw

@+shape:List<&2, U32> -> @+count:U32 -> @+fill:F32 -> @+workers:U32 -> @mask:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32> -> @yes:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32> -> @no:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32> -> 0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32>

def shapes source · line 74 · raw

@same:Bool -> @+shape:List<&2, U32> -> @count:Maybe<&2, U32> -> @mask_shape:List<&2, U32> -> @yes_shape:List<&2, U32> -> @no_shape:List<&2, U32> -> @fill:F32 -> @workers:U32 -> @mask:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32> -> @yes:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32> -> @no:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32> -> 0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32>

def masked_shape source · line 83 · raw

@common:Maybe<&2, List<&2, U32>> -> @+mask_shape:List<&2, U32> -> @+yes_shape:List<&2, U32> -> @+no_shape:List<&2, U32> -> @values_shape:List<&2, U32> -> @fill:F32 -> @workers:U32 -> @mask:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32> -> @yes:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32> -> @no:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32> -> 0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32>

def values_shape source · line 90 · raw

@common:Maybe<&2, List<&2, U32>> -> @+mask_shape:List<&2, U32> -> @+yes_shape:List<&2, U32> -> @no_shape:List<&2, U32> -> @fill:F32 -> @workers:U32 -> @mask:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32> -> @yes:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32> -> @no:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32> -> 0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32>

def status source · line 96 · raw

@status:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Status -> @mask_shape:List<&2, U32> -> @+yes_shape:List<&2, U32> -> @+no_shape:List<&2, U32> -> @fill:F32 -> @workers:U32 -> @mask:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32> -> @yes:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32> -> @no:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32> -> 0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32>

def select source · line 102 · raw

@mask:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32> -> @yes:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32> -> @no:0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32> -> 0x9be3b13bf249759bc81e1958bcd1a4c0/tensor.Tensor<F32>