Utilities used by other modules

xla_imported[source]

xla_imported()

Check whether the torch_xla module has been successfully imported

xla_imported is a utility method that is used to check if the torch_xla module has been successfully imported.

print_aten_ops[source]

print_aten_ops()

print out xla aten operations (from xla debug metrics report torch_xla.debug.metrics)

One of the problems we have hit testing different models and transforms is that sometimes it is slower on TPUs compared to running on CPUs, but this happens because we hit operations on Pytorch XLA that are only handled by the CPU and not by the accelerator.

print_aten_ops calls directly some pytorch metrics which outputs to stdout, so the only way to get that info is capture it.

#test that torch_xla has been imported on colab
assert xla_imported()