Fix TPU tests for metrics_impl branch (#1277)

* update accuracy to accumulate _num_correct in a tensor on the right device

* update loss metric to accumulate _sum in a tensor on the right device

* update mae metric to accumulate in a tensor on the right device

* update mpd metric to accumulate in a tensor on the right device

* update mse metric to accumulate in a tensor on the right device

* update top k accuracy  metric to accumulate in a tensor on the right device

* update precision and recall metrics to accumulate in tensors on the right device

* .....

* black formatting

* reverted run*.sh

* change all metrics default device to cpu except running_average

* Update ignite/metrics/precision.py

Co-authored-by: vfdev <[email protected]>

* remove Optional type from metric devices since default is cpu

* add comment explaining lack of detach in accuracy metrics

* update docstrings and docs

* Update ignite/metrics/accumulation.py

Co-authored-by: vfdev <[email protected]>

* Update ignite/metrics/accumulation.py

Co-authored-by: vfdev <[email protected]>

* Update ignite/metrics/accumulation.py

Co-authored-by: vfdev <[email protected]>

* Update ignite/metrics/accuracy.py

Co-authored-by: vfdev <[email protected]>

* Update ignite/metrics/fbeta.py

Co-authored-by: vfdev <[email protected]>

* Update ignite/metrics/loss.py

Co-authored-by: vfdev <[email protected]>

* Update ignite/metrics/metric.py

Co-authored-by: vfdev <[email protected]>

* Update ignite/metrics/precision.py

Co-authored-by: vfdev <[email protected]>

* Update ignite/metrics/recall.py

Co-authored-by: vfdev <[email protected]>

* add comment explaining lack of detach in metrics docs

* support device argument for running_average

* update support for device argumenet for accumulation

* fix and improve device tests for metrics

* fix and improve device tests for metrics

* fix TPU tests

* Apply suggestions from code review

* Apply suggestions from code review

* detach tensors earlier in update

* remove redundant to() call

* ensure metrics aren't created on XLA devices

* Fixed isort

* move xla check to Metric.__init__ instead of individual metrics

* update xla tests

* replace deleted callable check

* remove redundant precision and recall __init__

* replace precision/recall __init__ for docs rendering

* add support for metrics_lambda with components on diff devices

* fix epoch_metric xla test

Co-authored-by: vfdev <[email protected]>
Co-authored-by: n2cholas <[email protected]>

tests/ignite/metrics/test_epoch_metric.py

@@ -161,7 +161,7 @@ def compute_fn(y_preds, y_targets):
 def _test_distrib_integration(device=None):
 
     if device is None:
-        device = idist.device()
+        device = idist.device() if idist.device().type != "xla" else "cpu"
 
     rank = idist.get_rank()
     torch.manual_seed(12)