extension Collection
where
  Self: Differentiable, Element: Differentiable,
  TangentVector: MutableCollection,
  TangentVector.Element == Element.TangentVector
{
  @differentiable(wrt: (self, initialResult))
  public func differentiableReduce<Result: Differentiable>(
    _ initialResult: Result,
    _ nextPartialResult: @differentiable (Result, Element) -> Result
  ) -> Result {
    reduce(initialResult, nextPartialResult)
  }

  @inlinable
  @derivative(of: differentiableReduce)
  internal func _vjpDifferentiableReduce<Result: Differentiable>(
    _ initialResult: Result,
    _ nextPartialResult: @differentiable (Result, Element) -> Result
  ) -> (
    value: Result,
    pullback: (Result.TangentVector)
      -> (TangentVector, Result.TangentVector)
  ) {
    var pullbacks:
      [(Result.TangentVector) -> (Result.TangentVector, Element.TangentVector)] =
        []
    let count = self.count
    pullbacks.reserveCapacity(count)
    var result = initialResult
    for element in self {
      let (y, pb) =
        valueWithPullback(at: result, element, in: nextPartialResult)
      result = y
      pullbacks.append(pb)
    }
    return (
      value: result,
      pullback: { [selfTan = zeroTangentVector] tangent in
        var selfTangent = selfTan
        var resultTangent = tangent
        var index = selfTangent.startIndex
        for pullback in pullbacks.reversed() {
          let (newResultTangent, elementTangent) = pullback(resultTangent)
          resultTangent = newResultTangent
          selfTangent[index] += elementTangent
          index = selfTangent.index(after: index)
        }
        return (selfTangent, resultTangent)
      }
    )
  }
}

// Example:
@differentiable
func foo(_ array: [Float]) -> Float {
  return array.differentiableReduce(0, +)
}
print(gradient(at: [1, 2, 3, 4], in: foo))

// Doesn't currently work because `Array.TangentVector` is not a `MutableCollection`.