{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4K4FCNvH8dtm"
      },
      "source": [
        "# JAX-CFD channel flow demo\n",
        "\n",
        "This demonstration shows how to use JAX-CFD to simulate 2d flow between parallel no-slip walls."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "executionInfo": {
          "elapsed": 5628,
          "status": "ok",
          "timestamp": 1655317705710,
          "user": {
            "displayName": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "n2J40uBCo9By"
      },
      "outputs": [],
      "source": [
        "import jax\n",
        "import jax.numpy as jnp\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import seaborn\n",
        "import jax_cfd.base as cfd\n",
        "import xarray\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "executionInfo": {
          "elapsed": 5028,
          "status": "ok",
          "timestamp": 1655317714655,
          "user": {
            "displayName": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "KcJcpDavpAoi"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "2024-05-21 16:38:58.783926: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_iota:\n",
            "2024-05-21 16:38:58.783962: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:38:58.783976: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:38:58.783981: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:38:58.783986: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:38:58.783990: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:38:58.783993: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:38:58.783997: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:38:58.784001: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:38:58.784005: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:38:58.784009: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:38:58.784012: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:38:58.784016: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:38:58.784020: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:38:58.784023: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:38:58.784027: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:38:58.784031: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:38:58.784034: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:38:58.784038: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:38:58.784041: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:38:58.784045: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:38:58.784048: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:38:58.784052: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:38:58.784055: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:38:58.784059: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:38:58.784063: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:38:58.784066: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:38:58.784070: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:38:58.784073: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:38:58.784077: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:38:58.784080: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:38:58.784084: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:38:58.784088: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:38:58.784091: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:38:58.784095: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:38:58.784098: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:38:58.784104: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:38:58.784107: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:38:58.784111: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:38:58.784115: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:38:58.784118: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:38:58.784122: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:38:58.784126: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:38:58.784129: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:38:58.784133: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:38:58.784137: 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:38:58.784140: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:38:58.784144: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:38:58.784148: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:38:58.784151: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:38:58.784155: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:38:58.784159: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:38:58.784163: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:38:58.784167: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:38:58.784171: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:38:58.784174: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:38:58.784178: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:38:58.784182: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:38:58.784185: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:38:58.784189: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:38:58.784193: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:38:58.784196: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:38:58.784201: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:38:58.784204: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:38:58.784208: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:39:01.515588: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_fn:\n",
            "2024-05-21 16:39:01.515634: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:39:01.515643: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:39:01.515648: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:39:01.515652: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:39:01.515656: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:39:01.515660: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:39:01.515664: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:39:01.515668: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:39:01.515671: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:39:01.515675: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:39:01.515678: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:39:01.515682: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:39:01.515686: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:39:01.515689: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:39:01.515693: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:39:01.515697: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:39:01.515700: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:39:01.515704: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:39:01.515707: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:01.515711: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:39:01.515715: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:39:01.515718: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:39:01.515722: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:39:01.515725: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:39:01.515729: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:39:01.515732: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:39:01.515736: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:39:01.515740: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:39:01.515743: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:39:01.515747: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:39:01.515750: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:39:01.515754: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:39:01.515758: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:39:01.515761: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:39:01.515765: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:39:01.515771: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:39:01.515775: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:39:01.515778: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:39:01.515782: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:39:01.515785: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:39:01.515789: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:39:01.515793: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:01.515796: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:01.515800: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:39:01.515803: 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:39:01.515807: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:39:01.515811: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:39:01.515814: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:39:01.515818: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:39:01.515821: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:39:01.515825: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:39:01.515829: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:39:01.515832: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:39:01.515836: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:39:01.515839: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:39:01.515843: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:39:01.515847: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:39:01.515850: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:39:01.515854: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:39:01.515858: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:39:01.515861: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:39:01.515865: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:39:01.515868: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:39:01.515872: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:39:01.657169: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_fn:\n",
            "2024-05-21 16:39:01.657198: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:39:01.657203: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:39:01.657209: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:39:01.657213: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:39:01.657217: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:39:01.657221: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:39:01.657225: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:39:01.657230: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:39:01.657234: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:39:01.657237: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:39:01.657241: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:39:01.657245: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:39:01.657248: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:39:01.657252: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:39:01.657256: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:39:01.657260: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:39:01.657264: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:39:01.657268: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:39:01.657272: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:01.657276: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:39:01.657279: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:39:01.657283: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:39:01.657288: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:39:01.657293: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:39:01.657297: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:39:01.657301: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:39:01.657305: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:39:01.657309: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:39:01.657314: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:39:01.657318: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:39:01.657322: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:39:01.657326: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:39:01.657329: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:39:01.657333: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:39:01.657337: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:39:01.657341: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:39:01.657349: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:39:01.657353: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:39:01.657357: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:39:01.657361: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:39:01.657365: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:39:01.657372: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:01.657376: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:01.657379: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:39:01.657383: 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:39:01.657387: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:39:01.657391: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:39:01.657396: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:39:01.657400: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:39:01.657403: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:39:01.657407: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:39:01.657412: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:39:01.657415: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:39:01.657419: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:39:01.657424: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:39:01.657428: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:39:01.657432: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:39:01.657436: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:39:01.657440: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:39:01.657443: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:39:01.657447: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:39:01.657461: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:39:01.657465: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:39:01.657469: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:39:01.793002: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_fn:\n",
            "2024-05-21 16:39:01.793034: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:39:01.793040: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:39:01.793045: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:39:01.793050: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:39:01.793053: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:39:01.793057: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:39:01.793061: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:39:01.793065: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:39:01.793069: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:39:01.793073: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:39:01.793077: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:39:01.793082: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:39:01.793086: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:39:01.793089: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:39:01.793093: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:39:01.793097: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:39:01.793101: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:39:01.793105: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:39:01.793109: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:01.793112: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:39:01.793116: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:39:01.793119: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:39:01.793123: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:39:01.793127: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:39:01.793130: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:39:01.793134: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:39:01.793138: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:39:01.793142: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:39:01.793145: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:39:01.793149: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:39:01.793153: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:39:01.793157: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:39:01.793160: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:39:01.793164: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:39:01.793167: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:39:01.793175: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:39:01.793178: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:39:01.793182: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:39:01.793186: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:39:01.793189: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:39:01.793193: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:39:01.793197: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:01.793201: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:01.793204: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:39:01.793208: 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:39:01.793211: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:39:01.793215: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:39:01.793219: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:39:01.793222: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:39:01.793226: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:39:01.793230: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:39:01.793233: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:39:01.793237: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:39:01.793241: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:39:01.793245: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:39:01.793248: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:39:01.793252: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:39:01.793256: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:39:01.793259: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:39:01.793263: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:39:01.793267: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:39:01.793270: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:39:01.793274: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:39:01.793277: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:39:01.926581: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_iota:\n",
            "2024-05-21 16:39:01.926615: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:39:01.926621: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:39:01.926627: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:39:01.926631: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:39:01.926635: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:39:01.926639: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:39:01.926643: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:39:01.926647: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:39:01.926650: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:39:01.926654: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:39:01.926658: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:39:01.926661: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:39:01.926665: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:39:01.926669: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:39:01.926672: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:39:01.926676: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:39:01.926679: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:39:01.926683: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:39:01.926688: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:01.926692: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:39:01.926695: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:39:01.926699: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:39:01.926703: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:39:01.926707: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:39:01.926710: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:39:01.926714: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:39:01.926719: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:39:01.926723: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:39:01.926727: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:39:01.926731: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:39:01.926734: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:39:01.926738: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:39:01.926742: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:39:01.926746: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:39:01.926750: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:39:01.926757: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:39:01.926761: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:39:01.926765: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:39:01.926768: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:39:01.926772: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:39:01.926775: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:39:01.926780: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:01.926784: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:01.926787: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:39:01.926791: 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:39:01.926795: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:39:01.926798: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:39:01.926803: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:39:01.926807: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:39:01.926811: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:39:01.926815: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:39:01.926819: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:39:01.926822: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:39:01.926827: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:39:01.926831: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:39:01.926835: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:39:01.926838: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:39:01.926842: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:39:01.926846: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:39:01.926851: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:39:01.926855: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:39:01.926859: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:39:01.926864: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:39:01.926867: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:39:02.064027: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_fn:\n",
            "2024-05-21 16:39:02.064066: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:39:02.064071: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:39:02.064075: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:39:02.064079: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:39:02.064083: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:39:02.064087: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:39:02.064091: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:39:02.064095: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:39:02.064098: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:39:02.064102: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:39:02.064106: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:39:02.064110: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:39:02.064113: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:39:02.064117: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:39:02.064121: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:39:02.064124: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:39:02.064128: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:39:02.064131: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:39:02.064135: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:02.064139: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:39:02.064142: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:39:02.064146: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:39:02.064150: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:39:02.064153: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:39:02.064157: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:39:02.064161: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:39:02.064164: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:39:02.064168: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:39:02.064172: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:39:02.064175: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:39:02.064179: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:39:02.064183: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:39:02.064186: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:39:02.064190: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:39:02.064194: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:39:02.064200: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:39:02.064204: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:39:02.064208: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:39:02.064211: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:39:02.064215: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:39:02.064219: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:39:02.064222: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:02.064226: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:02.064230: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:39:02.064234: 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:39:02.064237: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:39:02.064241: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:39:02.064245: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:39:02.064250: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:39:02.064253: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:39:02.064257: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:39:02.064261: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:39:02.064265: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:39:02.064268: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:39:02.064272: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:39:02.064276: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:39:02.064279: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:39:02.064283: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:39:02.064287: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:39:02.064290: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:39:02.064294: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:39:02.064298: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:39:02.064301: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:39:02.064305: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:39:02.201897: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_fn:\n",
            "2024-05-21 16:39:02.201931: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:39:02.201938: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:39:02.201943: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:39:02.201947: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:39:02.201951: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:39:02.201955: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:39:02.201958: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:39:02.201962: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:39:02.201966: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:39:02.201969: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:39:02.201973: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:39:02.201976: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:39:02.201980: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:39:02.201984: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:39:02.201987: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:39:02.201991: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:39:02.201995: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:39:02.201998: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:39:02.202002: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:02.202006: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:39:02.202009: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:39:02.202013: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:39:02.202017: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:39:02.202020: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:39:02.202024: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:39:02.202028: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:39:02.202031: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:39:02.202035: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:39:02.202039: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:39:02.202043: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:39:02.202046: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:39:02.202050: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:39:02.202054: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:39:02.202057: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:39:02.202061: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:39:02.202070: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:39:02.202074: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:39:02.202079: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:39:02.202083: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:39:02.202086: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:39:02.202090: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:39:02.202094: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:02.202097: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:02.202101: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:39:02.202104: 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:39:02.202108: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:39:02.202112: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:39:02.202115: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:39:02.202119: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:39:02.202123: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:39:02.202126: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:39:02.202130: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:39:02.202134: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:39:02.202137: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:39:02.202141: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:39:02.202145: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:39:02.202148: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:39:02.202152: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:39:02.202156: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:39:02.202159: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:39:02.202163: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:39:02.202167: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:39:02.202170: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:39:02.202174: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:39:02.336680: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_fn:\n",
            "2024-05-21 16:39:02.336714: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:39:02.336719: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:39:02.336723: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:39:02.336727: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:39:02.336731: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:39:02.336735: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:39:02.336738: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:39:02.336742: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:39:02.336746: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:39:02.336749: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:39:02.336753: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:39:02.336757: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:39:02.336760: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:39:02.336764: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:39:02.336768: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:39:02.336771: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:39:02.336775: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:39:02.336779: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:39:02.336782: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:02.336786: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:39:02.336789: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:39:02.336793: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:39:02.336797: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:39:02.336800: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:39:02.336804: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:39:02.336807: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:39:02.336811: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:39:02.336815: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:39:02.336818: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:39:02.336822: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:39:02.336826: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:39:02.336829: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:39:02.336833: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:39:02.336837: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:39:02.336840: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:39:02.336848: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:39:02.336851: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:39:02.336855: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:39:02.336859: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:39:02.336862: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:39:02.336866: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:39:02.336870: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:02.336873: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:02.336877: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:39:02.336880: 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:39:02.336884: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:39:02.336888: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:39:02.336892: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:39:02.336896: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:39:02.336899: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:39:02.336903: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:39:02.336906: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:39:02.336910: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:39:02.336914: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:39:02.336917: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:39:02.336921: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:39:02.336924: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:39:02.336928: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:39:02.336932: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:39:02.336935: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:39:02.336939: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:39:02.336942: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:39:02.336946: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:39:02.336949: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:39:02.476479: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_broadcast_in_dim:\n",
            "2024-05-21 16:39:02.476512: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:39:02.476520: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:39:02.476524: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:39:02.476528: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:39:02.476532: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:39:02.476536: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:39:02.476539: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:39:02.476543: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:39:02.476547: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:39:02.476551: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:39:02.476554: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:39:02.476558: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:39:02.476561: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:39:02.476565: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:39:02.476569: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:39:02.476572: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:39:02.476576: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:39:02.476579: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:39:02.476583: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:02.476586: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:39:02.476590: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:39:02.476594: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:39:02.476597: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:39:02.476601: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:39:02.476604: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:39:02.476608: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:39:02.476612: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:39:02.476615: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:39:02.476619: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:39:02.476622: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:39:02.476626: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:39:02.476629: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:39:02.476633: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:39:02.476636: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:39:02.476640: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:39:02.476647: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:39:02.476650: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:39:02.476654: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:39:02.476657: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:39:02.476661: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:39:02.476665: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:39:02.476668: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:02.476672: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:02.476675: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:39:02.476679: 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:39:02.476683: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:39:02.476686: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:39:02.476690: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:39:02.476693: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:39:02.476697: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:39:02.476701: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:39:02.476704: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:39:02.476708: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:39:02.476711: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:39:02.476715: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:39:02.476718: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:39:02.476722: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:39:02.476726: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:39:02.476729: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:39:02.476733: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:39:02.476736: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:39:02.476740: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:39:02.476743: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:39:02.476747: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:39:02.615609: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_broadcast_in_dim:\n",
            "2024-05-21 16:39:02.615644: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:39:02.615649: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:39:02.615654: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:39:02.615658: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:39:02.615663: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:39:02.615667: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:39:02.615671: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:39:02.615674: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:39:02.615678: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:39:02.615682: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:39:02.615685: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:39:02.615689: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:39:02.615693: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:39:02.615696: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:39:02.615700: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:39:02.615704: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:39:02.615707: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:39:02.615711: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:39:02.615715: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:02.615718: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:39:02.615722: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:39:02.615725: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:39:02.615729: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:39:02.615733: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:39:02.615736: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:39:02.615740: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:39:02.615744: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:39:02.615747: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:39:02.615751: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:39:02.615754: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:39:02.615758: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:39:02.615762: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:39:02.615765: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:39:02.615769: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:39:02.615772: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:39:02.615776: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:39:02.615783: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:39:02.615786: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:39:02.615790: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:39:02.615793: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:39:02.615797: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:39:02.615801: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:02.615804: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:02.615808: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:39:02.615811: 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:39:02.615815: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:39:02.615818: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:39:02.615822: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:39:02.615826: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:39:02.615829: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:39:02.615833: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:39:02.615837: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:39:02.615840: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:39:02.615844: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:39:02.615848: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:39:02.615851: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:39:02.615855: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:39:02.615859: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:39:02.615862: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:39:02.615866: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:39:02.615869: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:39:02.615873: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:39:02.615876: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:39:02.615880: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:39:02.758171: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_fn:\n",
            "2024-05-21 16:39:02.758214: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:39:02.758222: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:39:02.758229: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:39:02.758234: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:39:02.758240: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:39:02.758246: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:39:02.758252: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:39:02.758258: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:39:02.758263: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:39:02.758270: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:39:02.758276: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:39:02.758281: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:39:02.758288: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:39:02.758293: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:39:02.758299: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:39:02.758306: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:39:02.758312: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:39:02.758317: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:39:02.758324: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:02.758329: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:39:02.758335: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:39:02.758342: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:39:02.758350: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:39:02.758355: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:39:02.758361: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:39:02.758366: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:39:02.758371: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:39:02.758377: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:39:02.758382: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:39:02.758387: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:39:02.758393: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:39:02.758399: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:39:02.758405: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:39:02.758410: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:39:02.758417: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:39:02.758423: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:39:02.758429: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:39:02.758435: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:39:02.758441: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:39:02.758446: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:39:02.758463: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:39:02.758469: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:02.758476: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:02.758484: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:39:02.758490: 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:39:02.758495: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:39:02.758500: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:39:02.758505: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:39:02.758511: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:39:02.758516: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:39:02.758522: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:39:02.758527: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:39:02.758533: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:39:02.758539: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:39:02.758544: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:39:02.758551: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:39:02.758558: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:39:02.758563: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:39:02.758569: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:39:02.758576: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:39:02.758581: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:39:02.758587: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:39:02.758592: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:39:02.758598: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:39:02.901002: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_convert_element_type:\n",
            "2024-05-21 16:39:02.901038: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:39:02.901043: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:39:02.901047: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:39:02.901051: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:39:02.901055: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:39:02.901059: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:39:02.901062: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:39:02.901066: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:39:02.901070: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:39:02.901073: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:39:02.901077: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:39:02.901081: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:39:02.901084: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:39:02.901088: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:39:02.901092: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:39:02.901095: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:39:02.901099: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:39:02.901102: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:39:02.901106: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:02.901110: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:39:02.901113: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:39:02.901117: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:39:02.901120: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:39:02.901124: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:39:02.901128: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:39:02.901131: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:39:02.901135: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:39:02.901138: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:39:02.901142: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:39:02.901146: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:39:02.901149: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:39:02.901153: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:39:02.901156: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:39:02.901160: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:39:02.901164: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:39:02.901171: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:39:02.901175: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:39:02.901178: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:39:02.901182: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:39:02.901185: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:39:02.901189: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:39:02.901192: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:02.901196: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:02.901200: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:39:02.901203: 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:39:02.901207: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:39:02.901210: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:39:02.901214: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:39:02.901217: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:39:02.901221: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:39:02.901225: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:39:02.901228: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:39:02.901232: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:39:02.901236: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:39:02.901239: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:39:02.901243: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:39:02.901246: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:39:02.901250: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:39:02.901254: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:39:02.901257: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:39:02.901261: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:39:02.901264: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:39:02.901268: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:39:02.901272: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:39:03.030678: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_broadcast_in_dim:\n",
            "2024-05-21 16:39:03.030715: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:39:03.030720: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:39:03.030724: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:39:03.030728: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:39:03.030732: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:39:03.030736: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:39:03.030740: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:39:03.030743: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:39:03.030747: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:39:03.030751: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:39:03.030754: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:39:03.030758: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:39:03.030762: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:39:03.030765: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:39:03.030769: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:39:03.030773: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:39:03.030776: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:39:03.030780: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:39:03.030784: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:03.030787: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:39:03.030791: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:39:03.030794: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:39:03.030798: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:39:03.030802: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:39:03.030805: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:39:03.030809: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:39:03.030812: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:39:03.030816: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:39:03.030820: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:39:03.030823: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:39:03.030827: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:39:03.030830: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:39:03.030834: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:39:03.030837: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:39:03.030841: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:39:03.030849: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:39:03.030853: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:39:03.030857: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:39:03.030860: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:39:03.030864: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:39:03.030867: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:39:03.030871: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:03.030875: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:03.030878: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:39:03.030882: 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:39:03.030885: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:39:03.030889: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:39:03.030893: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:39:03.030896: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:39:03.030900: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:39:03.030903: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:39:03.030907: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:39:03.030911: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:39:03.030914: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:39:03.030918: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:39:03.030921: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:39:03.030925: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:39:03.030929: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:39:03.030932: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:39:03.030936: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:39:03.030939: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:39:03.030943: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:39:03.030947: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:39:03.030950: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:39:03.523264: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_scan:\n",
            "2024-05-21 16:39:03.523299: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:39:03.523304: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:39:03.523308: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:39:03.523312: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:39:03.523316: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:39:03.523320: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:39:03.523323: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:39:03.523327: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:39:03.523331: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:39:03.523334: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:39:03.523338: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:39:03.523342: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:39:03.523346: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:39:03.523349: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:39:03.523353: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:39:03.523357: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:39:03.523360: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:39:03.523364: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:39:03.523368: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:03.523371: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:39:03.523375: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:39:03.523379: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:39:03.523382: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:39:03.523386: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:39:03.523389: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:39:03.523393: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:39:03.523397: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:39:03.523400: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:39:03.523404: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:39:03.523407: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:39:03.523411: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:39:03.523415: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:39:03.523419: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:39:03.523422: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:39:03.523426: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:39:03.523431: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:39:03.523435: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:39:03.523438: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:39:03.523442: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:39:03.523445: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:39:03.523455: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:39:03.523460: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:03.523464: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:03.523467: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:39:03.523471: 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:39:03.523474: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:39:03.523478: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:39:03.523481: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:39:03.523485: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:39:03.523488: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:39:03.523492: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:39:03.523496: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:39:03.523499: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:39:03.523503: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:39:03.523506: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:39:03.523510: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:39:03.523513: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_detailed_logging: true\n",
            "2024-05-21 16:39:03.523517: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:39:03.523521: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:39:03.523524: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:39:03.523528: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:39:03.523531: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:39:03.523535: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:39:03.523539: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:39:03.523542: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:39:03.856154: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit__pad:\n",
            "2024-05-21 16:39:03.856190: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:39:03.856196: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:39:03.856200: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:39:03.856204: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:39:03.856208: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:39:03.856212: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:39:03.856215: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:39:03.856219: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:39:03.856223: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:39:03.856226: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:39:03.856230: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:39:03.856234: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:39:03.856237: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:39:03.856241: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:39:03.856245: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:39:03.856248: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:39:03.856252: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:39:03.856256: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:39:03.856259: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:03.856263: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:39:03.856266: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:39:03.856270: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:39:03.856274: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:39:03.856277: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:39:03.856281: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:39:03.856284: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:39:03.856288: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:39:03.856292: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:39:03.856295: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:39:03.856299: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:39:03.856302: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:39:03.856306: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:39:03.856310: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:39:03.856313: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:39:03.856317: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:39:03.856326: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:39:03.856331: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:39:03.856337: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:39:03.856343: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:39:03.856352: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:39:03.856361: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:39:03.856369: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:03.856379: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:03.856386: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:39:03.856391: 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:39:03.856395: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:39:03.856399: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:39:03.856402: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:39:03.856406: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:39:03.856410: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:39:03.856413: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:39:03.856417: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:39:03.856421: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:39:03.856424: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:39:03.856428: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:39:03.856431: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:39:03.856435: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:39:03.856439: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:39:03.856442: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:39:03.856446: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:39:03.856457: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:39:03.856461: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:39:03.856465: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:39:03.856468: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:39:03.994744: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_slice:\n",
            "2024-05-21 16:39:03.994781: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:39:03.994786: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:39:03.994791: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:39:03.994795: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:39:03.994799: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:39:03.994802: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:39:03.994806: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:39:03.994810: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:39:03.994813: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:39:03.994817: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:39:03.994820: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:39:03.994824: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:39:03.994828: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:39:03.994831: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:39:03.994835: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:39:03.994839: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:39:03.994843: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:39:03.994846: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:39:03.994850: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:03.994853: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:39:03.994857: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:39:03.994861: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:39:03.994864: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:39:03.994868: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:39:03.994871: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:39:03.994875: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:39:03.994879: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:39:03.994882: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:39:03.994886: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:39:03.994889: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:39:03.994893: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:39:03.994897: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:39:03.994900: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:39:03.994904: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:39:03.994908: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:39:03.994917: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:39:03.994921: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:39:03.994924: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:39:03.994928: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:39:03.994931: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:39:03.994935: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:39:03.994939: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:03.994942: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:03.994946: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:39:03.994949: 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:39:03.994953: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:39:03.994956: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:39:03.994960: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:39:03.994964: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:39:03.994967: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:39:03.994971: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:39:03.994975: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:39:03.994978: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:39:03.994982: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:39:03.994986: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:39:03.994990: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:39:03.994993: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:39:03.994997: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:39:03.995000: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:39:03.995004: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:39:03.995008: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:39:03.995011: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:39:03.995015: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:39:03.995018: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:39:04.129419: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit__lambda_:\n",
            "2024-05-21 16:39:04.129461: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:39:04.129466: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:39:04.129470: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:39:04.129474: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:39:04.129478: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:39:04.129482: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:39:04.129486: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:39:04.129489: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:39:04.129493: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:39:04.129496: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:39:04.129500: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:39:04.129504: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:39:04.129507: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:39:04.129511: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:39:04.129515: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:39:04.129518: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:39:04.129522: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:39:04.129526: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:39:04.129529: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:04.129533: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:39:04.129536: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:39:04.129540: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:39:04.129543: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:39:04.129547: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:39:04.129551: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:39:04.129554: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:39:04.129558: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:39:04.129562: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:39:04.129565: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:39:04.129569: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:39:04.129572: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:39:04.129576: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:39:04.129579: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:39:04.129583: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:39:04.129587: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:39:04.129597: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:39:04.129600: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:39:04.129604: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:39:04.129608: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:39:04.129611: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:39:04.129615: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:39:04.129618: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:04.129622: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:04.129626: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:39:04.129629: 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:39:04.129633: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:39:04.129637: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:39:04.129640: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:39:04.129644: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:39:04.129648: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:39:04.129652: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:39:04.129655: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:39:04.129659: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:39:04.129663: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:39:04.129667: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:39:04.129671: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:39:04.129674: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:39:04.129678: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:39:04.129682: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:39:04.129685: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:39:04.129689: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:39:04.129693: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:39:04.129696: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:39:04.129700: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:39:04.268692: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_fn:\n",
            "2024-05-21 16:39:04.268730: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:39:04.268735: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:39:04.268739: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:39:04.268743: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:39:04.268746: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:39:04.268750: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:39:04.268754: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:39:04.268757: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:39:04.268761: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:39:04.268765: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:39:04.268769: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:39:04.268772: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:39:04.268776: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:39:04.268780: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:39:04.268783: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:39:04.268787: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:39:04.268791: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:39:04.268794: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:39:04.268798: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:04.268801: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:39:04.268805: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:39:04.268808: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:39:04.268812: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:39:04.268816: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:39:04.268819: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:39:04.268823: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:39:04.268827: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:39:04.268830: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:39:04.268834: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:39:04.268837: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:39:04.268841: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:39:04.268845: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:39:04.268848: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:39:04.268853: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:39:04.268856: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:39:04.268863: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:39:04.268867: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:39:04.268871: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:39:04.268874: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:39:04.268878: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:39:04.268882: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:39:04.268885: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:04.268889: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:04.268892: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:39:04.268896: 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:39:04.268899: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:39:04.268903: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:39:04.268907: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:39:04.268910: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:39:04.268914: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:39:04.268917: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:39:04.268921: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:39:04.268925: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:39:04.268928: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:39:04.268932: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:39:04.268936: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:39:04.268939: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:39:04.268943: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:39:04.268947: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:39:04.268950: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:39:04.268954: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:39:04.268957: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:39:04.268961: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:39:04.268965: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:39:04.412291: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_fn:\n",
            "2024-05-21 16:39:04.412326: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:39:04.412332: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:39:04.412336: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:39:04.412340: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:39:04.412344: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:39:04.412348: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:39:04.412351: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:39:04.412355: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:39:04.412358: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:39:04.412362: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:39:04.412366: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:39:04.412369: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:39:04.412373: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:39:04.412377: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:39:04.412381: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:39:04.412384: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:39:04.412388: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:39:04.412391: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:39:04.412395: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:04.412399: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:39:04.412402: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:39:04.412406: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:39:04.412409: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:39:04.412413: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:39:04.412417: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:39:04.412420: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:39:04.412424: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:39:04.412428: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:39:04.412431: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:39:04.412435: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:39:04.412438: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:39:04.412442: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:39:04.412446: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:39:04.412456: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:39:04.412460: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:39:04.412467: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:39:04.412470: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:39:04.412474: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:39:04.412477: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:39:04.412481: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:39:04.412485: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:39:04.412488: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:04.412492: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:04.412495: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:39:04.412499: 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:39:04.412502: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:39:04.412506: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:39:04.412510: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:39:04.412513: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:39:04.412517: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:39:04.412521: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:39:04.412524: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:39:04.412528: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:39:04.412532: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:39:04.412535: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:39:04.412539: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:39:04.412542: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:39:04.412546: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:39:04.412550: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:39:04.412553: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:39:04.412557: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:39:04.412560: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:39:04.412564: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:39:04.412567: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:39:04.443888: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_true_divide:\n",
            "2024-05-21 16:39:04.443921: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:39:04.443926: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:39:04.443931: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:39:04.443935: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:39:04.443939: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:39:04.443942: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:39:04.443946: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:39:04.443950: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:39:04.443954: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:39:04.443957: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:39:04.443961: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:39:04.443964: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:39:04.443968: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:39:04.443972: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:39:04.443976: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:39:04.443980: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:39:04.443983: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:39:04.443987: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:39:04.443991: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:04.443994: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:39:04.443998: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:39:04.444002: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:39:04.444006: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:39:04.444009: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:39:04.444013: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:39:04.444017: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:39:04.444020: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:39:04.444024: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:39:04.444028: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:39:04.444031: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:39:04.444035: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:39:04.444039: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:39:04.444042: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:39:04.444046: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:39:04.444050: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:39:04.444056: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:39:04.444060: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:39:04.444064: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:39:04.444068: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:39:04.444071: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:39:04.444075: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:39:04.444078: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:04.444082: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:04.444086: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:39:04.444089: 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:39:04.444093: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:39:04.444097: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:39:04.444100: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:39:04.444104: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:39:04.444108: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:39:04.444111: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:39:04.444115: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:39:04.444119: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:39:04.444123: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:39:04.444126: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:39:04.444130: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:39:04.444134: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:39:04.444137: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:39:04.444141: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:39:04.444145: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:39:04.444148: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:39:04.444152: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:39:04.444156: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:39:04.444159: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:39:04.729074: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit__pad:\n",
            "2024-05-21 16:39:04.729110: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:39:04.729115: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:39:04.729120: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:39:04.729124: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:39:04.729127: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:39:04.729131: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:39:04.729135: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:39:04.729139: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:39:04.729142: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:39:04.729146: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:39:04.729149: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:39:04.729153: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:39:04.729157: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:39:04.729160: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:39:04.729164: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:39:04.729168: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:39:04.729172: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:39:04.729175: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:39:04.729179: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:04.729182: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:39:04.729186: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:39:04.729190: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:39:04.729193: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:39:04.729197: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:39:04.729201: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:39:04.729204: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:39:04.729208: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:39:04.729212: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:39:04.729215: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:39:04.729219: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:39:04.729223: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:39:04.729226: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:39:04.729230: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:39:04.729233: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:39:04.729237: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:39:04.729246: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:39:04.729250: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:39:04.729253: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:39:04.729257: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:39:04.729260: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:39:04.729264: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:39:04.729268: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:04.729271: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:04.729275: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:39:04.729278: 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:39:04.729282: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:39:04.729286: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:39:04.729289: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:39:04.729293: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:39:04.729296: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:39:04.729300: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:39:04.729304: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:39:04.729307: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:39:04.729311: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:39:04.729315: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:39:04.729319: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:39:04.729322: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_detailed_logging: true\n",
            "2024-05-21 16:39:04.729326: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:39:04.729330: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:39:04.729333: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:39:04.729337: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:39:04.729340: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:39:04.729344: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:39:04.729348: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:39:04.729351: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n",
            "2024-05-21 16:39:04.868298: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_slice:\n",
            "2024-05-21 16:39:04.868335: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:39:04.868341: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:39:04.868345: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:39:04.868349: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:39:04.868353: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:39:04.868357: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:39:04.868360: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:39:04.868364: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:39:04.868368: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:39:04.868371: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:39:04.868375: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:39:04.868379: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:39:04.868383: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:39:04.868386: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:39:04.868390: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:39:04.868394: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:39:04.868397: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:39:04.868401: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:39:04.868405: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:04.868408: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:39:04.868412: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:39:04.868416: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:39:04.868419: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:39:04.868423: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:39:04.868426: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:39:04.868430: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:39:04.868434: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:39:04.868437: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:39:04.868441: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:39:04.868445: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:39:04.868448: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:39:04.868460: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:39:04.868463: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:39:04.868467: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:39:04.868471: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:39:04.868478: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:39:04.868482: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:39:04.868486: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:39:04.868489: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:39:04.868493: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:39:04.868496: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:39:04.868500: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:04.868504: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:39:04.868507: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:39:04.868511: 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:39:04.868515: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:39:04.868518: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:39:04.868522: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:39:04.868526: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:39:04.868529: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:39:04.868533: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:39:04.868537: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:39:04.868540: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:39:04.868544: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:39:04.868548: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:39:04.868551: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:39:04.868555: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:39:04.868558: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:39:04.868562: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:39:04.868566: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:39:04.868569: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:39:04.868573: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:39:04.868576: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:39:04.868580: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n"
          ]
        }
      ],
      "source": [
        "size = (512, 128)\n",
        "domain = ((0, 8), (0, 2))\n",
        "density = 1.\n",
        "viscosity = 1e-3  # kinematic visocity\n",
        "pressure_gradient = 2e-3  # uniform dP/dx\n",
        "\n",
        "# Define the physical dimensions of the simulation.\n",
        "grid = cfd.grids.Grid(size, domain=domain)\n",
        "\n",
        "# Specify periodic BC on x-boundaries, and no-slip walls on the y-boundaries.\n",
        "velocity_bc = (cfd.boundaries.channel_flow_boundary_conditions(ndim=2),\n",
        "               cfd.boundaries.channel_flow_boundary_conditions(ndim=2))\n",
        "\n",
        "# Specify zero velocity initial conditions\n",
        "vx_fn = lambda x, y: jnp.zeros_like(x + y)\n",
        "vy_fn = lambda x, y: jnp.zeros_like(x + y)\n",
        "\n",
        "v0 = cfd.initial_conditions.initial_velocity_field(\n",
        "    velocity_fns=(vx_fn, vy_fn),\n",
        "    grid=grid,\n",
        "    velocity_bc=velocity_bc,\n",
        "    pressure_solve=cfd.pressure.solve_fast_diag_channel_flow,\n",
        "    iterations=5)\n",
        "\n",
        "# Confirm that the divergence of v0 is zero\n",
        "div = cfd.finite_differences.divergence(v0)\n",
        "np.testing.assert_allclose(div.data, 0)\n",
        "\n",
        "# Specify a fixed time step based on the convection and diffusion scales\n",
        "max_velocity = 1  # value selected from known equilibirum profile\n",
        "cfl_safety_factor = 0.5\n",
        "\n",
        "dt = cfd.equations.stable_time_step(\n",
        "    max_velocity, cfl_safety_factor, viscosity, grid)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "executionInfo": {
          "elapsed": 2,
          "status": "ok",
          "timestamp": 1655317720089,
          "user": {
            "displayName": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "oZokhSc6VMKv"
      },
      "outputs": [],
      "source": [
        "# Define a linear force due to the uniform pressure gradient\n",
        "def pressure_gradient_forcing(pressure_gradient: float):\n",
        "  \n",
        "  def forcing(v):\n",
        "    force_vector = (pressure_gradient, 0)\n",
        "    return tuple(cfd.grids.GridArray(f * jnp.ones_like(u.data), u.offset, u.grid)\n",
        "                 for f, u in zip(force_vector, v))\n",
        "  \n",
        "  return forcing"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "executionInfo": {
          "elapsed": 12362,
          "status": "ok",
          "timestamp": 1655317733789,
          "user": {
            "displayName": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "_5YttC-TVxk1",
        "outputId": "b4e346c5-365c-48fc-e946-5957872f9ea1"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "2024-05-21 16:44:43.173526: W external/xla/xla/service/gpu/gpu_compiler.cc:555] GpuCompilationEnvironment of hlo_module jit_multistep:\n",
            "2024-05-21 16:44:43.173565: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_backend_optimization_level: 3\n",
            "2024-05-21 16:44:43.173614: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_eliminate_hlo_implicit_broadcast: true\n",
            "2024-05-21 16:44:43.173619: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_multi_thread_eigen: true\n",
            "2024-05-21 16:44:43.173624: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cuda_data_dir: \"./cuda_sdk_lib\"\n",
            "2024-05-21 16:44:43.173628: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_alias_scope_metadata: true\n",
            "2024-05-21 16:44:43.173632: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_noalias_metadata: true\n",
            "2024-05-21 16:44:43.173635: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_llvm_enable_invariant_load_metadata: true\n",
            "2024-05-21 16:44:43.173639: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_force_host_platform_device_count: 1\n",
            "2024-05-21 16:44:43.173643: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_nans: true\n",
            "2024-05-21 16:44:43.173647: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_infs: true\n",
            "2024-05-21 16:44:43.173650: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_allow_excess_precision: true\n",
            "2024-05-21 16:44:43.173654: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_autotune_level: 4\n",
            "2024-05-21 16:44:43.173657: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_division: true\n",
            "2024-05-21 16:44:43.173661: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_fast_math_honor_functions: true\n",
            "2024-05-21 16:44:43.173665: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_max_hlo_modules: -1\n",
            "2024-05-21 16:44:43.173668: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_multiheap_size_constraint_per_heap: -1\n",
            "2024-05-21 16:44:43.173672: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_async_all_reduce: true\n",
            "2024-05-21 16:44:43.173676: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_strict_conv_algorithm_picker: true\n",
            "2024-05-21 16:44:43.173679: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_reduce_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:44:43.173683: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_frontend: true\n",
            "2024-05-21 16:44:43.173686: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_nccl_termination_timeout_seconds: -1\n",
            "2024-05-21 16:44:43.173690: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_shared_constants: true\n",
            "2024-05-21 16:44:43.173694: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_scratch_max_megabytes: 4096\n",
            "2024-05-21 16:44:43.173697: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_simplify_all_fp_conversions: true\n",
            "2024-05-21 16:44:43.173701: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_xla_runtime_executable: true\n",
            "2024-05-21 16:44:43.173704: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_shape_checks: RUNTIME\n",
            "2024-05-21 16:44:43.173708: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_normalize_layouts: true\n",
            "2024-05-21 16:44:43.173711: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_tiling_and_fusion: true\n",
            "2024-05-21 16:44:43.173715: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_dump_enable_mlir_pretty_form: true\n",
            "2024-05-21 16:44:43.173718: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_triton_gemm: true\n",
            "2024-05-21 16:44:43.173723: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cudnn_int8x32_convolution_reordering: true\n",
            "2024-05-21 16:44:43.173726: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_experimental_deallocation: true\n",
            "2024-05-21 16:44:43.173730: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_enable_mlir_fusion_outlining: true\n",
            "2024-05-21 16:44:43.173734: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_m_dim: 8\n",
            "2024-05-21 16:44:43.173737: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_n_dim: 8\n",
            "2024-05-21 16:44:43.173746: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_cpu_matmul_tiling_k_dim: 8\n",
            "2024-05-21 16:44:43.173749: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_num_runs_to_instantiate: -1\n",
            "2024-05-21 16:44:43.173753: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_lhs_enable_gpu_async_tracker: true\n",
            "2024-05-21 16:44:43.173756: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_inflation_factor: 1\n",
            "2024-05-21 16:44:43.173760: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_min_graph_size: 5\n",
            "2024-05-21 16:44:43.173763: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reassociation_for_converted_ar: true\n",
            "2024-05-21 16:44:43.173767: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_all_gather_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:44:43.173771: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_reduce_scatter_combine_threshold_bytes: 31457280\n",
            "2024-05-21 16:44:43.173774: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_highest_priority_async_stream: true\n",
            "2024-05-21 16:44:43.173778: 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:44:43.173781: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_redzone_padding_bytes: 8388608\n",
            "2024-05-21 16:44:43.173785: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_triton_fusion_level: 2\n",
            "2024-05-21 16:44:43.173788: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_graph_eviction_timeout_seconds: 60\n",
            "2024-05-21 16:44:43.173792: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_gpu2_hal: true\n",
            "2024-05-21 16:44:43.173795: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_copy_insertion_use_region_analysis: true\n",
            "2024-05-21 16:44:43.173799: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807\n",
            "2024-05-21 16:44:43.173803: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_split_k_autotuning: true\n",
            "2024-05-21 16:44:43.173807: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduction_epilogue_fusion: true\n",
            "2024-05-21 16:44:43.173810: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_cublas_fallback: true\n",
            "2024-05-21 16:44:43.173814: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_filter_kernels_spilling_registers_on_autotuning: true\n",
            "2024-05-21 16:44:43.173817: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_debug_buffer_assignment_show_max: 15\n",
            "2024-05-21 16:44:43.173821: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_detailed_logging: true\n",
            "2024-05-21 16:44:43.173824: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_enable_dumping: true\n",
            "2024-05-21 16:44:43.173828: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_all_gather_combine_by_dim: true\n",
            "2024-05-21 16:44:43.173831: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_reduce_scatter_combine_by_dim: true\n",
            "2024-05-21 16:44:43.173835: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: FUSION\n",
            "2024-05-21 16:44:43.173839: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_command_buffer: CUBLAS\n",
            "2024-05-21 16:44:43.173842: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_enable_cub_radix_sort: true\n",
            "2024-05-21 16:44:43.173846: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_memory_limit_slop_factor: 95\n",
            "2024-05-21 16:44:43.173849: W external/xla/xla/service/gpu/gpu_compiler.cc:555] xla_gpu_threshold_for_windowed_einsum_mib: 100000\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "CPU times: user 23.2 s, sys: 12.1 s, total: 35.3 s\n",
            "Wall time: 21.5 s\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "array([[[0.        , 0.        , 0.        , ..., 0.        ,\n",
              "         0.        , 0.        ],\n",
              "        [0.        , 0.        , 0.        , ..., 0.        ,\n",
              "         0.        , 0.        ],\n",
              "        [0.        , 0.        , 0.        , ..., 0.        ,\n",
              "         0.        , 0.        ],\n",
              "        ...,\n",
              "        [0.        , 0.        , 0.        , ..., 0.        ,\n",
              "         0.        , 0.        ],\n",
              "        [0.        , 0.        , 0.        , ..., 0.        ,\n",
              "         0.        , 0.        ],\n",
              "        [0.        , 0.        , 0.        , ..., 0.        ,\n",
              "         0.        , 0.        ]],\n",
              "\n",
              "       [[0.00492697, 0.01430035, 0.02320853, ..., 0.02320853,\n",
              "         0.01430035, 0.00492697],\n",
              "        [0.00492697, 0.01430035, 0.02320853, ..., 0.02320853,\n",
              "         0.01430035, 0.00492697],\n",
              "        [0.00492697, 0.01430035, 0.02320853, ..., 0.02320853,\n",
              "         0.01430035, 0.00492697],\n",
              "        ...,\n",
              "        [0.00492697, 0.01430035, 0.02320853, ..., 0.02320853,\n",
              "         0.01430035, 0.00492697],\n",
              "        [0.00492697, 0.01430035, 0.02320853, ..., 0.02320853,\n",
              "         0.01430035, 0.00492697],\n",
              "        [0.00492697, 0.01430035, 0.02320853, ..., 0.02320853,\n",
              "         0.01430035, 0.00492697]],\n",
              "\n",
              "       [[0.00696648, 0.02041658, 0.03339467, ..., 0.03339467,\n",
              "         0.02041658, 0.00696648],\n",
              "        [0.00696648, 0.02041658, 0.03339467, ..., 0.03339467,\n",
              "         0.02041658, 0.00696648],\n",
              "        [0.00696648, 0.02041658, 0.03339467, ..., 0.03339467,\n",
              "         0.02041658, 0.00696648],\n",
              "        ...,\n",
              "        [0.00696648, 0.02041658, 0.03339467, ..., 0.03339467,\n",
              "         0.02041658, 0.00696648],\n",
              "        [0.00696648, 0.02041658, 0.03339467, ..., 0.03339467,\n",
              "         0.02041658, 0.00696648],\n",
              "        [0.00696648, 0.02041658, 0.03339467, ..., 0.03339467,\n",
              "         0.02041658, 0.00696648]],\n",
              "\n",
              "       ...,\n",
              "\n",
              "       [[0.01338679, 0.03967345, 0.06547584, ..., 0.06547584,\n",
              "         0.03967345, 0.01338679],\n",
              "        [0.01338679, 0.03967345, 0.06547584, ..., 0.06547584,\n",
              "         0.03967345, 0.01338679],\n",
              "        [0.01338679, 0.03967345, 0.06547584, ..., 0.06547584,\n",
              "         0.03967345, 0.01338679],\n",
              "        ...,\n",
              "        [0.01338679, 0.03967345, 0.06547584, ..., 0.06547584,\n",
              "         0.03967345, 0.01338679],\n",
              "        [0.01338679, 0.03967345, 0.06547584, ..., 0.06547584,\n",
              "         0.03967345, 0.01338679],\n",
              "        [0.01338679, 0.03967345, 0.06547584, ..., 0.06547584,\n",
              "         0.03967345, 0.01338679]],\n",
              "\n",
              "       [[0.01377827, 0.04084764, 0.06743206, ..., 0.06743206,\n",
              "         0.04084764, 0.01377827],\n",
              "        [0.01377827, 0.04084764, 0.06743206, ..., 0.06743206,\n",
              "         0.04084764, 0.01377827],\n",
              "        [0.01377827, 0.04084764, 0.06743206, ..., 0.06743206,\n",
              "         0.04084764, 0.01377827],\n",
              "        ...,\n",
              "        [0.01377827, 0.04084764, 0.06743206, ..., 0.06743206,\n",
              "         0.04084764, 0.01377827],\n",
              "        [0.01377827, 0.04084764, 0.06743206, ..., 0.06743206,\n",
              "         0.04084764, 0.01377827],\n",
              "        [0.01377827, 0.04084764, 0.06743206, ..., 0.06743206,\n",
              "         0.04084764, 0.01377827]],\n",
              "\n",
              "       [[0.01410103, 0.04181574, 0.06904502, ..., 0.06904502,\n",
              "         0.04181574, 0.01410103],\n",
              "        [0.01410103, 0.04181574, 0.06904502, ..., 0.06904502,\n",
              "         0.04181574, 0.01410103],\n",
              "        [0.01410103, 0.04181574, 0.06904502, ..., 0.06904502,\n",
              "         0.04181574, 0.01410103],\n",
              "        ...,\n",
              "        [0.01410103, 0.04181574, 0.06904502, ..., 0.06904502,\n",
              "         0.04181574, 0.01410103],\n",
              "        [0.01410103, 0.04181574, 0.06904502, ..., 0.06904502,\n",
              "         0.04181574, 0.01410103],\n",
              "        [0.01410103, 0.04181574, 0.06904502, ..., 0.06904502,\n",
              "         0.04181574, 0.01410103]]], dtype=float32)"
            ]
          },
          "execution_count": 9,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# time steps per output\n",
        "inner_steps = 10_000\n",
        "\n",
        "# number of outputs\n",
        "outer_steps = 12\n",
        "\n",
        "# TODO(pnorgaard): Fix Van-Leer limiters for Dirichlet velocity BC \n",
        "def convect(v):\n",
        "  return tuple(\n",
        "      cfd.advection.advect_van_leer(u, v, dt) for u in v)\n",
        "  \n",
        "# Define a step function and use it to compute a trajectory.\n",
        "step_fn = cfd.funcutils.repeated(\n",
        "    cfd.equations.semi_implicit_navier_stokes(\n",
        "        density=density,\n",
        "        viscosity=viscosity,\n",
        "        dt=dt,\n",
        "        grid=grid,\n",
        "        convect=convect,\n",
        "        pressure_solve=cfd.pressure.solve_fast_diag_channel_flow,\n",
        "        forcing=pressure_gradient_forcing(pressure_gradient),\n",
        "        ),\n",
        "    steps=inner_steps)\n",
        "rollout_fn = jax.jit(cfd.funcutils.trajectory(\n",
        "    step_fn, outer_steps, start_with_input=True))\n",
        "\n",
        "%time _, trajectory = jax.device_get(rollout_fn(v0))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 30,
      "metadata": {
        "colab": {
          "height": 573
        },
        "executionInfo": {
          "elapsed": 1812,
          "status": "ok",
          "timestamp": 1655317762576,
          "user": {
            "displayName": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "htvk6BN2XB-H",
        "outputId": "6c383e03-538f-491d-d772-ca1d761b1e8f"
      },
      "outputs": [
        {
          "data": {
            "image/png": "[encoded data, 92304 characters]",
            "text/plain": [
              "<Figure size 2500x900 with 13 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "# load into xarray for visualization and analysis\n",
        "\n",
        "ds = xarray.Dataset(\n",
        "    {\n",
        "        'u': (('time', 'x', 'y'), trajectory[0].data),\n",
        "        'v': (('time', 'x', 'y'), trajectory[1].data),\n",
        "    },\n",
        "    coords={\n",
        "        'x': grid.axes()[0],\n",
        "        'y': grid.axes()[1],\n",
        "        'time': dt * inner_steps * np.arange(1, outer_steps + 1)\n",
        "    }\n",
        ")\n",
        "\n",
        "# # Plot the x-velocity\n",
        "ds.pipe(lambda ds: ds.u).plot.imshow(x='x', y='y', col='time',cmap=seaborn.cm.rocket, robust=True, col_wrap=4, aspect=2);"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 31,
      "metadata": {
        "colab": {
          "height": 295
        },
        "executionInfo": {
          "elapsed": 392,
          "status": "ok",
          "timestamp": 1655317767581,
          "user": {
            "displayName": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "8V5D8pXRV-8Y",
        "outputId": "1dca792a-b2ba-4ddb-9835-2abe11b185ea"
      },
      "outputs": [
        {
          "data": {
            "image/png": "[encoded data, 131896 characters]",
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "def vx_at_x_upper_boundary(ds):\n",
        "  return ds.u.isel(x=-1)  # slice operator\n",
        "\n",
        "p = ds.pipe(vx_at_x_upper_boundary).plot(hue='time')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kS4FNx0cPiVA"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "last_runtime": {
        "build_target": "//learning/deepmind/public/tools/ml_python:ml_notebook",
        "kind": "private"
      },
      "name": "JAX-CFD channel flow demo.ipynb",
      "provenance": [
        {
          "file_id": "/piper/depot/google3/research/simulation/whirl/cfd/notebooks/whirl_demo.ipynb",
          "timestamp": 1594817389794
        },
        {
          "file_id": "1YsLs93M4TVtYHOOMilpryHqr1xV2N7Vx",
          "timestamp": 1591818482026
        },
        {
          "file_id": "/piper/depot/google3/research/simulation/whirl/cfd/notebooks/whirl_demo.ipynb",
          "timestamp": 1591811585084
        },
        {
          "file_id": "1E_FfLTw46OmO9elqraoSSlYJW9r1Gsav",
          "timestamp": 1584980552314
        }
      ]
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
