code: def test_pytorchxla_distribution(
      time,
      name_from_base,
      sagemaker_session,
      pytorch_training_compiler_version,
      instance_class,
      pytorch_training_compiler_py_version,
  ):
      if Version(pytorch_training_compiler_version) < Version("1.12"):
          pytest.skip(
              "This test is intended for PyTorch 1.12 and above"
          )
      compiler_config = TrainingCompilerConfig()
      instance_type = f"ml.{instance_class}.xlarge"
  
      pt = PyTorch(
          py_version=pytorch_training_compiler_py_version,
          entry_point=SCRIPT_PATH,
          role=ROLE,
          sagemaker_session=sagemaker_session,
          instance_count=2,
          instance_type=instance_type,
          framework_version=pytorch_training_compiler_version,
          enable_sagemaker_metrics=False,
          compiler_config=TrainingCompilerConfig(),
          distribution={"pytorchxla": {"enabled": True}},
      )
  
      inputs = "s3://mybucket/train"
  
      pt.fit(inputs=inputs, experiment_config=EXPERIMENT_CONFIG)
  
      sagemaker_call_names = [
          c[0] for c in sagemaker_session.method_calls
      ]
      assert sagemaker_call_names == ["train", "logs_for_job"]
      boto_call_names = [
          c[0] for c in sagemaker_session.boto_session.method_calls
      ]
      assert boto_call_names == ["resource"]
  
      expected_train_args = _create_train_job(
          pytorch_training_compiler_version,
          instance_type,
          compiler_config,
          pytorch_training_compiler_py_version,
          instance_count=2,
      )
      expected_train_args["input_config"][0]["DataSource"][
          "S3DataSource"
      ]["S3Uri"] = inputs
      expected_train_args["enable_sagemaker_metrics"] = False
      expected_train_args["hyperparameters"][
          TrainingCompilerConfig.HP_ENABLE_COMPILER
      ] = json.dumps(True)
      expected_train_args["hyperparameters"][
          PyTorch.LAUNCH_PT_XLA_ENV_NAME
      ] = json.dumps(True)
      expected_train_args["hyperparameters"][
          TrainingCompilerConfig.HP_ENABLE_DEBUG
      ] = json.dumps(False)
  
      actual_train_args = sagemaker_session.method_calls[0][2]
      assert (
          actual_train_args == expected_train_args
      ), f"{json.dumps(actual_train_args, indent=2)} != {json.dumps(expected_train_args, indent=2)}"
  
imports: from __future__ import absolute_import
  import logging
  import json
  import os
  import pytest
  from mock import MagicMock, Mock, patch, ANY
  from packaging.version import Version
  from sagemaker import image_uris
  from sagemaker.pytorch import PyTorch, TrainingCompilerConfig
  from sagemaker.pytorch.model import PyTorchModel
  from sagemaker.instance_group import InstanceGroup
  from sagemaker.session_settings import SessionSettings
  from tests.unit.sagemaker.training_compiler import EC2_GPU_INSTANCE_CLASSES
label_desc: Use image node_id 12 for calling a global function with example usage: test_pytorchxla_distribution(time, name_from_base, sagemaker_session, pytorch_training_compiler_version, instance_class, pytorch_training_compiler_py_version) without return types