[AutoDiff] populate diff witnesses during differentiation #28402
Conversation
|
@swift-ci please test tensorflow |
| // Get lowered argument indices. | ||
| auto *paramIndices = A->getParameterIndices(); | ||
| assert(paramIndices && "Parameter indices should have been resolved"); | ||
| auto *loweredParamIndices = autodiff::getLoweredParameterIndices( | ||
| paramIndices, decl->getInterfaceType()->castTo<AnyFunctionType>()); | ||
| // NOTE(TF-893): Extending capacity is necessary when `origSilFnType` has |
There was a problem hiding this comment.
I wonder if the only case that triggers this "capacity extending due to captured variables" logic is the @differentiable attribute on func original nested in func differentiableFunction(from:)?
Full context: there are known issues regarding differentiation and local variable capture (TF-881). @rxwei mentioned disallowing @differentiable attribute on nested functions for now and creating a builtin to support func differentiableFunction(from:). One known user of differentiableFunction(from:) is the custom differentiation tutorial.
There was a problem hiding this comment.
Yes, that's the only thing (in the stdlib + swift-apis + tests, at least) that triggers this problem.
It would indeed be nice to remove that and forbid @differentiable on nested functions for now.
| /// Sets the differentiability witness JVP and VJP to the JVP and VJP in `attr`. | ||
| /// | ||
| /// `attr` must have a JVP and VJP. | ||
| static void fillDifferentiabilityWitness(SILModule &module, |
There was a problem hiding this comment.
Minor: how about settling on the verb "canonicalize" instead of "fill"/"populate"/"process" for [differentiable] attributes, differentiability witnesses, differentiable_function instructions?
processDifferentiableAttribute and processDifferentiableFunctionInst can also be renamed later.
There was a problem hiding this comment.
Sounds good. Done for this one.
The main point of this PR is to populate differentiability witnesses during the differentiation pass, while still using attributes for everything, so that we can incrementally switch everything over to using witnesses.
A few adjustments in other things were necessary to get this to work:
I'm going to run a toolchain build and swift-apis tests on this before merging, because I'm not super confident that all the assertions I added will always pass.