cc @JackCaoG @miladm

Note: this PR can, possibly, create more partitions than before, affecting performance. Maybe we should make sure there are no regressions before actually landing it.

I think the definition of the fallback is that we have to execute this op on fallback devices(uusually) cpu. for the cases you mentioned

1. nodes which resulted in a non-XLA tensor weren't flagged as fallback
   Does the operation being executed on cpu or gpu/tpu?

2. nodes whose arguments were a container with non-XLA tensors weren't flagged as fallback (e.g. stack)
   I guess this one is technically a fallback, since op will be executed on cpu. It will show up in dynamo fallback messages, but shouldn't be in the XLA's metrics aten counter I think?

I don't fully understand what 3 means, can you give ma an example?

The problem was that FallBackNodeCollector was doing more than just collecting operations that were executed on CPU. It was also flagging nodes, that had non-XLA tensors arguments, as fallback (even though they were executed on XLA).

Reason: my guess is that it did so, in order to guarantee that the input/output of the generated partitions were all XLA tensors (which I believe extract_internal assumes).

Assuming that's the case, it still missed the 2 cases I mentioned:

• nodes which result in non-XLA tensor: might end up as output of some partition
• nodes whose arguments were a container with non-XLA tensors: arguments might end up as the input of some partition

Solution:

• disambiguate fallback from unsupported nodes (from the perspective of the partitioner)
• add the missing 2 cases

@JackCaoG let me know what you think.

yea that make sense, is this pr ready for review?

Yes, it is.

@JackCaoG could you approve this PR?