xla

43116740 - Add unsharded computation test for virtual device (#4215)

Add unsharded computation test for virtual device (#4215) * Add more tests for virtual device * Remove duplicate log line Co-authored-by: root <root@t1v-n-bc4358d0-w-0.us-central2-b.c.tpu-pytorch.internal>