law sum_dims_keepdim

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

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
import bendtorch-core@0.1.0.0/src/tensor/elementwise_totality.bend as Totality
import bendtorch-core@0.1.0.0/src/autograd/reduce.bend as Reduce
import bendtorch-core@0.1.0.0/src/autograd/sum_dims.bend as SumDims

154# With keepdim, the sum is exactly reduce_to onto the kept shape.
155law sum_dims_keepdim:
156  for input: Core.Dynamic
157  for +axes: List<&2, U32>
158  for +kept: Shape.Shape
159  for +dropped: List<&2, U32>
160  for planned: {SumDims.plan(Totality.dynamic_shape(input), axes) == Done{SumDims.Plan{kept, dropped}} : Result<&2, &2, SumDims.Error, SumDims.Plan>}
161  {SumDims.sum_dims(input, axes, True{}) == as_sum(Reduce.reduce_to(input, kept)) : Core.Dynamic & Result<&2, &1, SumDims.Error, Core.Dynamic>}