tensor_shape.bend source

tensor_shape.bend on the hub · documented module

# Storage-independent extents. View derives this descriptor from its axes.import Baseimport ./machine_limits.bend as Machineimport ./math_checked_word.bend as CheckedWordtype Shape is Data:  Shape{extents: List<&2,U32>}def rank(shape: Shape) -> Nat:  Shape{extents} = shape  List.length(&2,U32,extents)def extent(shape: Shape,axis: Nat) -> Maybe<&2,U32>:  Shape{extents} = shape  List.get(&2,U32,extents,axis)def same_extents(left: List<&2,U32>,right: List<&2,U32>) -> Bool:  match left right:    case Nil{} Nil{}: True{}    case Con{x,xs} Con{y,ys}: U32.is_eq(x,y) && same_extents(xs,ys)    case left right: False{}def same(left: Shape,right: Shape) -> Bool:  Shape{xs} = left  Shape{ys} = right  same_extents(xs,ys)def has_zero(extents: List<&2,U32>) -> Bool:  match extents:    case Nil{}: False{}    case Con{x,xs}: U32.is_zero(x) || has_zero(xs)def checked_multiply(left: U32,right: U32) -> Maybe<&2,U32>:  CheckedWord.multiply(left,right)def product(extents: List<&2,U32>,accumulator: Maybe<&2,U32>) -> Maybe<&2,U32>:  match extents accumulator:    case Nil{} result: result    case Con{x,xs} Some{value}: product(xs,checked_multiply(value,x))    case Con{x,xs} None{}: None{}def count_nonempty(empty: Bool,extents: List<&2,U32>) -> Maybe<&2,U32>:  match empty:    case True{}: Some{0}    case False{}: product(extents,Some{1})# Scalar [] has one element. Any zero extent makes an empty shape, even when# an earlier partial product would overflow. Nonempty overflow is rejected.def numel(shape: Shape) -> Maybe<&2,U32>:  Shape{+extents} = shape  count_nonempty(has_zero(extents),extents)# NumPy broadcasting of one axis: equal extents, or one of them is 1.def broadcast_extent(same: Bool,left_one: Bool,right_one: Bool,left: U32,right: U32) -> Maybe<&2,U32>:  match same left_one right_one:    case True{} left_one right_one: Some{left}    case False{} True{} right_one: Some{right}    case False{} False{} True{}: Some{left}    case False{} False{} False{}: None{}def prepend_extent(extent: Maybe<&2,U32>,rest: Maybe<&2,List<&2,U32>>) -> Maybe<&2,List<&2,U32>>:  match extent rest:    case Some{extent} Some{rest}: Some{extent <> rest}    case extent rest: None{}def broadcast_same_rank(left: List<&2,U32>,right: List<&2,U32>) -> Maybe<&2,List<&2,U32>>:  match left right:    case Nil{} Nil{}: Some{[]}    case Con{+x,xs} Con{+y,ys}:      prepend_extent(broadcast_extent(U32.is_eq(x,y),U32.is_eq(x,1),U32.is_eq(y,1),x,y),broadcast_same_rank(xs,ys))    case left right: None{}# The extents with leading axes of extent 1 added up to rank.def with_rank(+extents: List<&2,U32>,rank: Nat) -> List<&2,U32>:  List.append(&2,U32,List.replicate(U32,Nat.sub(rank,List.length(&2,U32,extents)),1),extents)# The common extents of two operands under NumPy broadcasting: the extents are# aligned at the last axis, a missing axis counts as 1, and on each axis the# extents are equal or one of them is 1.def broadcast(+left: List<&2,U32>,+right: List<&2,U32>) -> Maybe<&2,List<&2,U32>>:  broadcast_same_rank(with_rank(left,List.length(&2,U32,right)),with_rank(right,List.length(&2,U32,left)))# The extents without one axis.def without_axis(+extents: List<&2,U32>,+axis: Nat) -> List<&2,U32>:  List.append(&2,U32,List.take(&2,U32,extents,axis),List.drop(&2,U32,extents,1n+axis))def sum_when(valid: Bool,value: U32) -> Maybe<&2,U32>:  match valid:    case True{}: Some{value}    case False{}: None{}def checked_add(+left: U32,+right: U32) -> Maybe<&2,U32>:  sum_when(U32.is_le(right,(4294967295 - left : U32)),(left + right : U32))def checked_add_to(left: Maybe<&2,U32>,right: U32) -> Maybe<&2,U32>:  match left:    case Some{left}: checked_add(left,right)    case None{}: None{}# Each extent with before elements added in front and after elements behind.# The three lists have one entry per axis and no sum may wrap.def padded(extents: List<&2,U32>,before: List<&2,U32>,after: List<&2,U32>) -> Maybe<&2,List<&2,U32>>:  match extents before after:    case Nil{} Nil{} Nil{}: Some{[]}    case Con{extent,extents} Con{front,before} Con{back,after}:      prepend_extent(checked_add_to(checked_add(front,extent),back),padded(extents,before,after))    case extents before after: None{}