def PlanKeeps

bendtorch-core@0.1.0.0 · src/autograd/sum_dims_LAWS.bend:42

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

42# What a plan holds: the kept shape's dimensions and the dropped ones.
43def PlanKeeps(+axes: List<&2, U32>, +dimensions: List<&2, U32>, result: Result<&2, &2, SumDims.Error, SumDims.Plan>) -> Type:
44  match result:
45    case Fail{error}:
46      Unit
47    case Done{SumDims.Plan{kept, dropped}}:
48      {Shape.dimensions(kept) == SumDims.kept_dimensions(axes, dimensions) : List<&2, U32>} &
49        {dropped == SumDims.dropped_dimensions(axes, dimensions) : List<&2, U32>}