Hm, I'm finding it a bit hard to understand what's going on in these tests.

testing suggestion: throughout these tests, call into a utility that produces a non-contiguous tensor from a given tensor. You can then run the contiguous and non-contiguous inputs through the op under test and verify output equivalence. wdyt about this?

My idea here is to compare nested tensor op with padded op:

• create a nested tensor, mess up its contiguity, then unbind
• pad the nested tensor, do same reshaping and transposing, then (kind of) unbind

Correct me if I'm wrong: are you suggesting comparing 2 nested tensors who appears to have same entries but in fact different memory layout?

Correct me if I'm wrong: are you suggesting comparing 2 nested tensors who appears to have same entries but in fact different memory layout?

Right, we'd expect that if we run the same operation on each of these, their outputs would also appear to have the same entries (and possibly different memory layout, but this depends on the op). Existing tests cover nested contiguous vs. padded, and here the focus would be on nested contiguous vs. nested non-contiguous. wdyt, are we losing any coverage this way?

I kinda doubt this is related to the inference_mode() issue, but this looks a bit suspicious. Do we want NestedTensor::reshape to be able to work with autograd? If we do, some quick comments:

The way that normal reshape() works today is that it sometimes returns a of the input, and sometimes doesn't depending on the arguments to reshape. We can't actually support an autograd formula for that directly, so what happens instead of that reshape() decomposes into a view op or a clone, and autograd handles those individual ops instead. Does the nested tensor implementation of reshape() also "sometimes return a view"?

@bdhirsh Yes that's correct, AFAICT the NT implementation of reshape uses the same underlying logic to determine when to return a view and when to return a copy. Do you see it being possible to define a NT-specific reshape_backward that deals with this properly?

Note that we need a NT-specific version of reshape because the semantics of -1s are slightly different due to the possibility for ragged dims.

Nested tensor implementation can return a view but I guess that's not the "view" autograd engine uses:

• I heard that autograd would handle "view" differently as "fast path"
• I also heard that this "fast path" is disabled for nested tensor: it would record nested tensor view just as another function
• So based on sharing memory it can return a view, but based on "fast path" it's not an autograd view I think

Could that trigger some edge case?

@bdhirsh Yes that's correct, AFAICT the NT implementation of reshape uses the same underlying logic to determine when to return a view and when to return a copy. Do you see it being possible to define a NT-specific reshape_backward that deals with this properly?

Note that we need a NT-specific version of reshape because the semantics of -1s are slightly different due to the possibility for ragged dims.

Yes the way I implemented reshape_backward is just like Driss's way for linear. Although I can also do the ugly "is_nested()" version to maintain compositeimplicit-ness

Yes the way I implemented reshape_backward is just like Driss's way for linear.

This is good, but I'm curious how this backward deals with the forward sometimes returning a view and sometimes returning a copy?

Yes the way I implemented reshape_backward is just like Driss's way for linear.

This is good, but I'm curious how this backward deals with the forward sometimes returning a view and sometimes returning a copy?

Basically no special treatment: simply reshape grad as input

Why is this schema different from the reshape above?
This function is also return sometimes a view and sometimes not a view.

So this schema is wrong (and autograd is going to generate silently wrong code for it!)

It is NOT possible to register autograd for a function that has such a property so you will need to keep reshape CompositImplicitAutograd.
Since the computeStride function there will not work with nested, you can simply create a new native function to abstract it away: You can make it a native function that takes self as input directly and you will be able to provide a custom Nested implementation for it.

btw this wrong schema is also most likely the reason for the inference mode issue you had as the autograd kernel is not properly registered!

I'm not sure to understand? Why can't you just override the explicit kernels you don't want?

This is Driss's work that we thought would be a necessary fix here. The original discussion is here

@drisspg @bdhirsh would know better than I do

This was added in order to allow for composite implicit kernels to work with nested_tensors. Some functions like softmax are composite implicit functions that in turn call non composite ops like _softmax . In order to be able to register at _softmax the nested tensor needs to be able to run by composite function. Similiarly Reshape as an example of one of these ops whose top most function is composite implicit that quickly calls a non-composite op

I understand that you want it to get the Implicit kernels, but this comment seem to indicate that you do special logic (compared to the other classical backend keys) to prevent it from getting the Explicit kernels. Did I read this comment wrong?

Yeah that is correct, work for implicit don't work for explicit. I went back and fourth about this, practically it appears that there are alot more composite explicit keys that end up calling sizes and producing strange hard to debug errors for users. This possibility exists for composite implicit as well, however was happening with less frequency for those kernels.

Is your suggestion to enable for both composite kernels types and whenever we find a composite explicit doing something wrong just register a Torch_check(false) kernel?

Is your suggestion to enable for both composite kernels types and whenever we find a composite explicit doing something wrong just register a Torch_check(false) kernel?

That would have been my first guess yes.
But I am not sure:

• how involved is the current "workaround" to get Implicit but not Explicit kernels? cc @bdhirsh
• What happens when you want to use some of the Explicit kernels? You now have to copy/paste them?

Doing such a check is usually a code smell. You use the dispatcher to call reshape so if you had the right dispatch for it, you should have ended up in your implementation directly.