{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "/home/modelzoo/jax/jax-cfd/notebooks\n",
            "/home/modelzoo/jax/jax-cfd/notebooks\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.10/site-packages/IPython/core/magics/osm.py:417: UserWarning: This is now an optional IPython functionality, setting dhist requires you to install the `pickleshare` library.\n",
            "  self.shell.db['dhist'] = compress_dhist(dhist)[-100:]\n"
          ]
        }
      ],
      "source": [
        "!source ~/.bashrc\n",
        "%cd /home/modelzoo/jax/jax-cfd/notebooks/\n",
        "!pwd\n",
        "!export HIP_VISIBLE_DEVICES=3"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "WFXyg2qVtPHM"
      },
      "outputs": [],
      "source": [
        "dataset_name = 'kolmogorov_re_1000'  #@param ['kolmogorov_re_1000', 'decaying', 'kolmogorov_re_4000'] {type: \"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "kUtRFyDgzyN9",
        "outputId": "651d3d2e-abda-4cc7-fe14-329048f4dcd6"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "/bin/bash: gsutil: command not found\n",
            "CPU times: user 1.74 ms, sys: 3.17 ms, total: 4.91 ms\n",
            "Wall time: 276 ms\n"
          ]
        }
      ],
      "source": [
        "%time ! gsutil -m cp gs://gresearch/jax-cfd/public_eval_datasets/{dataset_name}/eval_*.nc ../content"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 33,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "AiC9-QA4Deni",
        "outputId": "9f0d6b12-5894-4a76-a3a2-b0318287f4ec"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "/bin/bash: gsutil: command not found\n",
            "CPU times: user 3.34 ms, sys: 10.8 ms, total: 14.2 ms\n",
            "Wall time: 289 ms\n"
          ]
        }
      ],
      "source": [
        "%time ! gsutil -m cp gs://gresearch/jax-cfd/public_models/*.pkl ../content"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "MzEWHd5e0Jwf",
        "outputId": "315b1e4d-47c3-46cb-b67f-1dbc96e28506"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "/home/modelzoo/jax/jax-cfd/notebooks\n",
            "EPD_ckpt.pkl\t\t eval_2048x2048_64x64.nc  eval_64x64_64x64.nc\n",
            "eval_1024x1024_64x64.nc  eval_256x256_64x64.nc\t  LC_ckpt.pkl\n",
            "eval_128x128_64x64.nc\t eval_512x512_64x64.nc\t  LI_ckpt.pkl\n"
          ]
        }
      ],
      "source": [
        "! pwd\n",
        "! ls ../content"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "cellView": "form",
        "id": "wTfEmQPZ1YiT"
      },
      "outputs": [],
      "source": [
        "#@title Imports { form-width: \"30%\" }\n",
        "\n",
        "import warnings\n",
        "warnings.simplefilter('ignore')\n",
        "\n",
        "import os\n",
        "import functools\n",
        "import pickle\n",
        "\n",
        "import gin\n",
        "import jax\n",
        "import jax.numpy as jnp\n",
        "import numpy as np\n",
        "import haiku as hk\n",
        "\n",
        "import xarray\n",
        "import seaborn\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import jax_cfd.base as cfd\n",
        "import jax_cfd.data as cfd_data\n",
        "import jax_cfd.ml as ml\n",
        "\n",
        "model_builder = ml.model_builder\n",
        "model_utils = ml.model_utils\n",
        "optimizer_modules = ml.optimizer_modules\n",
        "physics_specifications = ml.physics_specifications"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "cellView": "form",
        "id": "0aVRpqBzEOzS"
      },
      "outputs": [],
      "source": [
        "#@title Helper functions\n",
        "\n",
        "# shape_structure = lambda tree: jax.tree_util.tree_map(lambda x: x.shape, tree)\n",
        "shape_structure = lambda tree: jax.tree_util.tree_map(lambda x: x.shape, tree)\n",
        "\n",
        "\n",
        "def xarray_open(path):\n",
        "  return xarray.open_dataset(path, chunks={'time': '100MB'})\n",
        "\n",
        "\n",
        "def strip_imports(s):\n",
        "  out_lines = []\n",
        "  for line in s.splitlines():\n",
        "    if not line.startswith('import'):\n",
        "      out_lines.append(line)\n",
        "  return '\\n'.join(out_lines)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ljEfE4rdVQ3o"
      },
      "source": [
        "# Selecting evaluation dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "id": "xiQG5F213nN1"
      },
      "outputs": [],
      "source": [
        "#@title Paths to evaluation datasets\n",
        "\n",
        "base_path = '../content/'\n",
        "\n",
        "kolmogorov_re_1000 = {\n",
        "    f'baseline_{i}x{i}': os.path.join(base_path, f'eval_{i}x{i}_64x64.nc')\n",
        "    for i in [64, 128, 256, 512, 1024, 2048]\n",
        "}\n",
        "decaying = {\n",
        "    f'baseline_{i}x{i}': os.path.join(base_path, f'eval_{i}x{i}_64x64.nc')\n",
        "    for i in [64, 128, 256, 512, 1024, 2048]\n",
        "}\n",
        "kolmogorov_re_4000 = {\n",
        "    f'baseline_{i}x{i}': os.path.join(base_path, f'eval_{i}x{i}_128x128.nc')\n",
        "    for i in [128, 256, 512, 1024, 2048, 4096]\n",
        "}\n",
        "\n",
        "all_datasets = {\n",
        "    'kolmogorov_re_1000': kolmogorov_re_1000,\n",
        "    'decaying': decaying,\n",
        "    'kolmogorov_re_4000': kolmogorov_re_4000,\n",
        "}\n",
        "\n",
        "reference_names = {\n",
        "    'kolmogorov_re_1000': 'baseline_2048x2048',\n",
        "    'decaying': 'baseline_2048x2048',\n",
        "    'kolmogorov_re_4000': 'baseline_4096x4096',\n",
        "}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "id": "I_Mvh9yI3XJn"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "{'baseline_64x64': '../content/eval_64x64_64x64.nc', 'baseline_128x128': '../content/eval_128x128_64x64.nc', 'baseline_256x256': '../content/eval_256x256_64x64.nc', 'baseline_512x512': '../content/eval_512x512_64x64.nc', 'baseline_1024x1024': '../content/eval_1024x1024_64x64.nc', 'baseline_2048x2048': '../content/eval_2048x2048_64x64.nc'}\n"
          ]
        }
      ],
      "source": [
        "#@title Loading evaluation dataset {run: \"auto\"}\n",
        "\n",
        "dataset_paths = all_datasets[dataset_name]\n",
        "print(dataset_paths)\n",
        "datasets = {k: xarray_open(v) for k, v in dataset_paths.items()}\n",
        "reference_ds = datasets[reference_names[dataset_name]]\n",
        "\n",
        "grid = cfd_data.xarray_utils.grid_from_attrs(reference_ds.attrs)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "cellView": "form",
        "id": "BcRUrBGb6v-U"
      },
      "outputs": [],
      "source": [
        "#@title Selecting initial conditions and baseline trajectories.\n",
        "\n",
        "sample_id = 0\n",
        "time_id = 0\n",
        "length = 200  # length of the trajectory.\n",
        "inner_steps = 10  # since we deal with subsampled datasets\n",
        "\n",
        "initial_conditions = tuple(\n",
        "    reference_ds[velocity_name].isel(\n",
        "        sample=slice(sample_id, sample_id + 1),\n",
        "        time=slice(time_id, time_id + 1)\n",
        "    ).values\n",
        "    for velocity_name in cfd_data.xarray_utils.XR_VELOCITY_NAMES[:grid.ndim]\n",
        ")\n",
        "\n",
        "target_ds = reference_ds.isel(\n",
        "    sample=slice(sample_id, sample_id + 1),\n",
        "    time=slice(time_id, time_id + length))\n",
        "\n",
        "\n",
        "datasets = {\n",
        "    k: v.isel(sample=slice(sample_id, sample_id + 1),\n",
        "              time=slice(time_id, time_id + length))\n",
        "    for k, v in datasets.items()\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WASsDE3iVXKs"
      },
      "source": [
        "# Selecting model checkpoint to load"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "id": "--a-KrfBHlp-"
      },
      "outputs": [],
      "source": [
        "class CheckpointState:\n",
        "  \"\"\"Object to package up the state we load and restore.\"\"\"\n",
        "\n",
        "  def __init__(self, **kwargs):\n",
        "    for name, value in kwargs.items():\n",
        "      setattr(self, name, value)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "id": "7NFw1ynrVXKt"
      },
      "outputs": [],
      "source": [
        "checkpoint_paths = {\n",
        "    'LI': \"../content/LI_ckpt.pkl\",\n",
        "    'LC': \"../content/LC_ckpt.pkl\",\n",
        "    'EPD': \"../content/EPD_ckpt.pkl\",\n",
        "}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "id": "aMBjDP-yVXKt"
      },
      "outputs": [],
      "source": [
        "#@title selecting model to evaluate {run: \"auto\"}\n",
        "\n",
        "model_name = \"LI\" #@param ['LI', 'LC', 'EPD',] {type: \"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "id": "3rBkCx9OVXKt"
      },
      "outputs": [],
      "source": [
        "#@title Loading the checkpoint\n",
        "\n",
        "ckpt_path = checkpoint_paths[model_name]\n",
        "with open(ckpt_path, 'rb') as f:\n",
        "  ckpt = pickle.load(f)\n",
        "params = ckpt.eval_params"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "sqZfsPHT65dl",
        "outputId": "2e4de23b-c634-4f75-9a8b-4e941ecba5dc"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'modular_step_model/~/navier_stokes_step_fn/fused_learned_interpolation/periodic_conv_2d/~/conv2_d': {'b': (64,),\n",
              "  'w': (3, 3, 2, 64)},\n",
              " 'modular_step_model/~/navier_stokes_step_fn/fused_learned_interpolation/periodic_conv_2d_1/~/conv2_d': {'b': (64,),\n",
              "  'w': (3, 3, 64, 64)},\n",
              " 'modular_step_model/~/navier_stokes_step_fn/fused_learned_interpolation/periodic_conv_2d_2/~/conv2_d': {'b': (64,),\n",
              "  'w': (3, 3, 64, 64)},\n",
              " 'modular_step_model/~/navier_stokes_step_fn/fused_learned_interpolation/periodic_conv_2d_3/~/conv2_d': {'b': (64,),\n",
              "  'w': (3, 3, 64, 64)},\n",
              " 'modular_step_model/~/navier_stokes_step_fn/fused_learned_interpolation/periodic_conv_2d_4/~/conv2_d': {'b': (64,),\n",
              "  'w': (3, 3, 64, 64)},\n",
              " 'modular_step_model/~/navier_stokes_step_fn/fused_learned_interpolation/periodic_conv_2d_5/~/conv2_d': {'b': (64,),\n",
              "  'w': (3, 3, 64, 64)},\n",
              " 'modular_step_model/~/navier_stokes_step_fn/fused_learned_interpolation/periodic_conv_2d_6/~/conv2_d': {'b': (120,),\n",
              "  'w': (3, 3, 64, 120)}}"
            ]
          },
          "execution_count": 22,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "shape_structure(params)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oF3_xji9VbsP"
      },
      "source": [
        "# Model inference"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "metadata": {
        "id": "VBk07Syx6v6V"
      },
      "outputs": [],
      "source": [
        "#@title Setting up model configuration from the checkpoint;\n",
        "\n",
        "gin.clear_config()\n",
        "gin.parse_config(ckpt.model_config_str)\n",
        "gin.parse_config(strip_imports(reference_ds.attrs['physics_config_str']))\n",
        "dt = ckpt.model_time_step\n",
        "physics_specs = physics_specifications.get_physics_specs()\n",
        "model_cls = model_builder.get_model_cls(grid, dt, physics_specs)\n",
        "\n",
        "\n",
        "def compute_trajectory_fwd(x):\n",
        "  solver = model_cls()\n",
        "  x = solver.encode(x)\n",
        "  final, trajectory = solver.trajectory(\n",
        "      x, length, inner_steps, start_with_input=True, post_process_fn=solver.decode)\n",
        "  return trajectory\n",
        "\n",
        "\n",
        "model = hk.without_apply_rng(hk.transform(compute_trajectory_fwd))\n",
        "trajectory_fn = functools.partial(model.apply, params)\n",
        "trajectory_fn = jax.vmap(trajectory_fn)  # predict a batch of trajectories;"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8Ns9kfMJ7SD7",
        "outputId": "99b0c0e9-7dde-4fb3-a110-9c92d755c673"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "2024-05-21 16:24:30.604219: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_iota:\n",
            "2024-05-21 16:24:30.604268: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:24:30.604274: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:24:30.604278: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:24:30.604282: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:24:30.604286: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:24:30.604292: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:24:30.604296: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:24:30.604299: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:24:30.604304: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:24:30.604308: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:24:30.604312: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:24:30.604316: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:24:30.604321: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:24:30.604324: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:24:30.604328: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:24:30.604332: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:24:30.604335: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:24:30.604339: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:24:30.604342: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:30.604346: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:24:30.604350: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:24:30.604353: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:24:30.604357: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:24:30.604360: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:24:30.604364: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:24:30.604368: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:24:30.604371: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:24:30.604375: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:24:30.604378: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:24:30.604382: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:24:30.604386: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:24:30.604389: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:24:30.604393: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:24:30.604396: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:24:30.604400: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:24:30.604405: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:24:30.604408: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:24:30.604412: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:24:30.604416: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:24:30.604419: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:24:30.604423: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:24:30.604427: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:30.604430: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:30.604434: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:24:30.604438: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1\n",
            "2024-05-21 16:24:30.604442: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:24:30.604446: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:24:30.604505: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:24:30.604510: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:24:30.604513: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:24:30.604517: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:24:30.604522: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:24:30.604528: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:24:30.604534: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:24:30.604542: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:24:30.604549: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:24:30.604560: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:24:30.604572: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:24:30.604586: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:24:30.604611: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:24:30.604617: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:24:30.604621: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:24:30.604626: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:24:30.604630: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:24:34.017536: E external/xla/xla/stream_executor/rocm/rocm_driver.cc:1268] failed to allocate 47.99GiB (51527024640 bytes) from device: HIP_ERROR_OutOfMemory\n",
            "2024-05-21 16:24:34.017757: E external/xla/xla/stream_executor/rocm/rocm_driver.cc:1268] failed to allocate 43.19GiB (46374322176 bytes) from device: HIP_ERROR_OutOfMemory\n",
            "2024-05-21 16:24:34.017937: E external/xla/xla/stream_executor/rocm/rocm_driver.cc:1268] failed to allocate 38.87GiB (41736888320 bytes) from device: HIP_ERROR_OutOfMemory\n",
            "2024-05-21 16:24:34.018076: E external/xla/xla/stream_executor/rocm/rocm_driver.cc:1268] failed to allocate 34.98GiB (37563199488 bytes) from device: HIP_ERROR_OutOfMemory\n",
            "2024-05-21 16:24:34.018186: E external/xla/xla/stream_executor/rocm/rocm_driver.cc:1268] failed to allocate 31.48GiB (33806878720 bytes) from device: HIP_ERROR_OutOfMemory\n",
            "2024-05-21 16:24:34.018291: E external/xla/xla/stream_executor/rocm/rocm_driver.cc:1268] failed to allocate 28.34GiB (30426189824 bytes) from device: HIP_ERROR_OutOfMemory\n",
            "2024-05-21 16:24:34.018408: E external/xla/xla/stream_executor/rocm/rocm_driver.cc:1268] failed to allocate 25.50GiB (27383570432 bytes) from device: HIP_ERROR_OutOfMemory\n",
            "2024-05-21 16:24:34.018527: E external/xla/xla/stream_executor/rocm/rocm_driver.cc:1268] failed to allocate 22.95GiB (24645212160 bytes) from device: HIP_ERROR_OutOfMemory\n",
            "2024-05-21 16:24:34.018638: E external/xla/xla/stream_executor/rocm/rocm_driver.cc:1268] failed to allocate 20.66GiB (22180689920 bytes) from device: HIP_ERROR_OutOfMemory\n",
            "2024-05-21 16:24:34.018731: E external/xla/xla/stream_executor/rocm/rocm_driver.cc:1268] failed to allocate 18.59GiB (19962619904 bytes) from device: HIP_ERROR_OutOfMemory\n",
            "2024-05-21 16:24:34.018858: E external/xla/xla/stream_executor/rocm/rocm_driver.cc:1268] failed to allocate 16.73GiB (17966356480 bytes) from device: HIP_ERROR_OutOfMemory\n",
            "2024-05-21 16:24:34.288634: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_fn:\n",
            "2024-05-21 16:24:34.288676: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:24:34.288683: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:24:34.288690: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:24:34.288696: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:24:34.288702: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:24:34.288708: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:24:34.288714: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:24:34.288719: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:24:34.288725: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:24:34.288731: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:24:34.288736: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:24:34.288744: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:24:34.288749: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:24:34.288754: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:24:34.288759: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:24:34.288765: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:24:34.288771: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:24:34.288777: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:24:34.288783: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:34.288790: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:24:34.288794: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:24:34.288799: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:24:34.288803: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:24:34.288808: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:24:34.288814: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:24:34.288819: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:24:34.288825: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:24:34.288831: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:24:34.288837: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:24:34.288842: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:24:34.288847: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:24:34.288853: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:24:34.288858: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:24:34.288865: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:24:34.288873: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:24:34.288886: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:24:34.288897: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:24:34.288904: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:24:34.288909: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:24:34.288913: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:24:34.288917: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:24:34.288921: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:34.288925: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:34.288929: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:24:34.288933: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1\n",
            "2024-05-21 16:24:34.288936: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:24:34.288942: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:24:34.288947: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:24:34.288952: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:24:34.288959: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:24:34.288970: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:24:34.288982: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:24:34.288996: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:24:34.289012: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:24:34.289018: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:24:34.289023: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:24:34.289027: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:24:34.289030: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:24:34.289034: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:24:34.289038: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:24:34.289041: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:24:34.289045: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:24:34.289048: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:24:34.289052: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:24:34.423721: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_fn:\n",
            "2024-05-21 16:24:34.423760: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:24:34.423765: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:24:34.423770: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:24:34.423773: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:24:34.423777: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:24:34.423781: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:24:34.423785: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:24:34.423788: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:24:34.423792: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:24:34.423796: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:24:34.423799: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:24:34.423804: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:24:34.423808: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:24:34.423812: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:24:34.423815: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:24:34.423819: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:24:34.423823: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:24:34.423826: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:24:34.423830: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:34.423833: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:24:34.423837: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:24:34.423840: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:24:34.423844: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:24:34.423848: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:24:34.423851: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:24:34.423855: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:24:34.423859: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:24:34.423863: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:24:34.423867: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:24:34.423871: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:24:34.423874: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:24:34.423878: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:24:34.423881: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:24:34.423885: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:24:34.423888: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:24:34.423894: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:24:34.423898: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:24:34.423902: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:24:34.423905: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:24:34.423909: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:24:34.423913: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:24:34.423918: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:34.423922: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:34.423925: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:24:34.423929: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1\n",
            "2024-05-21 16:24:34.423932: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:24:34.423936: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:24:34.423939: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:24:34.423943: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:24:34.423947: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:24:34.423950: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:24:34.423954: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:24:34.423957: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:24:34.423961: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:24:34.423965: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:24:34.423969: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:24:34.423973: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:24:34.423977: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:24:34.423981: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:24:34.423984: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:24:34.423988: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:24:34.423991: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:24:34.423995: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:24:34.423999: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:24:34.559189: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_fn:\n",
            "2024-05-21 16:24:34.559233: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:24:34.559238: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:24:34.559243: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:24:34.559247: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:24:34.559250: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:24:34.559254: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:24:34.559258: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:24:34.559261: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:24:34.559266: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:24:34.559270: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:24:34.559274: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:24:34.559278: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:24:34.559281: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:24:34.559285: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:24:34.559289: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:24:34.559292: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:24:34.559296: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:24:34.559299: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:24:34.559303: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:34.559307: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:24:34.559310: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:24:34.559314: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:24:34.559317: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:24:34.559321: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:24:34.559325: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:24:34.559329: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:24:34.559334: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:24:34.559337: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:24:34.559341: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:24:34.559344: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:24:34.559348: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:24:34.559354: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:24:34.559358: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:24:34.559363: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:24:34.559368: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:24:34.559379: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:24:34.559386: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:24:34.559392: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:24:34.559397: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:24:34.559403: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:24:34.559408: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:24:34.559413: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:34.559419: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:34.559424: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:24:34.559429: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1\n",
            "2024-05-21 16:24:34.559435: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:24:34.559442: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:24:34.559448: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:24:34.559466: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:24:34.559472: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:24:34.559477: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:24:34.559482: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:24:34.559487: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:24:34.559493: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:24:34.559499: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:24:34.559504: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:24:34.559510: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:24:34.559516: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:24:34.559521: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:24:34.559526: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:24:34.559531: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:24:34.559537: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:24:34.559542: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:24:34.559548: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:24:34.701821: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_broadcast_in_dim:\n",
            "2024-05-21 16:24:34.701868: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:24:34.701874: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:24:34.701881: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:24:34.701886: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:24:34.701892: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:24:34.701898: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:24:34.701904: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:24:34.701911: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:24:34.701917: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:24:34.701922: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:24:34.701928: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:24:34.701934: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:24:34.701939: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:24:34.701944: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:24:34.701951: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:24:34.701957: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:24:34.701963: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:24:34.701968: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:24:34.701973: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:34.701979: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:24:34.701984: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:24:34.701990: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:24:34.701995: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:24:34.702000: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:24:34.702007: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:24:34.702012: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:24:34.702017: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:24:34.702023: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:24:34.702028: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:24:34.702034: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:24:34.702040: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:24:34.702045: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:24:34.702050: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:24:34.702056: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:24:34.702061: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:24:34.702070: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:24:34.702076: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:24:34.702082: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:24:34.702087: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:24:34.702092: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:24:34.702098: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:24:34.702103: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:34.702109: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:34.702115: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:24:34.702120: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1\n",
            "2024-05-21 16:24:34.702126: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:24:34.702131: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:24:34.702136: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:24:34.702141: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:24:34.702148: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:24:34.702153: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:24:34.702159: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:24:34.702165: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:24:34.702170: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:24:34.702176: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:24:34.702181: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:24:34.702186: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:24:34.702191: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:24:34.702196: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:24:34.702202: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:24:34.702208: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:24:34.702214: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:24:34.702220: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:24:34.702225: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:24:34.829112: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_broadcast_in_dim:\n",
            "2024-05-21 16:24:34.829155: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:24:34.829162: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:24:34.829168: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:24:34.829173: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:24:34.829180: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:24:34.829186: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:24:34.829191: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:24:34.829198: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:24:34.829203: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:24:34.829208: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:24:34.829215: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:24:34.829220: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:24:34.829226: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:24:34.829232: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:24:34.829237: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:24:34.829244: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:24:34.829249: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:24:34.829255: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:24:34.829261: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:34.829266: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:24:34.829270: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:24:34.829276: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:24:34.829281: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:24:34.829287: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:24:34.829292: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:24:34.829298: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:24:34.829303: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:24:34.829308: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:24:34.829313: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:24:34.829319: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:24:34.829325: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:24:34.829331: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:24:34.829334: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:24:34.829338: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:24:34.829342: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:24:34.829349: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:24:34.829353: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:24:34.829356: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:24:34.829360: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:24:34.829363: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:24:34.829367: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:24:34.829371: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:34.829375: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:34.829379: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:24:34.829382: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1\n",
            "2024-05-21 16:24:34.829386: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:24:34.829390: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:24:34.829393: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:24:34.829397: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:24:34.829400: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:24:34.829404: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:24:34.829409: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:24:34.829414: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:24:34.829421: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:24:34.829425: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:24:34.829428: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:24:34.829432: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:24:34.829436: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:24:34.829439: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:24:34.829443: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:24:34.829446: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:24:34.829456: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:24:34.829460: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:24:34.829464: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:24:34.963207: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_fn:\n",
            "2024-05-21 16:24:34.963251: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:24:34.963258: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:24:34.963264: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:24:34.963270: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:24:34.963277: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:24:34.963282: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:24:34.963288: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:24:34.963293: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:24:34.963299: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:24:34.963304: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:24:34.963310: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:24:34.963316: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:24:34.963321: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:24:34.963327: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:24:34.963333: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:24:34.963339: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:24:34.963344: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:24:34.963349: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:24:34.963354: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:34.963360: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:24:34.963366: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:24:34.963372: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:24:34.963378: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:24:34.963383: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:24:34.963388: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:24:34.963394: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:24:34.963399: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:24:34.963404: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:24:34.963410: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:24:34.963416: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:24:34.963422: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:24:34.963427: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:24:34.963432: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:24:34.963437: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:24:34.963445: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:24:34.963465: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:24:34.963471: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:24:34.963476: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:24:34.963482: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:24:34.963487: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:24:34.963493: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:24:34.963498: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:34.963504: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:34.963509: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:24:34.963515: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1\n",
            "2024-05-21 16:24:34.963521: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:24:34.963527: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:24:34.963532: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:24:34.963538: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:24:34.963543: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:24:34.963548: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:24:34.963553: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:24:34.963559: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:24:34.963564: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:24:34.963571: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:24:34.963576: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:24:34.963582: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:24:34.963587: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:24:34.963592: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:24:34.963597: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:24:34.963602: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:24:34.963608: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:24:34.963614: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:24:34.963620: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:24:35.096494: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit__lambda_:\n",
            "2024-05-21 16:24:35.096538: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:24:35.096544: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:24:35.096551: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:24:35.096557: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:24:35.096563: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:24:35.096568: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:24:35.096574: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:24:35.096581: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:24:35.096586: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:24:35.096592: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:24:35.096598: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:24:35.096604: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:24:35.096610: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:24:35.096615: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:24:35.096621: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:24:35.096627: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:24:35.096633: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:24:35.096638: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:24:35.096643: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:35.096647: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:24:35.096653: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:24:35.096657: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:24:35.096663: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:24:35.096668: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:24:35.096673: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:24:35.096678: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:24:35.096683: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:24:35.096689: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:24:35.096694: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:24:35.096700: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:24:35.096707: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:24:35.096715: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:24:35.096718: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:24:35.096722: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:24:35.096726: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:24:35.096734: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:24:35.096738: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:24:35.096742: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:24:35.096746: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:24:35.096750: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:24:35.096754: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:24:35.096757: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:35.096761: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:35.096765: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:24:35.096768: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1\n",
            "2024-05-21 16:24:35.096772: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:24:35.096776: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:24:35.096779: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:24:35.096783: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:24:35.096786: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:24:35.096790: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:24:35.096795: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:24:35.096802: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:24:35.096807: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:24:35.096811: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:24:35.096815: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:24:35.096818: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:24:35.096822: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:24:35.096825: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:24:35.096829: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:24:35.096833: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:24:35.096836: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:24:35.096840: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:24:35.096844: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:24:35.311994: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_fn:\n",
            "2024-05-21 16:24:35.312044: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:24:35.312052: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:24:35.312059: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:24:35.312068: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:24:35.312074: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:24:35.312079: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:24:35.312083: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:24:35.312088: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:24:35.312094: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:24:35.312100: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:24:35.312105: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:24:35.312113: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:24:35.312119: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:24:35.312124: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:24:35.312130: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:24:35.312137: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:24:35.312142: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:24:35.312148: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:24:35.312154: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:35.312159: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:24:35.312165: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:24:35.312170: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:24:35.312175: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:24:35.312181: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:24:35.312186: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:24:35.312192: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:24:35.312198: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:24:35.312203: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:24:35.312208: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:24:35.312213: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:24:35.312219: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:24:35.312225: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:24:35.312231: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:24:35.312237: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:24:35.312242: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:24:35.312251: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:24:35.312257: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:24:35.312262: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:24:35.312267: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:24:35.312273: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:24:35.312282: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:24:35.312291: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:35.312299: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:35.312308: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:24:35.312317: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1\n",
            "2024-05-21 16:24:35.312324: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:24:35.312331: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:24:35.312335: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:24:35.312339: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:24:35.312343: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:24:35.312347: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:24:35.312351: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:24:35.312354: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:24:35.312359: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:24:35.312364: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:24:35.312369: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:24:35.312376: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:24:35.312385: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:24:35.312397: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:24:35.312414: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:24:35.312430: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:24:35.312437: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:24:35.312442: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:24:35.312446: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:24:35.447428: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_convert_element_type:\n",
            "2024-05-21 16:24:35.447482: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:24:35.447488: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:24:35.447495: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:24:35.447501: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:24:35.447506: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:24:35.447513: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:24:35.447519: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:24:35.447525: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:24:35.447530: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:24:35.447536: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:24:35.447541: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:24:35.447547: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:24:35.447552: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:24:35.447559: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:24:35.447565: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:24:35.447570: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:24:35.447575: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:24:35.447581: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:24:35.447586: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:35.447592: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:24:35.447597: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:24:35.447603: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:24:35.447609: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:24:35.447615: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:24:35.447620: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:24:35.447625: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:24:35.447631: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:24:35.447636: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:24:35.447642: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:24:35.447647: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:24:35.447653: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:24:35.447660: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:24:35.447665: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:24:35.447670: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:24:35.447675: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:24:35.447688: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:24:35.447694: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:24:35.447700: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:24:35.447705: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:24:35.447711: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:24:35.447716: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:24:35.447721: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:35.447727: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:35.447733: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:24:35.447738: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1\n",
            "2024-05-21 16:24:35.447743: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:24:35.447749: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:24:35.447754: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:24:35.447759: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:24:35.447765: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:24:35.447771: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:24:35.447777: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:24:35.447782: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:24:35.447787: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:24:35.447792: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:24:35.447798: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:24:35.447805: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:24:35.447811: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:24:35.447817: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:24:35.447822: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:24:35.447828: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:24:35.447834: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:24:35.447839: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:24:35.447844: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:24:35.575316: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_broadcast_in_dim:\n",
            "2024-05-21 16:24:35.575361: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:24:35.575367: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:24:35.575373: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:24:35.575379: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:24:35.575386: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:24:35.575392: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:24:35.575398: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:24:35.575404: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:24:35.575409: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:24:35.575415: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:24:35.575421: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:24:35.575426: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:24:35.575431: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:24:35.575437: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:24:35.575443: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:24:35.575459: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:24:35.575465: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:24:35.575471: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:24:35.575476: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:35.575480: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:24:35.575486: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:24:35.575490: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:24:35.575497: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:24:35.575503: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:24:35.575508: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:24:35.575514: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:24:35.575519: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:24:35.575525: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:24:35.575530: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:24:35.575536: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:24:35.575542: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:24:35.575547: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:24:35.575552: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:24:35.575558: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:24:35.575561: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:24:35.575569: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:24:35.575572: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:24:35.575576: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:24:35.575580: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:24:35.575584: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:24:35.575588: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:24:35.575592: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:35.575595: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:35.575599: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:24:35.575602: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1\n",
            "2024-05-21 16:24:35.575606: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:24:35.575610: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:24:35.575613: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:24:35.575617: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:24:35.575620: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:24:35.575624: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:24:35.575628: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:24:35.575633: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:24:35.575638: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:24:35.575644: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:24:35.575647: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:24:35.575651: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:24:35.575655: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:24:35.575658: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:24:35.575662: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:24:35.575665: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:24:35.575669: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:24:35.575672: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:24:35.575676: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:24:35.713553: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_convert_element_type:\n",
            "2024-05-21 16:24:35.713599: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:24:35.713606: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:24:35.713612: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:24:35.713618: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:24:35.713624: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:24:35.713630: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:24:35.713636: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:24:35.713642: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:24:35.713647: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:24:35.713653: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:24:35.713659: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:24:35.713665: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:24:35.713671: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:24:35.713677: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:24:35.713682: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:24:35.713687: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:24:35.713693: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:24:35.713698: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:24:35.713703: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:35.713709: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:24:35.713714: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:24:35.713720: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:24:35.713727: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:24:35.713732: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:24:35.713737: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:24:35.713742: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:24:35.713748: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:24:35.713754: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:24:35.713760: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:24:35.713765: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:24:35.713771: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:24:35.713777: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:24:35.713783: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:24:35.713788: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:24:35.713793: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:24:35.713803: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:24:35.713809: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:24:35.713814: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:24:35.713820: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:24:35.713826: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:24:35.713832: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:24:35.713837: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:35.713842: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:35.713847: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:24:35.713853: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1\n",
            "2024-05-21 16:24:35.713858: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:24:35.713864: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:24:35.713870: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:24:35.713877: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:24:35.713882: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:24:35.713887: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:24:35.713893: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:24:35.713899: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:24:35.713904: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:24:35.713909: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:24:35.713915: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:24:35.713920: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:24:35.713926: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:24:35.713932: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:24:35.713937: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:24:35.713943: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:24:35.713949: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:24:35.713954: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:24:35.713960: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:24:35.846618: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_convert_element_type:\n",
            "2024-05-21 16:24:35.846663: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:24:35.846669: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:24:35.846674: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:24:35.846680: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:24:35.846686: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:24:35.846691: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:24:35.846698: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:24:35.846704: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:24:35.846709: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:24:35.846715: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:24:35.846721: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:24:35.846727: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:24:35.846732: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:24:35.846737: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:24:35.846743: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:24:35.846749: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:24:35.846755: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:24:35.846761: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:24:35.846766: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:35.846771: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:24:35.846775: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:24:35.846780: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:24:35.846785: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:24:35.846792: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:24:35.846797: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:24:35.846803: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:24:35.846808: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:24:35.846813: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:24:35.846818: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:24:35.846823: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:24:35.846830: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:24:35.846839: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:24:35.846842: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:24:35.846846: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:24:35.846850: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:24:35.846857: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:24:35.846861: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:24:35.846865: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:24:35.846868: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:24:35.846872: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:24:35.846875: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:24:35.846879: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:35.846884: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:35.846887: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:24:35.846891: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1\n",
            "2024-05-21 16:24:35.846894: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:24:35.846898: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:24:35.846901: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:24:35.846905: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:24:35.846909: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:24:35.846912: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:24:35.846917: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:24:35.846922: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:24:35.846928: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:24:35.846932: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:24:35.846936: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:24:35.846939: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:24:35.846943: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:24:35.846946: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:24:35.846950: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:24:35.846954: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:24:35.846957: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:24:35.846961: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:24:35.846964: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:24:35.875345: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_fn:\n",
            "2024-05-21 16:24:35.875391: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:24:35.875397: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:24:35.875404: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:24:35.875410: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:24:35.875415: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:24:35.875421: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:24:35.875426: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:24:35.875432: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:24:35.875437: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:24:35.875442: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:24:35.875461: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:24:35.875467: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:24:35.875473: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:24:35.875478: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:24:35.875483: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:24:35.875489: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:24:35.875494: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:24:35.875500: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:24:35.875505: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:35.875510: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:24:35.875516: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:24:35.875521: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:24:35.875527: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:24:35.875532: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:24:35.875538: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:24:35.875544: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:24:35.875550: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:24:35.875555: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:24:35.875560: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:24:35.875566: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:24:35.875571: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:24:35.875577: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:24:35.875582: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:24:35.875588: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:24:35.875593: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:24:35.875603: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:24:35.875609: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:24:35.875615: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:24:35.875620: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:24:35.875626: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:24:35.875631: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:24:35.875637: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:35.875642: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:35.875648: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:24:35.875654: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1\n",
            "2024-05-21 16:24:35.875659: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:24:35.875664: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:24:35.875670: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:24:35.875675: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:24:35.875680: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:24:35.875685: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:24:35.875692: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:24:35.875698: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:24:35.875705: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:24:35.875711: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:24:35.875716: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:24:35.875721: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:24:35.875726: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:24:35.875731: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:24:35.875737: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:24:35.875742: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:24:35.875748: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:24:35.875753: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:24:35.875760: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:24:36.009497: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_fn:\n",
            "2024-05-21 16:24:36.009542: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:24:36.009548: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:24:36.009554: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:24:36.009560: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:24:36.009565: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:24:36.009571: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:24:36.009577: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:24:36.009583: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:24:36.009589: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:24:36.009595: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:24:36.009601: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:24:36.009607: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:24:36.009612: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:24:36.009617: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:24:36.009624: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:24:36.009629: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:24:36.009634: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:24:36.009640: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:24:36.009645: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:36.009650: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:24:36.009656: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:24:36.009662: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:24:36.009667: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:24:36.009673: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:24:36.009679: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:24:36.009684: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:24:36.009689: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:24:36.009695: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:24:36.009700: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:24:36.009706: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:24:36.009712: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:24:36.009718: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:24:36.009724: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:24:36.009730: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:24:36.009736: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:24:36.009746: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:24:36.009753: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:24:36.009758: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:24:36.009764: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:24:36.009770: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:24:36.009776: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:24:36.009781: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:36.009787: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:36.009792: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:24:36.009797: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1\n",
            "2024-05-21 16:24:36.009803: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:24:36.009808: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:24:36.009814: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:24:36.009820: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:24:36.009826: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:24:36.009832: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:24:36.009837: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:24:36.009842: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:24:36.009848: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:24:36.009854: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:24:36.009859: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:24:36.009864: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:24:36.009869: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:24:36.009875: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:24:36.009882: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:24:36.009887: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:24:36.009893: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:24:36.009898: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:24:36.009904: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:24:36.145028: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_select_n:\n",
            "2024-05-21 16:24:36.145070: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:24:36.145075: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:24:36.145080: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:24:36.145085: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:24:36.145088: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:24:36.145092: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:24:36.145096: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:24:36.145100: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:24:36.145103: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:24:36.145107: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:24:36.145111: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:24:36.145114: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:24:36.145118: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:24:36.145122: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:24:36.145125: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:24:36.145129: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:24:36.145133: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:24:36.145137: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:24:36.145140: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:36.145144: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:24:36.145148: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:24:36.145151: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:24:36.145155: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:24:36.145158: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:24:36.145162: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:24:36.145166: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:24:36.145169: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:24:36.145173: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:24:36.145177: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:24:36.145181: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:24:36.145185: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:24:36.145188: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:24:36.145192: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:24:36.145196: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:24:36.145199: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:24:36.145206: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:24:36.145210: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:24:36.145213: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:24:36.145217: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:24:36.145220: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:24:36.145224: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:24:36.145228: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:36.145232: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:36.145236: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:24:36.145239: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1\n",
            "2024-05-21 16:24:36.145243: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:24:36.145247: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:24:36.145250: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:24:36.145254: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:24:36.145257: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:24:36.145261: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:24:36.145265: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:24:36.145268: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:24:36.145272: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:24:36.145276: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:24:36.145280: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:24:36.145284: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:24:36.145287: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:24:36.145291: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:24:36.145295: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:24:36.145298: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:24:36.145302: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:24:36.145305: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:24:36.145309: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:24:36.278506: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_broadcast_in_dim:\n",
            "2024-05-21 16:24:36.278552: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:24:36.278557: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:24:36.278562: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:24:36.278566: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:24:36.278569: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:24:36.278573: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:24:36.278577: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:24:36.278580: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:24:36.278584: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:24:36.278587: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:24:36.278591: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:24:36.278595: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:24:36.278599: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:24:36.278603: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:24:36.278607: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:24:36.278611: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:24:36.278614: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:24:36.278618: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:24:36.278621: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:36.278625: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:24:36.278628: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:24:36.278632: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:24:36.278636: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:24:36.278639: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:24:36.278643: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:24:36.278647: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:24:36.278651: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:24:36.278654: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:24:36.278658: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:24:36.278662: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:24:36.278665: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:24:36.278669: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:24:36.278672: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:24:36.278676: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:24:36.278679: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:24:36.278689: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:24:36.278694: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:24:36.278698: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:24:36.278702: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:24:36.278705: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:24:36.278709: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:24:36.278712: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:36.278716: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:36.278720: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:24:36.278723: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1\n",
            "2024-05-21 16:24:36.278727: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:24:36.278730: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:24:36.278734: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:24:36.278739: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:24:36.278743: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:24:36.278747: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:24:36.278750: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:24:36.278754: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:24:36.278758: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:24:36.278761: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:24:36.278765: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:24:36.278768: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:24:36.278772: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:24:36.278776: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:24:36.278779: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:24:36.278783: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:24:36.278786: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:24:36.278790: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:24:36.278794: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:24:36.408883: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_concatenate:\n",
            "2024-05-21 16:24:36.408920: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:24:36.408925: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:24:36.408930: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:24:36.408934: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:24:36.408937: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:24:36.408941: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:24:36.408945: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:24:36.408949: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:24:36.408953: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:24:36.408956: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:24:36.408960: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:24:36.408963: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:24:36.408967: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:24:36.408971: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:24:36.408974: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:24:36.408978: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:24:36.408982: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:24:36.408985: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:24:36.408989: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:36.408994: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:24:36.408997: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:24:36.409001: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:24:36.409004: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:24:36.409008: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:24:36.409011: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:24:36.409015: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:24:36.409019: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:24:36.409022: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:24:36.409026: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:24:36.409029: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:24:36.409033: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:24:36.409037: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:24:36.409041: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:24:36.409045: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:24:36.409048: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:24:36.409058: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:24:36.409064: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:24:36.409070: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:24:36.409075: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:24:36.409081: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:24:36.409085: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:24:36.409089: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:36.409093: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:36.409096: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:24:36.409100: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1\n",
            "2024-05-21 16:24:36.409103: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:24:36.409107: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:24:36.409110: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:24:36.409114: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:24:36.409117: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:24:36.409121: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:24:36.409125: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:24:36.409128: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:24:36.409132: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:24:36.409136: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:24:36.409140: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:24:36.409143: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:24:36.409147: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:24:36.409150: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:24:36.409154: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:24:36.409157: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:24:36.409161: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:24:36.409164: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:24:36.409168: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:24:36.543600: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_gather:\n",
            "2024-05-21 16:24:36.543640: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:24:36.543645: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:24:36.543649: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:24:36.543653: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:24:36.543657: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:24:36.543661: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:24:36.543664: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:24:36.543668: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:24:36.543673: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:24:36.543676: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:24:36.543680: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:24:36.543683: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:24:36.543687: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:24:36.543691: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:24:36.543694: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:24:36.543698: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:24:36.543702: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:24:36.543705: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:24:36.543709: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:36.543712: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:24:36.543717: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:24:36.543722: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:24:36.543727: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:24:36.543733: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:24:36.543741: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:24:36.543748: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:24:36.543753: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:24:36.543759: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:24:36.543766: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:24:36.543776: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:24:36.543783: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:24:36.543794: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:24:36.543807: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:24:36.543821: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:24:36.543826: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:24:36.543832: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:24:36.543838: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:24:36.543844: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:24:36.543849: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:24:36.543856: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:24:36.543869: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:24:36.543878: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:36.543886: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:36.543897: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:24:36.543904: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1\n",
            "2024-05-21 16:24:36.543909: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:24:36.543915: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:24:36.543920: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:24:36.543927: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:24:36.543932: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:24:36.543941: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:24:36.543951: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:24:36.543965: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:24:36.543976: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:24:36.543992: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:24:36.544001: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:24:36.544010: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:24:36.544015: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:24:36.544020: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:24:36.544028: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:24:36.544038: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:24:36.544049: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:24:36.544057: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:24:36.544068: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:24:36.709719: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_squeeze:\n",
            "2024-05-21 16:24:36.709764: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:24:36.709769: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:24:36.709773: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:24:36.709777: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:24:36.709781: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:24:36.709786: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:24:36.709790: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:24:36.709793: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:24:36.709797: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:24:36.709801: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:24:36.709804: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:24:36.709808: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:24:36.709811: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:24:36.709815: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:24:36.709819: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:24:36.709822: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:24:36.709826: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:24:36.709829: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:24:36.709833: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:36.709837: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:24:36.709841: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:24:36.709845: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:24:36.709848: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:24:36.709852: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:24:36.709855: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:24:36.709859: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:24:36.709862: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:24:36.709866: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:24:36.709869: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:24:36.709873: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:24:36.709877: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:24:36.709880: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:24:36.709884: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:24:36.709888: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:24:36.709892: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:24:36.709901: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:24:36.709907: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:24:36.709912: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:24:36.709917: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:24:36.709923: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:24:36.709926: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:24:36.709930: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:36.709934: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:36.709938: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:24:36.709942: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1\n",
            "2024-05-21 16:24:36.709946: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:24:36.709949: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:24:36.709953: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:24:36.709956: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:24:36.709960: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:24:36.709963: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:24:36.709967: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:24:36.709970: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:24:36.709974: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:24:36.709977: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:24:36.709981: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:24:36.709985: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:24:36.709989: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:24:36.709993: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:24:36.709996: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:24:36.710000: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:24:36.710003: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:24:36.710007: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:24:36.710010: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:24:37.223923: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_scan:\n",
            "2024-05-21 16:24:37.223966: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:24:37.223971: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:24:37.223975: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:24:37.223979: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:24:37.223983: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:24:37.223987: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:24:37.223990: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:24:37.223994: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:24:37.223998: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:24:37.224002: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:24:37.224005: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:24:37.224009: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:24:37.224013: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:24:37.224016: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:24:37.224020: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:24:37.224024: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:24:37.224027: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:24:37.224031: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:24:37.224034: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:37.224038: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:24:37.224042: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:24:37.224046: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:24:37.224050: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:24:37.224053: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:24:37.224057: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:24:37.224060: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:24:37.224064: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:24:37.224068: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:24:37.224071: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:24:37.224075: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:24:37.224078: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:24:37.224082: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:24:37.224086: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:24:37.224090: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:24:37.224094: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:24:37.224100: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:24:37.224104: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:24:37.224107: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:24:37.224111: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:24:37.224115: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:24:37.224118: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:24:37.224122: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:37.224126: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:37.224131: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:24:37.224137: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1\n",
            "2024-05-21 16:24:37.224142: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:24:37.224148: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:24:37.224154: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:24:37.224160: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:24:37.224167: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:24:37.224178: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:24:37.224192: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:24:37.224208: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:24:37.224225: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:24:37.224231: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:24:37.224235: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:24:37.224238: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_detailed_logging: true\n",
            "2024-05-21 16:24:37.224242: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:24:37.224246: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:24:37.224249: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:24:37.224253: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:24:37.224256: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:24:37.224260: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:24:37.224264: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:24:37.224267: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:24:38.152667: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_transpose:\n",
            "2024-05-21 16:24:38.152713: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:24:38.152718: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:24:38.152723: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:24:38.152728: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:24:38.152732: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:24:38.152737: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:24:38.152741: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:24:38.152745: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:24:38.152749: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:24:38.152754: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:24:38.152758: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:24:38.152763: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:24:38.152767: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:24:38.152771: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:24:38.152775: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:24:38.152779: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:24:38.152783: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:24:38.152788: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:24:38.152792: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:38.152796: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:24:38.152800: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:24:38.152804: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:24:38.152808: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:24:38.152813: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:24:38.152817: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:24:38.152822: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:24:38.152826: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:24:38.152830: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:24:38.152835: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:24:38.152839: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:24:38.152843: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:24:38.152847: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:24:38.152851: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:24:38.152855: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:24:38.152859: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:24:38.152867: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:24:38.152871: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:24:38.152875: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:24:38.152880: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:24:38.152885: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:24:38.152890: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:24:38.152894: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:38.152898: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:24:38.152902: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:24:38.152906: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1\n",
            "2024-05-21 16:24:38.152910: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:24:38.152914: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:24:38.152919: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:24:38.152923: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:24:38.152927: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:24:38.152932: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:24:38.152936: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:24:38.152940: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:24:38.152944: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:24:38.152948: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:24:38.152953: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:24:38.152957: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:24:38.152962: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:24:38.152966: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:24:38.152970: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:24:38.152974: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:24:38.152979: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:24:38.152983: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:24:38.152987: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n"
          ]
        }
      ],
      "source": [
        "#@title Running inference;\n",
        "\n",
        "prediction = trajectory_fn(initial_conditions)\n",
        "prediction_ds = cfd_data.xarray_utils.velocity_trajectory_to_xarray(\n",
        "    prediction, grid, samples=True)\n",
        "\n",
        "# roundoff error in coordinates sometimes leads to wrong alignment results;\n",
        "prediction_ds.coords['x'] = target_ds.coords['x']\n",
        "prediction_ds.coords['y'] = target_ds.coords['y']\n",
        "prediction_ds.coords['time'] = target_ds.coords['time']\n",
        "\n",
        "datasets[model_name] = prediction_ds"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2TPzWKdWWTpc"
      },
      "source": [
        "# Computing summaries\n",
        "\n",
        "**Note:** Evaluations in this notebook are demonstrative and performed over a single sample and shorter times than those used in the paper;"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "metadata": {
        "id": "oQQUugPn7XbC"
      },
      "outputs": [],
      "source": [
        "summary = xarray.concat([\n",
        "    cfd_data.evaluation.compute_summary_dataset(ds, target_ds)\n",
        "    for ds in datasets.values()\n",
        "], dim='model')\n",
        "summary.coords['model'] = list(datasets.keys())\n",
        "\n",
        "correlation = summary.vorticity_correlation.compute()\n",
        "spectrum = summary.energy_spectrum_mean.mean('time').compute()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "metadata": {
        "id": "PCmgDeDqc4zr"
      },
      "outputs": [],
      "source": [
        "baseline_palette = seaborn.color_palette('YlGnBu', n_colors=7)[1:]\n",
        "models_color = seaborn.xkcd_palette(['burnt orange', 'taupe', 'greenish blue'])\n",
        "palette = baseline_palette + models_color[:(len(datasets.keys()) - 6)]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 27,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 405
        },
        "id": "v0HoNYna9Stp",
        "outputId": "e801f83f-d470-4e6d-e348-bf3c3b8e96bc"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(0.0, 15.0)"
            ]
          },
          "execution_count": 27,
          "metadata": {},
          "output_type": "execute_result"
        },
        {
          "data": {
            "image/png": "[encoded data, 119828 characters]",
            "text/plain": [
              "<Figure size 700x600 with 1 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "#@title Vorticity correlation as a function of time\n",
        "\n",
        "plt.figure(figsize=(7, 6))\n",
        "for color, model in zip(palette, summary['model'].data):\n",
        "  style = '-' if 'baseline' in model else '--'\n",
        "  correlation.sel(model=model).plot.line(\n",
        "      color=color, linestyle=style, label=model, linewidth=3);\n",
        "plt.axhline(y=0.95, xmin=0, xmax=20, color='gray')\n",
        "plt.legend();\n",
        "plt.title('')\n",
        "plt.xlim(0, 15)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 392
        },
        "id": "-EhXq3cX9SoB",
        "outputId": "1b8976d4-c268-42e5-aadf-cd206295ece0"
      },
      "outputs": [
        {
          "data": {
            "image/png": "[encoded data, 114228 characters]",
            "text/plain": [
              "<Figure size 1000x600 with 1 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "#@title Energy spectrum\n",
        "\n",
        "plt.figure(figsize=(10, 6))\n",
        "for color, model in zip(palette, summary['model'].data):\n",
        "  style = '-' if 'baseline' in model else '--'\n",
        "  (spectrum.k ** 5 * spectrum).sel(model=model).plot.line(\n",
        "      color=color, linestyle=style, label=model, linewidth=3);\n",
        "plt.legend();\n",
        "plt.yscale('log')\n",
        "plt.xscale('log')\n",
        "plt.title('')\n",
        "plt.xlim(3.5, None)\n",
        "if dataset_name == 'kolmogorov_re_4000':\n",
        "  plt.ylim(5e8, None)\n",
        "elif dataset_name == 'kolmogorov_re_1000':\n",
        "  plt.ylim(1e9, None)\n",
        "elif dataset_name == 'decaying':\n",
        "  plt.ylim(2e8, None)\n",
        "else:\n",
        "  raise ValueError('Unrecognized dataset')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 29,
      "metadata": {
        "id": "LULDCYOc8FQv"
      },
      "outputs": [],
      "source": [
        "vorticities = xarray.concat(\n",
        "    [cfd_data.xarray_utils.vorticity_2d(ds) for ds in datasets.values()],\n",
        "    dim='model'\n",
        ").to_dataset()\n",
        "vorticities.coords['model'] = list(datasets.keys())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 30,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "MpYjDWfE9SjC",
        "outputId": "83ab0c48-f2a2-49ed-e692-1103ea258863"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "<xarray.plot.facetgrid.FacetGrid at 0x7fa2081bc2e0>"
            ]
          },
          "execution_count": 30,
          "metadata": {},
          "output_type": "execute_result"
        },
        {
          "data": {
            "image/png": "[encoded data, 859456 characters]",
            "text/plain": [
              "<Figure size 1600x2100 with 36 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "#@title Visualizing model unrolls { form-width: \"30%\", run: \"auto\"}\n",
        "time_range = {'min': 0, 'max': vorticities.sizes['time'], 'step': 1}\n",
        "\n",
        "last_step_to_plot = 200 #@param {type: \"slider\", min: 1, max: 200 , step: 5}\n",
        "num_to_show = 5 #@param {type: \"slider\", min: 1, max: 10, step: 1}\n",
        "time_slice = slice(None, last_step_to_plot, last_step_to_plot // num_to_show)\n",
        "\n",
        "(vorticities.isel({'time': time_slice, 'sample': 0})['vorticity']\n",
        " .plot.imshow(row='model', col='time', cmap=seaborn.cm.icefire, robust=True))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {
        "id": "Z-JMFbSKODM2"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "ml_model_inference_demo.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
