def drop_broadcast

bendtorch-core@0.1.0.0 · src/shape/broadcast_proofs.bend:194

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape

194def drop_broadcast(+count: Nat, +dims: List<&2, Nat>, +values: List<&2, Nat>) -> {List.drop(&2, Nat, Shape.broadcast_indices(Shape.pad(count, dims), values), count) == Shape.broadcast_indices(dims, List.drop(&2, Nat, values, count)) : List<&2, Nat>}:
195  match count:
196    case 0n:
197      %drop_zero(values) : {List.drop(&2, Nat, Shape.broadcast_indices(dims, values), 0n) == Shape.broadcast_indices(dims, _) : List<&2, Nat>}
198      %drop_zero(Shape.broadcast_indices(dims, values)) : {_ == Shape.broadcast_indices(dims, values) : List<&2, Nat>}
199      {==}
200    case 1n+rest:
201      match values:
202        case Nil{}:
203          broadcast_nil(dims)
204        case +head <> +tail:
205          drop_broadcast(rest, dims, tail)