Added support for a 'Tensor.gathering(where:)'.

[eaplatanios]

@rxwei How can we make this differentiable? The problem is with the line:

let posAxis = axis < 0 ? axis + rank : axis

but that should not affect differentiability of the function as a whole.

In general, I feel AD should be capable of checking whether certain statements depend on the variables wrt to which we're differentiating and not fail in cases where they do not (like this op). Is that something that's planned? It would also avoid using the thunk hack in many places (e.g., with .map calls on arrays for creating shapes).

[rxwei]

In general, I feel AD should be capable of checking whether certain statements depend on the variables wrt to which we're differentiating and not fail in cases where they do not (like this op). Is that something that's planned? It would also avoid using the thunk hack in many places (e.g., with .map calls on arrays for creating shapes).

Yes, AD has been doing this all along, powered by activity analysis. However, the issue with the ternary operator expression is that AD does not yet support functions with any control flow, even when the control flow is not part of the code that should be differentiated. A ternary operator expression is a control flow construct in Swift.

A workaround is to wrap the control flow in a thunk and call it directly to get a value.

let posAxis = { axis < 0 ? axis + rank : axis }()

[eaplatanios]

I see. That makes sense. Although I believe that the example in #157 does not involve control flow.

Also, the thunk approach does not work in this case. This is what I tried previously. I get this:

note: cannot differentiate through a non-differentiable result; do you want to add '.withoutDerivative()'?
        let posAxis = {axis < 0 ? axis + rank : axis}()
                      ^

[rxwei]

This is because the thunk captures self, which makes the activity analysis think that the thunk should be differentiated with respect to the captured variable, which we don't support yet. We can create a local variable rank for now.

let rank = rank
let posAxis = { axis < 0 ? axis + rank : axis }()

[eaplatanios]

Oh I see. That makes a lot more sense now. I updated this PR.

[rxwei]

Doc comments for non-mutating methods often starts with Returns a ... by <verb>ing ... instead of <Verb>s .... Quite a few doc comments in the library do not follow this guideline yet and should be updated.

[eaplatanios]

Sounds good. I'll try to go through them tomorrow.

[rxwei]

Filed #160 .

[rxwei]

All tests passed.