In-place argument updates with device affinities fails

[sogartar]

What happened?

This program in-place-sharded-tensor.zip fails to do a proper in-place update.

The program is equivalent to

arg0 += 1
arg1 += 1

What I observe is that the 2 tensor arguments don't change.
It is exported from the test ShardedTensorTest.testInPlaceUpdate from this PR.

Steps to reproduce your issue

1. Extract the archive and run compile.sh
2. You can't verify the result with iree-run-module as it does not support validating the arguments post execution. You can use the test from the referenced PR for that.

What component(s) does this issue relate to?

Compiler

Version information

You may need this patch. I was using it when testing.

Additional context

No response