[SR-14056] Generalize Array: Differentiable conformance to more Collection types

[dabrahams]

Previous ID	SR-14056
Radar	rdar://problem/73268333
Original Reporter	@dabrahams
Type	Improvement

Attachment: Download

Additional Detail from JIRA

Votes	0
Component/s
Labels	Improvement
Assignee	None
Priority	Medium

md5: d37a9a62c43e4bc5a8a7df00a5b517f9

Sub-Tasks:

• SR-14057 Implement Collection.differentiableReduce

blocks:

• TF-1199 ContiguousArray differentiability

is blocked by:

• TF-1149 Cannot differentiate active value with loadable type but address-only tangent type

Issue Description:

We have lots of Array-specific code in ArrayDifferentiation.swift that should be made generic, so that all that's required to make a collection with the right properties differentiable is to write a conditional conformance to `Differentiable`. Then we should add those conditional conformances for other standard library types like `ArraySlice`, `ContiguousArray`, etc.

[rxwei]

@swift-ci create

[philipturner]

Would it be acceptable if I gave Array conformance to AdditiveArithmetic, but used a @_disfavoredOverload for the AdditiveArithmetic.+ operator? This seems necessary to generalize the array-specific differentiation code to arbitrary collections.

Update: I just realized I have a workaround where I don’t have to conform a collection to AdditiveArithmetic. Please disregard this comment if you were notified about it.

[BradLarson]

@philipturner - (Replying here from your comment on SR-15584) An example of adding a conditional Differentiable conformance with Dictionary can be found in the SwiftFusion project , but that probably wouldn't be acceptable for incorporation into the standard library in its present form. Adding a direct conformance to AdditiveArithmetic for various collection types might not be accepted, so this should be reworked into using something similar to `Array.DifferentiableView`.

Beyond that one example, we want to see how generic some of the existing ArrayDifferentiation.swift code can be made, so it becomes much easier to add conditional Differentiable conformances to a number of other collection types.

[rxwei]

+1 on what Brad said. I want to point out that the DifferentiableView type itself is also very unfortunate, because ideally we'd just call this type `Array.TangentVector` (like `Optional.TangentVector`) and not introduce a new type name. We tried to rewrite `DifferentiableView` as `TangentVector` (rxwei@cd0ca57) but ran into some type checker issues (can't remember what exactly). (BTW please don't build on top of my commit as it's very out of date.)

[philipturner]

Would it be fine if I made the pull request draft without any modifications to the tests - just minimal validation on my behalf that some other collection types function correctly? I think adding tests might be loading on more content than is necessary for such a massive pull request draft.

Also, should I CC everyone since it's going to be such a massive change to the _Differentiation package? I think all 5 of us might want to be involved in the discussion. Perhaps after @rxwei does a little bit of validation first.

[philipturner]

I just got a pretty serious crash reproducer. I'm hesistant to make a new JIRA report since this morning, I filed a rather mundane report that I probably shouldn't have. I'm linking the files here, and it's under the same conditions as the report I filed this morning (search "[AutoDiff]" on JIRA).

Experimentation4.zip

I got the crash right after copying and pasting the code for _vjpInit(repeating:count:) and changing .base.reduce to .reduce. It seems I have no way to work around the crash, so you will not be able to differentiate RangeReplaceableCollection.init(repeating:count:) for the time being.

[rxwei]

To report a crash, you should attach a crash log (stack trace and any output related to the crash).

[philipturner]

The crash log is inside the zip file. There was no stack trace; I captured all the information that was available and explained how to reproduce it in my comment above.

[rxwei]

Ok. This is actually not a crash, but a compilation error about an internal inconsistency failure. (A crash would definitely show a stack trace.) The error is coming from TBDGen, but the issue is very, very nuanced and I don't have the bandwidth to look into this anytime soon.

To unblock you though, what you can do is taking off all public modifiers (along with any {{@usableFromInline}}, {{@inlinable}}) from the source code. That will get around this error and enable you to progress further to validate your prototype.

[philipturner]

I did find a way to continue work on my prototype, by adding the -Xfrontend -validate-tbd-against-ir=none flag to the compiler.

I'm going to pause work on MutableCollection autodiff for now. Just like tgmath cross-import overlays, standing issues with the Swift compiler negatively affect how I'm implementing this. This makes it unsuitable for testing and makes workarounds a waste of time. Just like building S4TF, it's better to come back to once a whole slew of compiler bugs are fixed.

I'm refraining from logging anything on JIRA until I've isolated the 3 bugs I found, confirmed they aren't duplicates, and made proper compiler crasher tests for them. It will be helpful to have more complicated edge-case crash reproducers because that's something we couldn't do with S4TF.

@BradLarson could I get an estimate on how much you'll be able to contribute to fixing autodiff bugs in the near future?

[philipturner]

The current prototype and its isolated blocking crashes/errors are published at swift-differentiable-collection-experiments. I don't have experience making tests for the Swift compiler, and one of them - the most problematic - is caused exclusively by RequirementMachine. I'm still on the January 9 nightly toolchain, and it might have been fixed since then.

I'll try messing around with what's on the main Swift branch currently, and see if each crash/error still exists. For the second one (what we discussed a few comments ago), is that type of error possible to catch with Swift compiler tests?

Please note that the questions on the prototype are preliminary, and not yet ready for discussion. I intend for them to be resolved on a future PR draft where each is answered via a code change suggestion.

[BradLarson]

The TBDGen issue looks similar to something we saw here, although with a different cause. In addition to the previous suggestions, you might also be able to work around that by manually defining custom empty JVPs to go alongside your custom VJPs in the places the compiler is throwing this assertion.

For your crash_2 case, I think that's another manifestation of SR-15666. If I'm interpreting that right, I think these are cases where the function really isn't differentiable, but a diagnostic to that effect is missing and thus the assertion failure. In the cases where I hit this, the code wouldn't work to begin with, so it hadn't been a high priority to fix. I just filed that reproducer because I had it on hand.

crash_3 looks like it's failing due to the circular references. RequirementMachine should probably do something other than provide an assertion failure in that case, but I can see why it has trouble with that. I don't know that that's entirely RequirementMachine's fault, but you could file that as a single-file reproducer for an assertion failure.

I unfortunately am on the road for the next week and won't be able to look at anything further until I'm back. After that, I should be able to start catching up on relevant issues.