code: def _run_xla_precision(fabric, expected_dtype):
      with fabric.init_module():
          model = BoringPrecisionModule(expected_dtype)
  
      optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
      model, optimizer = fabric.setup(model, optimizer)
  
      batch = torch.rand(2, 32, device=fabric.device)
  
      # TODO: This should be float16/bfloat16
      assert (
          model.layer.weight.dtype
          == model.layer.bias.dtype
          == torch.float32
      )
  
      assert batch.dtype == torch.float32
      output = model(batch)
      assert output.dtype == torch.float32
      loss = torch.nn.functional.mse_loss(
          output, torch.ones_like(output)
      )
      fabric.backward(loss)
      assert model.layer.weight.grad.dtype == torch.float32
      optimizer.step()
      optimizer.zero_grad()
  
file_path: pytorch-lightning/tests/tests_fabric/plugins/precision/test_xla_integration.py
imports: import os
  from unittest import mock
  import pytest
  import torch
  import torch.nn
  from lightning.fabric import Fabric
  from lightning.fabric.plugins import XLAPrecision
  from tests_fabric.helpers.runif import RunIf
label_desc: Use image node_id 1 for calling a global function with example usage: _run_xla_precision(fabric, expected_dtype) without return types