[AutoDiff] Use @derivative for derivative registration.

[dan-zheng]

Rewrite all derivative registration using @differentiable(jvp:vjp:) with
@derivative. This has no functional impact.

Keep original @differentiable attributes so that derivative functions are
publicly exposed.

When retroactive derivative registration is complete:

• @differentiable(jvp:vjp:) will be deprecated.
• @derivative attribute will be the canonical way to register derivatives.

Resolves TF-1076.

---

No remaining uses of @differentiable(vjp:) for derivative registration:

$ grep -nr "vjp:" *
# No remaining uses.

---

Example:

// Before: `@differentiable(vjp:)`.
extension Tensor where Scalar: Numeric {
    @differentiable(vjp: _vjpAdd where Scalar: TensorFlowFloatingPoint)
    public static func + (lhs: Tensor, rhs: Tensor) -> Tensor { ... }
}

extension Tensor where Scalar: TensorFlowFloatingPoint {
    @inlinable
    static func _vjpAdd(lhs: Tensor, rhs: Tensor) -> (
        value: Tensor, pullback: (Tensor) -> (Tensor, Tensor)
    ) { ... }
}

// After: `@derivative`.
extension Tensor where Scalar: Numeric {
    @differentiable(where Scalar: TensorFlowFloatingPoint)
    public static func + (lhs: Tensor, rhs: Tensor) -> Tensor { ... }
}

extension Tensor where Scalar: TensorFlowFloatingPoint {
    @inlinable
    @derivative(of: +)
    static func _vjpAdd(lhs: Tensor, rhs: Tensor) -> (
        value: Tensor, pullback: (Tensor) -> (Tensor, Tensor)
    ) { ... }
}

[saeta]

Wow, really nice work. Just one nit about the removed test, but otherwise LGTM. (Eagerly approving to not block you over the weekend.)

[burmako]

@dan-zheng Perhaps it would make sense to rename _vjpXXX functions to _derivativeXXX?

[dan-zheng]

I don't feel strongly! The best name is func _ (a private anonymous function), which is used in the differentiable programming manifesto but not yet implemented.

When linear functions and transposition are done, users will register only primitive JVP functions returning a @differentiable(linear) differential function, and primitive VJP functions will be removed. At that point, @derivative functions will unambiguously be JVP functions, so derivativeXXX names become appropriate and vjpXXX names become obsolete.