Improved derivative performance for broadcasted operations.

[eaplatanios]

Re-implementation of swiftlang/swift#24408.

@rxwei the reason you were getting errors is because you were differentiating a tensor-valued function (as opposed to scalar-valued) and you were assuming that the gradient seed has the same shape as the operation result. I fixed that by broadcasting the gradient seeds before computing the gradients.

[rxwei]

the reason you were getting errors is because you were differentiating a tensor-valued function (as opposed to scalar-valued) and you were assuming that the gradient seed has the same shape as the operation result. I fixed that by broadcasting the gradient seeds before computing the gradients.

We have been assuming that seeds are always shape-compatible (if not same-shape) as the original result, and existing implementations have not been doing broadcasting to the shape of the original result. I think there are other reasons for it, which may lead to a different solution.

[eaplatanios]

That's interesting! I actually traced the test failure and saw that the seed shape was not matching the result shape when the add VJP was being called. Given that this was being directly invoked by the gradient function, could there be a bug somewhere more core to AD support?

[rxwei]

In this test case, the use of gradient(at:in:) is not valid because gradient is only mathematically defined for functions that return a scalar. We can either make foo(_:_:) to a sum() or use pullback(at:in:).

[eaplatanios]

That's what I thought. This was actually the failing case in the other PR and that's why I copied it here but wasn't sure of the semantics of gradient. I'll add a call to sum() and remove the seed broadcasts then. :)

[rxwei]

Given that this was being directly invoked by the gradient function, could there be a bug somewhere more core to AD support?

AD is fully shape-agnostic, and there's some mathematical consistency to differential operators like gradient(at:in:) in that it should fail when the function's result is not semantically a scalar. I think a proper solution is to change the definition of gradient(at:in:), gradient(at:_:in:), Tensor.gradient(in:), valueWithGradient(at:in:), valueWithGradient(at:_:in:)and Tensor.valueWithGradient(in:) (all defined in DifferentialOperators.swift) to call precondition(_:_:file:line:) on the original function's result, for example:

@inlinable
public func valueWithGradient<T, R>(
    at x: T,
    in f: @differentiable (T) -> Tensor<R>
) -> (value: Tensor<R>, gradient: T.TangentVector)
    where T: Differentiable, R: TensorFlowFloatingPoint {
    let (y, pullback) = valueWithPullback(at: x, in: f)
    precondition(y.rank == 0)
    return (y, pullback(Tensor<R>(1)))
}

[eaplatanios]

I also prefer that. What's the cost of preconditions? Are they removed when compiling with optimizations enabled?

[rxwei]

I also prefer that. What's the cost of preconditions? Are they removed when compiling with optimizations enabled?

	-Onone	-O	-Ounchecked
assert()	YES	NO	NO
assertionFailure()	YES	NO	NO**
precondition()	YES	YES	NO
preconditionFailure()	YES	YES	YES**
fatalError()*	YES	YES	YES

From this blog post.

[eaplatanios]

Awesome, thanks! :) I'll be away for ~30' but will make those changes once I get back.

[eaplatanios]

I'm looking into this now. I can add the precondition for valueWithGradient, but not for gradient as I don't think it ever evaluates f directly.

[eaplatanios]

I made all necessary changes and all tests pass locally.

[rxwei]

I'm looking into this now. I can add the precondition for valueWithGradient, but not for gradient as I don't think it ever evaluates f directly.

It actually does because it calls pullback(at:in:) which calls valueWithPullback(at:in:). You can change gradient(at:in:) to call valueWithGradient(at:in:) and drop the value after precondition. Don't worry about performance since we can optimize those things away later.

[eaplatanios]

Sounds good. Done in the latest commit. :)

[rxwei]

All tests pass. Thank you!