def simd_mesh_setup(params, mesh_shape, layout_rules):
    """Constructs SimdMesh function - instructions on how to evenly split tensors across all TPU cores"""

    num_hosts = params["context"].num_hosts
    host_placement_fn = params["context"].tpu_host_placement_function
    device_list = [host_placement_fn(host_id=i) for i in range(num_hosts)]
    tf.logging.info(f"device_list = {device_list}")

    # TODO: Better estimation of replica cache size?
    replica_cache_size = 300 * 1000000  # 300M per replica

    # Worker 0 caches all the TPU binaries
    worker0_mem = replica_cache_size * params["context"].num_replicas
    devices_memory_usage = [worker0_mem] + [0] * (num_hosts - 1)
    var_placer = mtf.utils.BalancedVariablePlacer(device_list, devices_memory_usage)
    mesh_devices = [""] * mesh_shape.size
    mesh_impl = mtf.simd_mesh_impl.SimdMeshImpl(
        mesh_shape, layout_rules, mesh_devices, params["context"].device_assignment)

    return var_placer, mesh_impl
Constructs SimdMesh function - instructions on how to evenly split tensors across all TPU cores
utils.py