Skip to content

Derive ElementaryFunctions conformances for structs. #25500

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Jun 15, 2019

Conversation

dan-zheng
Copy link
Contributor

@dan-zheng dan-zheng commented Jun 15, 2019

ElementaryFunctions derived conformances enable elementary math functions to work with product spaces formed from ElementaryFunctions-conforming types.

Enables efficient, elegant mathematical optimizers.
Example:

public class Adam<Model: Differentiable>
where
    Model.TangentVector: VectorProtocol & ElementaryFunctions,
    Model.TangentVector.VectorSpaceScalar: BinaryFloatingPoint & ElementaryFunctions
{
    public typealias Scalar = Model.TangentVector.VectorSpaceScalar
    public var learningRate: Scalar
    public var beta1: Scalar
    public var beta2: Scalar
    public var epsilon: Scalar
    public var decay: Scalar
    public var step: Int = 0
    public var firstMoments: Model.TangentVector = .zero
    public var secondMoments: Model.TangentVector = .zero

    public init(
        for model: __shared Model,
        learningRate: Scalar = 1e-3,
        beta1: Scalar = 0.9,
        beta2: Scalar = 0.999,
        epsilon: Scalar = 1e-8,
        decay: Scalar = 0
    ) {
        precondition(learningRate >= 0, "Learning rate must be non-negative")
        precondition(0 <= beta1 && beta1 <= 1, "Beta parameter must be between 0 and 1")
        precondition(0 <= beta2 && beta2 <= 1, "Beta parameter must be between 0 and 1")
        precondition(decay >= 0, "Learning rate decay must be non-negative")

        self.learningRate = learningRate
        self.beta1 = beta1
        self.beta2 = beta2
        self.epsilon = epsilon
        self.decay = decay
    }

    public func update(_ model: inout Model, along direction: Model.TangentVector) {
        step += 1
        let learningRate = self.learningRate / (1 + decay * Scalar(step))
        // Note: `stepSize` is split into two lines to avoid the "compiler is unable to type-check
        // this expression in reasonable time" error.
        var stepSize = learningRate * Scalar.sqrt(1 - Scalar.pow(beta2, step))
        stepSize = stepSize / (1 - Scalar.pow(beta1, step))
        firstMoments = firstMoments * beta1 + (1 - beta1) * direction
        secondMoments = secondMoments * beta2 + (1 - beta2) * Model.TangentVector.pow(direction, 2)
        model.move(along: -stepSize * firstMoments / (Model.TangentVector.sqrt(secondMoments) + epsilon))
    }
}

Resolves TF-578.
Enables revamping optimizers in tensorflow/swift-apis#218.

`ElementaryFunctions` derived conformances enable elementary math
functions to work with product spaces formed from
`ElementaryFunctions`-conforming types.

Enables efficient, elegant mathematical optimizers.
…sible.

Conform `Differentiable` synthesized associated types to `ElementaryFunctions`
if possible. Similar to existing logic for `VectorProtocol`.
@dan-zheng dan-zheng added the tensorflow This is for "tensorflow" branch PRs. label Jun 15, 2019
@dan-zheng dan-zheng requested a review from rxwei June 15, 2019 07:52
@rxwei rxwei requested a review from stephentyrone June 15, 2019 08:05
@@ -625,6 +634,10 @@ getOrSynthesizeSingleAssociatedStruct(DerivedConformance &derived,
None))
inherited.push_back(kpIterableType);
}
// If all members conform to `ElementaryFunctions`, make the associated struct
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This logic (conform synthesized Differentiable associated types to ElementaryFunctions if possible) is not super principled. Similar logic exists for conforming synthesized types to VectorProtocol if possible.

We could remove this logic by constraining TangentVector to VectorProtocol and/or ElementaryFunctions, but it's not clear to me whether that's desirable. Constraining TangentVector to VectorProtocol seems more reasonable.

@@ -203,6 +210,14 @@ ValueDecl *DerivedConformance::getDerivableRequirement(TypeChecker &tc,

// Retrieve the requirement.
auto results = proto->lookupDirect(name);
// SWIFT_ENABLE_TENSORFLOW
// Filter requirements, if `filter` function is specified.
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note: adding a filter function here is necessary because, for the first time, derived conformances must handle two protocol requirements with the same name: pow(_ x: Self, _ y: Self) -> Self and pow(_ x: Self, _ n: Int) -> Self.

@dan-zheng
Copy link
Contributor Author

Verifying whether tests pass.
@swift-ci Please test tensorflow

@dan-zheng dan-zheng force-pushed the derive-elementaryfunctions branch from 2735b55 to fd11642 Compare June 15, 2019 20:13
@dan-zheng
Copy link
Contributor Author

@swift-ci Please test tensorflow

@dan-zheng
Copy link
Contributor Author

https://ci-external.swift.org is down, locally verified that tests pass.
AdditiveArithmetic and ElementaryFunctions derived conformances will be pitched soon via Swift Evolution.

@dan-zheng dan-zheng merged commit 95b85a9 into swiftlang:tensorflow Jun 15, 2019
@dan-zheng dan-zheng deleted the derive-elementaryfunctions branch June 15, 2019 20:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
tensorflow This is for "tensorflow" branch PRs.
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants