code: def test_tf_heterogeneous_cluster_distribution_config(
      sagemaker_session,
      tensorflow_training_version,
      tensorflow_training_py_version,
  ):
      if version.Version(tensorflow_training_version) < version.Version(
          "2.0"
      ):
          pytest.skip("This test is for TF 2.0 and higher.")
  
      training_group = InstanceGroup("train_group", "ml.c4.xlarge", 1)
      expected_return = {
          "mpi": {"enabled": True},
          "instance_groups": ["train_group"],
      }
      tf = _build_tf(
          sagemaker_session,
          framework_version=tensorflow_training_version,
          py_version=tensorflow_training_py_version,
          instance_groups=[training_group],
          distribution={
              "mpi": {"enabled": True},
              "instance_groups": [training_group],
          },
      )
      assert tf.distribution == expected_return
  
imports: from __future__ import absolute_import
  import json
  import logging
  import os
  from mock import patch, Mock, MagicMock
  from packaging import version
  import pytest
  from sagemaker.estimator import _TrainingJob
  from sagemaker.session_settings import SessionSettings
  from sagemaker.tensorflow import TensorFlow
  from sagemaker.instance_group import InstanceGroup
  from sagemaker.workflow.parameters import ParameterString, ParameterBoolean
  from tests.unit import DATA_DIR
label_desc: Use image node_id 18 for calling a global function with example usage: test_tf_heterogeneous_cluster_distribution_config(sagemaker_session, tensorflow_training_version, tensorflow_training_py_version) without return types