diff --git a/RELEASE.md b/RELEASE.md index 1c82b90434a66d..32c40bba961f9a 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -46,6 +46,14 @@ In `tensorflow/c/experimental/filesystem/filesystem_interface.h`, removed `TF_Tr previously failed with a lookup error because the `SoftsignGrad` backward op had no registered Python gradient. +* `tf.math.reciprocal` + + * Constrains the XLA registration of `Reciprocal` and `Inv` to the types + that have a device kernel, so `jit_compile=True` no longer silently + accepts the integer inputs that eager execution and autoclustering + reject. Fixes + [#126414](https://github.com/tensorflow/tensorflow/issues/126414). + * `tf.experimental.numpy` diff --git a/ci/official/envs/linux_x86_cuda b/ci/official/envs/linux_x86_cuda index dd93e2f4ded292..5e6ce4d7bef58d 100644 --- a/ci/official/envs/linux_x86_cuda +++ b/ci/official/envs/linux_x86_cuda @@ -14,7 +14,7 @@ # ============================================================================== source ci/official/envs/linux_x86 export TF_FORCE_GPU_ALLOW_GROWTH=true -TFCI_BAZEL_COMMON_ARGS="--repo_env=HERMETIC_PYTHON_VERSION=$TFCI_PYTHON_VERSION --repo_env=USE_PYWRAP_RULES=True --config release_gpu_linux --test_env=TF_FORCE_GPU_ALLOW_GROWTH=true" +TFCI_BAZEL_COMMON_ARGS="--repo_env=HERMETIC_PYTHON_VERSION=$TFCI_PYTHON_VERSION --repo_env=USE_PYWRAP_RULES=True --config release_gpu_linux --test_env=TF_FORCE_GPU_ALLOW_GROWTH=true --local_test_jobs=16" TFCI_BAZEL_HERMETIC_CUDA_UMD_ENABLE=1 TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_cuda TFCI_BUILD_PIP_PACKAGE_WHEEL_NAME_ARG="--repo_env=WHEEL_NAME=tensorflow" diff --git a/ci/official/envs/linux_x86_cuda13_nvcc b/ci/official/envs/linux_x86_cuda13_nvcc index a8306754723f4f..bc8fc063157615 100644 --- a/ci/official/envs/linux_x86_cuda13_nvcc +++ b/ci/official/envs/linux_x86_cuda13_nvcc @@ -13,7 +13,7 @@ # limitations under the License. # ============================================================================== source ci/official/envs/linux_x86 -TFCI_BAZEL_COMMON_ARGS="--repo_env=HERMETIC_PYTHON_VERSION=$TFCI_PYTHON_VERSION --repo_env=USE_PYWRAP_RULES=True --config release_gpu_linux --config=cuda_nvcc --config=cuda13_version" +TFCI_BAZEL_COMMON_ARGS="--repo_env=HERMETIC_PYTHON_VERSION=$TFCI_PYTHON_VERSION --repo_env=USE_PYWRAP_RULES=True --config release_gpu_linux --config=cuda_nvcc --config=cuda13_version --local_test_jobs=16" TFCI_BAZEL_HERMETIC_CUDA_UMD_ENABLE=1 TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_cuda_13_nvcc TFCI_BUILD_PIP_PACKAGE_WHEEL_NAME_ARG="--repo_env=WHEEL_NAME=tensorflow_cuda13" diff --git a/requirements_lock_3_10.txt b/requirements_lock_3_10.txt index e2298fa546a5bb..a0fe535475f49a 100644 --- a/requirements_lock_3_10.txt +++ b/requirements_lock_3_10.txt @@ -1077,9 +1077,9 @@ zstandard==0.25.0 \ # via -r ci/official/requirements_updater/requirements.in # The following packages are considered to be unsafe in a requirements file: -setuptools==78.1.1 \ - --hash=sha256:c3a9c4211ff4c309edb8b8c4f1cbfa7ae324c4ba9f91ff254e3d305b9fd54561 \ - --hash=sha256:fcc17fd9cd898242f6b4adfaca46137a9edef687f43e6f78469692a5e70d851d +setuptools==83.0.0 \ + --hash=sha256:025bccbbf0fa05b6192bc64ae1e7b16e001fd6d6d4d5de03c97b1c1ade523bef \ + --hash=sha256:29b23c360f22f414dc7336bb39178cc7bcbf6021ed2733cde173f09dba19abb3 # via # -r ci/official/requirements_updater/requirements.in # tb-nightly diff --git a/requirements_lock_3_11.txt b/requirements_lock_3_11.txt index df484dd450e71a..e2be36f47b976b 100644 --- a/requirements_lock_3_11.txt +++ b/requirements_lock_3_11.txt @@ -1076,9 +1076,9 @@ zstandard==0.25.0 \ # via -r ci/official/requirements_updater/requirements.in # The following packages are considered to be unsafe in a requirements file: -setuptools==78.1.1 \ - --hash=sha256:c3a9c4211ff4c309edb8b8c4f1cbfa7ae324c4ba9f91ff254e3d305b9fd54561 \ - --hash=sha256:fcc17fd9cd898242f6b4adfaca46137a9edef687f43e6f78469692a5e70d851d +setuptools==83.0.0 \ + --hash=sha256:025bccbbf0fa05b6192bc64ae1e7b16e001fd6d6d4d5de03c97b1c1ade523bef \ + --hash=sha256:29b23c360f22f414dc7336bb39178cc7bcbf6021ed2733cde173f09dba19abb3 # via # -r ci/official/requirements_updater/requirements.in # tb-nightly diff --git a/requirements_lock_3_12.txt b/requirements_lock_3_12.txt index af658e13cac1e0..d50e6bea7c6ae4 100644 --- a/requirements_lock_3_12.txt +++ b/requirements_lock_3_12.txt @@ -1076,9 +1076,9 @@ zstandard==0.25.0 \ # via -r ci/official/requirements_updater/requirements.in # The following packages are considered to be unsafe in a requirements file: -setuptools==78.1.1 \ - --hash=sha256:c3a9c4211ff4c309edb8b8c4f1cbfa7ae324c4ba9f91ff254e3d305b9fd54561 \ - --hash=sha256:fcc17fd9cd898242f6b4adfaca46137a9edef687f43e6f78469692a5e70d851d +setuptools==83.0.0 \ + --hash=sha256:025bccbbf0fa05b6192bc64ae1e7b16e001fd6d6d4d5de03c97b1c1ade523bef \ + --hash=sha256:29b23c360f22f414dc7336bb39178cc7bcbf6021ed2733cde173f09dba19abb3 # via # -r ci/official/requirements_updater/requirements.in # tb-nightly diff --git a/requirements_lock_3_13.txt b/requirements_lock_3_13.txt index da604d3fa08573..d86b43c3312010 100644 --- a/requirements_lock_3_13.txt +++ b/requirements_lock_3_13.txt @@ -1083,9 +1083,9 @@ zstandard==0.25.0 \ # via -r ci/official/requirements_updater/requirements.in # The following packages are considered to be unsafe in a requirements file: -setuptools==78.1.1 \ - --hash=sha256:c3a9c4211ff4c309edb8b8c4f1cbfa7ae324c4ba9f91ff254e3d305b9fd54561 \ - --hash=sha256:fcc17fd9cd898242f6b4adfaca46137a9edef687f43e6f78469692a5e70d851d +setuptools==83.0.0 \ + --hash=sha256:025bccbbf0fa05b6192bc64ae1e7b16e001fd6d6d4d5de03c97b1c1ade523bef \ + --hash=sha256:29b23c360f22f414dc7336bb39178cc7bcbf6021ed2733cde173f09dba19abb3 # via # -r ci/official/requirements_updater/requirements.in # tb-nightly diff --git a/requirements_lock_3_14.txt b/requirements_lock_3_14.txt index 0a4539230b6f6c..12edbb9376b9f4 100644 --- a/requirements_lock_3_14.txt +++ b/requirements_lock_3_14.txt @@ -1317,9 +1317,9 @@ zstandard==0.25.0 \ # via -r ci/official/requirements_updater/requirements.in # The following packages are considered to be unsafe in a requirements file: -setuptools==78.1.1 \ - --hash=sha256:c3a9c4211ff4c309edb8b8c4f1cbfa7ae324c4ba9f91ff254e3d305b9fd54561 \ - --hash=sha256:fcc17fd9cd898242f6b4adfaca46137a9edef687f43e6f78469692a5e70d851d +setuptools==83.0.0 \ + --hash=sha256:025bccbbf0fa05b6192bc64ae1e7b16e001fd6d6d4d5de03c97b1c1ade523bef \ + --hash=sha256:29b23c360f22f414dc7336bb39178cc7bcbf6021ed2733cde173f09dba19abb3 # via # -r ci/official/requirements_updater/requirements.in # tb-nightly diff --git a/requirements_lock_3_14_freethreaded.txt b/requirements_lock_3_14_freethreaded.txt index 3fbc748ba944e5..ce6b78a4f5f62c 100644 --- a/requirements_lock_3_14_freethreaded.txt +++ b/requirements_lock_3_14_freethreaded.txt @@ -1318,9 +1318,9 @@ zstandard==0.25.0 \ # via -r ci/official/requirements_updater/requirements.in # The following packages are considered to be unsafe in a requirements file: -setuptools==78.1.1 \ - --hash=sha256:c3a9c4211ff4c309edb8b8c4f1cbfa7ae324c4ba9f91ff254e3d305b9fd54561 \ - --hash=sha256:fcc17fd9cd898242f6b4adfaca46137a9edef687f43e6f78469692a5e70d851d +setuptools==83.0.0 \ + --hash=sha256:025bccbbf0fa05b6192bc64ae1e7b16e001fd6d6d4d5de03c97b1c1ade523bef \ + --hash=sha256:29b23c360f22f414dc7336bb39178cc7bcbf6021ed2733cde173f09dba19abb3 # via # -r ci/official/requirements_updater/requirements.in # tb-nightly diff --git a/tensorflow/compiler/mlir/lite/tests/optimize.mlir b/tensorflow/compiler/mlir/lite/tests/optimize.mlir index c6e18be30ae02d..1dca5b29a23b74 100644 --- a/tensorflow/compiler/mlir/lite/tests/optimize.mlir +++ b/tensorflow/compiler/mlir/lite/tests/optimize.mlir @@ -4067,6 +4067,46 @@ func.func @gelu_approximate1_with_mul1(%arg0: tensor<3xf32>) -> tensor<3xf32> { // CHECK: "tfl.gelu"(%arg0) <{approximate = true}> : (tensor<3xf32>) -> tensor<3xf32> } +func.func @gelu_approximate_with_mul3(%arg0: tensor<3xf32>) -> tensor<3xf32> { + %cst = arith.constant dense<0.797884583> : tensor + %cst_0 = arith.constant dense<5.000000e-01> : tensor + %cst_1 = arith.constant dense<1.000000e+00> : tensor + %cst_3 = arith.constant dense<4.471500e-02> : tensor + %0 = "tfl.mul"(%arg0, %cst_3) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor) -> tensor<3xf32> + %1 = "tfl.mul"(%0, %arg0) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor<3xf32>) -> tensor<3xf32> + %2 = "tfl.mul"(%1, %arg0) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor<3xf32>) -> tensor<3xf32> + %3 = "tfl.add"(%arg0, %2) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor<3xf32>) -> tensor<3xf32> + %4 = "tfl.mul"(%3, %cst) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor) -> tensor<3xf32> + %5 = "tfl.tanh"(%4) : (tensor<3xf32>) -> tensor<3xf32> + %6 = "tfl.add"(%5, %cst_1) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor) -> tensor<3xf32> + %7 = "tfl.mul"(%arg0, %cst_0) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor) -> tensor<3xf32> + %8 = "tfl.mul"(%7, %6) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor<3xf32>) -> tensor<3xf32> + func.return %8 : tensor<3xf32> + +// CHECK-LABEL:gelu_approximate +// CHECK: "tfl.gelu"(%arg0) <{approximate = true}> : (tensor<3xf32>) -> tensor<3xf32> +} + +func.func @gelu_approximate1_with_mul3(%arg0: tensor<3xf32>) -> tensor<3xf32> { + %cst = arith.constant dense<0.797884583> : tensor + %cst_0 = arith.constant dense<5.000000e-01> : tensor + %cst_1 = arith.constant dense<1.000000e+00> : tensor + %cst_3 = arith.constant dense<4.471500e-02> : tensor + %0 = "tfl.mul"(%arg0, %cst_3) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor) -> tensor<3xf32> + %1 = "tfl.mul"(%0, %arg0) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor<3xf32>) -> tensor<3xf32> + %2 = "tfl.mul"(%1, %arg0) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor<3xf32>) -> tensor<3xf32> + %3 = "tfl.add"(%arg0, %2) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor<3xf32>) -> tensor<3xf32> + %4 = "tfl.mul"(%3, %cst) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor) -> tensor<3xf32> + %5 = "tfl.tanh"(%4) : (tensor<3xf32>) -> tensor<3xf32> + %6 = "tfl.add"(%5, %cst_1) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor) -> tensor<3xf32> + %7 = "tfl.mul"(%6, %cst_0) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor) -> tensor<3xf32> + %8 = "tfl.mul"(%arg0, %7) {fused_activation_function = "NONE"} : (tensor<3xf32>, tensor<3xf32>) -> tensor<3xf32> + func.return %8 : tensor<3xf32> + +// CHECK-LABEL:gelu_approximate +// CHECK: "tfl.gelu"(%arg0) <{approximate = true}> : (tensor<3xf32>) -> tensor<3xf32> +} + func.func @gelu_approximate_no_match(%arg0: tensor<3xf32>) -> tensor<3xf32> { %cst = arith.constant dense<0.797884583> : tensor %cst_0 = arith.constant dense<5.000000e-01> : tensor diff --git a/tensorflow/compiler/mlir/lite/transforms/optimize_patterns.td b/tensorflow/compiler/mlir/lite/transforms/optimize_patterns.td index d1940e9aee4fec..195a978fa9b1ae 100644 --- a/tensorflow/compiler/mlir/lite/transforms/optimize_patterns.td +++ b/tensorflow/compiler/mlir/lite/transforms/optimize_patterns.td @@ -1566,6 +1566,41 @@ def MatchGeluApproximate_Mul2 : Pat< (HasOneUse $sqr_out), ]>; +// Alternate pattern for GeluApproximate to match mul(mul(mul(x, coeff), x), x), +// the cube shape produced by the left-associative spelling +// `0.044715 * x * x * x` (common in Python model code), where the coefficient +// is folded into the innermost mul so no bare x^3 subterm exists. +// 0.5 * x * ( 1 + tanh( sqrt_2dPi * ( x + mul(mul(mul(x, 0.044715), x), x) ) ) ) +def MatchGeluApproximate_Mul3 : Pat< + (TFL_MulOp + (TFL_MulOp:$mul_out $arg0, (Arith_ConstantOp F32ElementsAttr:$Cst_1_2), TFL_AF_None), + (TFL_AddOp:$add_out + (TFL_TanhOp:$tanh_out + (TFL_MulOp:$mul_out1 + (TFL_AddOp:$add_out1 $arg0, + (TFL_MulOp:$mul_out2 + (TFL_MulOp:$mul_out3 + (TFL_MulOp:$mul_out4 $arg0, + (Arith_ConstantOp F32ElementsAttr:$Coeff), TFL_AF_None), + $arg0, TFL_AF_None), + $arg0, TFL_AF_None), TFL_AF_None), + (Arith_ConstantOp F32ElementsAttr:$Cst_sqrt_2dPi), TFL_AF_None)), + (Arith_ConstantOp F32ElementsAttr:$Cst_1), TFL_AF_None), TFL_AF_None), + (TFL_GeluOp $arg0, ConstBoolAttrTrue), + [(FloatValueEquals<"0.5"> $Cst_1_2), + (FloatValueEquals<"1"> $Cst_1), + (FloatValueEquals<"0.797884583"> $Cst_sqrt_2dPi), + (FloatValueEquals<"0.044715"> $Coeff), + (HasOneUse $mul_out), + (HasOneUse $add_out), + (HasOneUse $tanh_out), + (HasOneUse $mul_out1), + (HasOneUse $add_out1), + (HasOneUse $mul_out2), + (HasOneUse $mul_out3), + (HasOneUse $mul_out4), + ]>; + // Alternate pattern for GeluApproximate (see different order for mul), replaces // x * ( 0.5 * ( 1 + tanh( sqrt_2dPi * ( x + 0.044715 * pow( x, 3 ) ) ) ) ) def MatchGeluApproximate1 : Pat< @@ -1656,6 +1691,38 @@ def MatchGeluApproximate1_Mul2 : Pat< (HasOneUse $sqr_out), ]>; +// Alternate pattern for GeluApproximate1 to match mul(mul(mul(x, coeff), x), x). +// x * ( 0.5 * ( 1 + tanh( sqrt_2dPi * ( x + mul(mul(mul(x, 0.044715), x), x) ) ) ) ) +def MatchGeluApproximate1_Mul3 : Pat< + (TFL_MulOp $arg0, + (TFL_MulOp:$mul_out + (TFL_AddOp:$add_out + (TFL_TanhOp:$tanh_out + (TFL_MulOp:$mul_out1 + (TFL_AddOp:$add_out1 $arg0, + (TFL_MulOp:$mul_out2 + (TFL_MulOp:$mul_out3 + (TFL_MulOp:$mul_out4 $arg0, + (Arith_ConstantOp F32ElementsAttr:$Coeff), TFL_AF_None), + $arg0, TFL_AF_None), + $arg0, TFL_AF_None), TFL_AF_None), + (Arith_ConstantOp F32ElementsAttr:$Cst_sqrt_2dPi), TFL_AF_None)), + (Arith_ConstantOp F32ElementsAttr:$Cst_1), TFL_AF_None), (Arith_ConstantOp F32ElementsAttr:$Cst_1_2), TFL_AF_None), TFL_AF_None), + (TFL_GeluOp $arg0, ConstBoolAttrTrue), + [(FloatValueEquals<"0.5"> $Cst_1_2), + (FloatValueEquals<"1"> $Cst_1), + (FloatValueEquals<"0.797884583"> $Cst_sqrt_2dPi), + (FloatValueEquals<"0.044715"> $Coeff), + (HasOneUse $mul_out), + (HasOneUse $add_out), + (HasOneUse $tanh_out), + (HasOneUse $mul_out1), + (HasOneUse $add_out1), + (HasOneUse $mul_out2), + (HasOneUse $mul_out3), + (HasOneUse $mul_out4), + ]>; + // For Gelu, replaces // 0.5 * x * ( 1 + erf( x * sqrt_1_2 ) ) def MatchGelu : Pat< diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD index 00cb78d3a71774..fbd222e4984c2f 100644 --- a/tensorflow/compiler/tests/BUILD +++ b/tensorflow/compiler/tests/BUILD @@ -1811,11 +1811,13 @@ tf_xla_py_strict_test( ], deps = [ ":xla_test", + "//tensorflow/python/eager:def_function", "//tensorflow/python/framework:constant_op", "//tensorflow/python/framework:dtypes", "//tensorflow/python/framework:errors", "//tensorflow/python/ops:array_ops", "//tensorflow/python/ops:list_ops", + "//tensorflow/python/ops:map_fn", "//tensorflow/python/platform:client_testlib", "//third_party/py/numpy", "@absl_py//absl/testing:parameterized", diff --git a/tensorflow/compiler/tests/slice_ops_test.py b/tensorflow/compiler/tests/slice_ops_test.py index 4328a65e5f0511..06df524bafb212 100644 --- a/tensorflow/compiler/tests/slice_ops_test.py +++ b/tensorflow/compiler/tests/slice_ops_test.py @@ -50,6 +50,28 @@ def testZeroSlice(self): self.assertAllEqual([], result) + def testSliceOfDynamicDimension(self): + # Regression test for GitHub issue 110789. The input's leading dimension + # is dynamic because it comes from where, and a partial slice of such a + # dimension returned wrong output sizes, and corrupted the heap on some + # platforms, instead of taking the requested slice. + with self.session(): + i = array_ops.placeholder(dtypes.int64, shape=[2, 2, 5]) + with self.test_scope(): + indices = array_ops.where(math_ops.not_equal(i, 0)) + sliced = array_ops.slice(indices, [0, 0], [2, 1]) + empty = array_ops.slice(indices, [0, 0], [0, 1]) + params = { + i: [ + [[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]], + [[10, 11, 12, 13, 14], [15, 16, 17, 18, 19]], + ], + } + # The first two nonzero elements sit at indices (0, 0, 1) and + # (0, 0, 2), so the slice holds the first coordinate of each. + self.assertAllEqual([[0], [0]], sliced.eval(feed_dict=params)) + self.assertAllEqual((0, 1), empty.eval(feed_dict=params).shape) + def test3D(self): for dtype in self.numeric_types: with self.session(): diff --git a/tensorflow/compiler/tests/tensor_list_ops_test.py b/tensorflow/compiler/tests/tensor_list_ops_test.py index 66bfd3008b783a..05afb8f781c451 100644 --- a/tensorflow/compiler/tests/tensor_list_ops_test.py +++ b/tensorflow/compiler/tests/tensor_list_ops_test.py @@ -19,17 +19,47 @@ from absl.testing import parameterized import numpy as np + from tensorflow.compiler.tests import xla_test +from tensorflow.python.eager import def_function from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors from tensorflow.python.ops import array_ops from tensorflow.python.ops import list_ops +from tensorflow.python.ops import map_fn from tensorflow.python.platform import test class ListOpsTest(parameterized.TestCase, xla_test.XLATestCase): + def testGetItemFromEmptyList(self): + # Regression test for GitHub issue 109648. Reading from a statically + # empty list appears in code that never runs, such as the body of a + # while loop with a zero trip count, but XLA compiles that code anyway + # and used to reject the read at compile time. It now yields zeros of + # the element shape. + with self.session() as sess, self.test_scope(): + l = list_ops.tensor_list_reserve( + element_shape=[2], element_dtype=dtypes.float32, num_elements=0 + ) + e = list_ops.tensor_list_get_item(l, 0, element_dtype=dtypes.float32) + self.assertAllEqual(sess.run(e), [0.0, 0.0]) + + def testMapFnOverEmptyTensor(self): + # End to end case for GitHub issue 109648: map_fn over a zero length + # tensor compiles its loop body even though it never runs, and the + # TensorListGetItem in that body used to fail compilation. The whole + # function is jit-compiled so the list stays inside XLA rather than + # crossing the XLA/TF boundary at the unstack and stack ops. + @def_function.function(jit_compile=True) + def f(x): + return map_fn.map_fn(lambda t: t + 1.0, x) + + with self.session() as sess: + x = array_ops.zeros([0], dtype=dtypes.float32) + self.assertAllEqual(sess.run(f(x)).shape, (0,)) + def testElementShape(self): with self.session() as sess, self.test_scope(): dim = array_ops.placeholder(dtypes.int32) diff --git a/tensorflow/compiler/tf2xla/BUILD b/tensorflow/compiler/tf2xla/BUILD index 92786fc8f30ed7..59283a2bbe0d59 100644 --- a/tensorflow/compiler/tf2xla/BUILD +++ b/tensorflow/compiler/tf2xla/BUILD @@ -1643,11 +1643,13 @@ tf_cc_test( srcs = ["xla_op_registry_test.cc"], deps = [ ":xla_compiler", + "//tensorflow/compiler/tf2xla/kernels:xla_ops", "//tensorflow/core:framework", "//tensorflow/core:protos_all_cc", "//tensorflow/core:test", "//tensorflow/core:test_main", "@com_google_absl//absl/log", + "@com_google_googletest//:gtest", ], ) diff --git a/tensorflow/compiler/tf2xla/kernels/BUILD b/tensorflow/compiler/tf2xla/kernels/BUILD index f1b6ab5f79ecb4..a2745ecb61b490 100644 --- a/tensorflow/compiler/tf2xla/kernels/BUILD +++ b/tensorflow/compiler/tf2xla/kernels/BUILD @@ -1587,6 +1587,7 @@ tf_kernel_library( "//tensorflow/core:lib", "@com_google_absl//absl/container:inlined_vector", "@com_google_absl//absl/types:span", + "@xla//xla:shape_util", "@xla//xla:xla_data_proto_cc", "@xla//xla/hlo/builder:value_inference", "@xla//xla/hlo/builder:xla_builder", @@ -1873,6 +1874,7 @@ tf_kernel_library( "//tensorflow/compiler/tf2xla:xla_resource", "//tensorflow/compiler/tf2xla/ops:xla_ops", "//tensorflow/core:framework", + "//tensorflow/core:protos_all_cc", "@com_google_absl//absl/status:statusor", "@xla//xla:xla_data_proto_cc", "@xla//xla/hlo/builder:xla_builder", diff --git a/tensorflow/compiler/tf2xla/kernels/slice_op.cc b/tensorflow/compiler/tf2xla/kernels/slice_op.cc index b0e337cec20c33..89fb4c06f26061 100644 --- a/tensorflow/compiler/tf2xla/kernels/slice_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/slice_op.cc @@ -26,6 +26,7 @@ limitations under the License. #include "xla/hlo/builder/lib/dynamic_shaped_ops.h" #include "xla/hlo/builder/value_inference.h" #include "xla/hlo/builder/xla_builder.h" +#include "xla/shape.h" #include "xla/xla_data.pb.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/op_requires.h" @@ -104,6 +105,30 @@ class SliceOp : public XlaOpKernel { } std::vector strides(begin.size(), 1); auto slice = xla::Slice(ctx->Input(0), begin, limits, strides); + + // If the input has dynamic dimensions, the dynamic size that XLA + // infers for a partial slice of such a dimension does not implement + // the slice semantics, so set the output size explicitly. A dimension + // holding `dynamic_size` valid elements contributes + // clamp(dynamic_size - begin, 0, size) elements to the slice. + auto input_xla_shape = ctx->builder()->GetShape(ctx->Input(0)); + OP_REQUIRES_OK(ctx, input_xla_shape.status()); + for (int64_t i = 0; i < input_dims; ++i) { + if (input_xla_shape->is_dynamic_dimension(i)) { + xla::XlaOp input_dynamic_size = + xla::GetDimensionSize(ctx->Input(0), i); + xla::XlaOp begin_size = xla::ConstantR0( + ctx->builder(), static_cast(begin[i])); + xla::XlaOp requested_size = xla::ConstantR0( + ctx->builder(), static_cast(wrapped_size[i])); + xla::XlaOp output_size = + xla::Clamp(xla::ConstantR0(ctx->builder(), 0), + input_dynamic_size - begin_size, requested_size); + slice = xla::RemoveDynamicDimension(slice, i); + slice = xla::SetDimensionSize(slice, output_size, i); + } + } + // Check for slice on dynamic dimensions. std::vector size_is_dynamic; OP_REQUIRES_OK( diff --git a/tensorflow/compiler/tf2xla/kernels/tensor_list_utils.cc b/tensorflow/compiler/tf2xla/kernels/tensor_list_utils.cc index 0a7297456fce8d..422785bc561f0a 100644 --- a/tensorflow/compiler/tf2xla/kernels/tensor_list_utils.cc +++ b/tensorflow/compiler/tf2xla/kernels/tensor_list_utils.cc @@ -549,6 +549,21 @@ absl::Status ExecuteTensorListGetItem(xla::XlaOp list, xla::XlaOp index, TF_ASSIGN_OR_RETURN(xla::Shape list_shape, b->GetShape(list)); const xla::Shape& buffer_shape = xla::ShapeUtil::GetTupleElementShape(list_shape, 0); + + if (buffer_shape.dimensions(0) == 0) { + // The list is statically empty, so this read can only appear in code + // that never executes at runtime, such as the body of a while loop + // with a zero trip count, which XLA still compiles. The slice below + // would fail compile-time shape inference, so return zeros of the + // element shape instead, mirroring how ExecuteTensorListSetItem + // ignores writes that cannot fit the list. + *result = xla::Broadcast( + xla::ConstantLiteral( + b, xla::LiteralUtil::Zero(buffer_shape.element_type())), + buffer_shape.dimensions().subspan(1)); + return absl::OkStatus(); + } + std::vector start_indices(buffer_shape.dimensions().size(), xla::ConstantR0(b, 0)); start_indices[0] = index; diff --git a/tensorflow/compiler/tf2xla/kernels/unary_ops.cc b/tensorflow/compiler/tf2xla/kernels/unary_ops.cc index 90a022f5111e9a..5e5f222e4ef0f3 100644 --- a/tensorflow/compiler/tf2xla/kernels/unary_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/unary_ops.cc @@ -15,6 +15,7 @@ limitations under the License. // Native XLA implementations of simple unary Ops +#include #include #include "absl/status/statusor.h" @@ -27,10 +28,15 @@ limitations under the License. #include "xla/tsl/platform/statusor.h" #include "xla/xla_data.pb.h" #include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/types.pb.h" namespace tensorflow { namespace { +constexpr std::array kReciprocalGpuTypes = { + {DT_HALF, DT_FLOAT, DT_DOUBLE, DT_BFLOAT16, DT_INT64, DT_COMPLEX64, + DT_COMPLEX128}}; + #define XLAJIT_MAKE_UNARY(NAME, COMPUTATION) \ class NAME##Op : public XlaOpKernel { \ public: \ @@ -45,6 +51,15 @@ namespace { }; \ REGISTER_XLA_OP(Name(#NAME), NAME##Op); +// We don't have kernels for all integer types, so we limit the supported types. +#define REGISTER_XLA_RECIPROCAL_OP(BUILDER, OP) \ + REGISTER_XLA_OP(BUILDER.TypeConstraint("T", kFloatAndComplexTypes) \ + .Device(DEVICE_CPU_XLA_JIT), \ + OP); \ + REGISTER_XLA_OP(BUILDER.TypeConstraint("T", kReciprocalGpuTypes) \ + .Device(DEVICE_GPU_XLA_JIT), \ + OP) + XLAJIT_MAKE_UNARY(ComplexAbs, xla::Abs(x)); XLAJIT_MAKE_UNARY(Angle, xla::Atan2(xla::Imag(x), xla::Real(x))); @@ -72,7 +87,9 @@ REGISTER_XLA_OP(Name("IsInf"), MlirXlaOpKernel); REGISTER_XLA_OP(Name("IsNan"), MlirXlaOpKernel); // Return 1/x XLAJIT_MAKE_UNARY(Inv, xla::ScalarLike(x, 1.0) / x); +REGISTER_XLA_RECIPROCAL_OP(Name("Inv"), InvOp); REGISTER_XLA_OP(Name("Reciprocal"), MlirXlaOpKernel); +REGISTER_XLA_RECIPROCAL_OP(Name("Reciprocal"), MlirXlaOpKernel); XLAJIT_MAKE_UNARY(Log, xla::Log(x)); REGISTER_XLA_OP(Name("Log1p"), MlirXlaOpKernel); diff --git a/tensorflow/compiler/tf2xla/xla_op_registry_test.cc b/tensorflow/compiler/tf2xla/xla_op_registry_test.cc index 13b648b78004ac..8dea74642ebe73 100644 --- a/tensorflow/compiler/tf2xla/xla_op_registry_test.cc +++ b/tensorflow/compiler/tf2xla/xla_op_registry_test.cc @@ -15,16 +15,24 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_op_registry.h" +#include + +#include #include "absl/log/log.h" #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" +#include "tensorflow/core/framework/node_def.pb.h" +#include "tensorflow/core/framework/node_def_util.h" #include "tensorflow/core/framework/op.h" #include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/types.h" #include "tensorflow/core/framework/types.pb.h" #include "tensorflow/core/platform/test.h" namespace tensorflow { namespace { +using ::testing::UnorderedElementsAre; + // This test is to verify the correctness of XLA op registration with specific // backend overrides. @@ -139,5 +147,41 @@ TEST(XlaOpRegistryTest, OpWithInfeasibleTypeConstraintIsNotRegistered) { } } +// Test that Inv and Reciprocal drop the types that have no CPU device kernel. +TEST(XlaOpRegistryTest, ReciprocalIsNotRegisteredForUnsupportedTypesOnCpu) { + XlaOpRegistry::RegisterCompilationKernels(); + int found = 0; + for (const KernelDef* kernel : + XlaOpRegistry::DeviceKernels(DEVICE_CPU_XLA_JIT, true)) { + if (kernel->op() != "Inv" && kernel->op() != "Reciprocal") continue; + ++found; + EXPECT_EQ(kernel->constraint_size(), 1); + EXPECT_EQ(kernel->constraint(0).name(), "T"); + EXPECT_THAT(kernel->constraint(0).allowed_values().list().type(), + UnorderedElementsAre(DT_HALF, DT_FLOAT, DT_DOUBLE, DT_BFLOAT16, + DT_COMPLEX64, DT_COMPLEX128)); + } + EXPECT_EQ(found, 2); +} + +// Test that Inv and Reciprocal have a CPU kernel for DT_FLOAT but not DT_INT32. +TEST(XlaOpRegistryTest, ReciprocalOnIntegerHasNoCpuCompilationKernel) { + XlaOpRegistry::RegisterCompilationKernels(); + for (const std::string& op : {"Inv", "Reciprocal"}) { + NodeDef node_def; + node_def.set_name(op); + node_def.set_op(op); + AddNodeAttr("T", DT_INT32, &node_def); + EXPECT_FALSE(FindKernelDef(DeviceType(DEVICE_CPU_XLA_JIT), node_def, + nullptr, nullptr) + .ok()); + node_def.clear_attr(); + AddNodeAttr("T", DT_FLOAT, &node_def); + EXPECT_TRUE(FindKernelDef(DeviceType(DEVICE_CPU_XLA_JIT), node_def, nullptr, + nullptr) + .ok()); + } +} + } // namespace } // namespace tensorflow diff --git a/tensorflow/core/framework/tensor.cc b/tensorflow/core/framework/tensor.cc index e703d5a7279062..3eee081d163e19 100644 --- a/tensorflow/core/framework/tensor.cc +++ b/tensorflow/core/framework/tensor.cc @@ -185,10 +185,11 @@ struct Helper { // Encoder of simple type T to a string. We do a copy. template - static void Encode(TensorBuffer* in, int64_t n, Destination* out) { + static bool Encode(TensorBuffer* in, int64_t n, Destination* out) { DCHECK_EQ(in->size(), sizeof(T) * n); port::AssignRefCounted( absl::string_view(in->base(), in->size()), in, out); + return true; } // Decoder of simple type T. Copy the bytes from "in" into the @@ -241,8 +242,9 @@ struct Helper { // Encodes "n" elements of type string stored in "in" into Cord // "out", which is usually the TensorProto::tensor_content. template - static void Encode(TensorBuffer* in, int64_t n, Destination* out) { + static bool Encode(TensorBuffer* in, int64_t n, Destination* out) { port::EncodeStringList(in->base(), n, out); + return true; } // Decodes "n" elements of type string from "in" and constructs a @@ -278,9 +280,10 @@ struct Helper { // Encodes "n" elements of type ResourceHandle stored in "in" into destination // "out", which is usually the TensorProto::tensor_content. template - static void Encode(TensorBuffer* in, int64_t n, Destination* out) { + static bool Encode(TensorBuffer* in, int64_t n, Destination* out) { EncodeResourceHandleList(in->base(), n, port::NewStringListEncoder(out)); + return true; } // Decodes "n" elements of type string from "in" and constructs a @@ -310,9 +313,9 @@ struct Helper { // Encodes "n" elements of type Variant stored in "in" into destination // "out", which is usually the TensorProto::tensor_content. template - static void Encode(TensorBuffer* in, int64_t n, Destination* out) { - EncodeVariantList(in->base(), n, - port::NewStringListEncoder(out)); + static bool Encode(TensorBuffer* in, int64_t n, Destination* out) { + return EncodeVariantList(in->base(), n, + port::NewStringListEncoder(out)); } // Decodes "n" elements of type Variant from "in" and constructs a @@ -462,14 +465,15 @@ struct ProtoHelper { static size_t NumElements(const TensorProto& proto) { return proto.variant_val().size(); } - static void Fill(const Variant* data, size_t n, TensorProto* proto) { + static bool Fill(const Variant* data, size_t n, TensorProto* proto) { auto* variant_values = proto->mutable_variant_val(); variant_values->Clear(); for (size_t i = 0; i < n; ++i) { VariantTensorData tmp; - data[i].Encode(&tmp); - tmp.ToProto(variant_values->Add()); + if (!data[i].Encode(&tmp)) return false; + if (!tmp.ToProto(variant_values->Add())) return false; } + return true; } }; @@ -916,13 +920,18 @@ TensorBuffer* FromProtoField(Allocator* a, const TensorProto& in, // Copies T[n] stored in the buffer "in" into the repeated field in // "out" corresponding to type T. template -void ToProtoField(const TensorBuffer& in, int64_t n, TensorProto* out) { +bool ToProtoField(const TensorBuffer& in, int64_t n, TensorProto* out) { const T* data = in.base(); // NOTE: T may not the same as // ProtoHelper::FieldType::value_type. E.g., T==int16, // ProtoHelper::FieldType::value_type==int32. If performance is // critical, we can specialize T=float and do memcpy directly. - ProtoHelper::Fill(data, n, out); + if constexpr (std::is_same_v) { + return ProtoHelper::Fill(data, n, out); + } else { + ProtoHelper::Fill(data, n, out); + return true; + } } void RefIfNonNull(core::RefCounted* buf) { @@ -1281,23 +1290,25 @@ bool Tensor::FromProto(Allocator* a, const TensorProto& proto) { return true; } -void Tensor::AsProtoField(TensorProto* proto) const { +bool Tensor::AsProtoField(TensorProto* proto) const { proto->Clear(); shape_.AsProto(proto->mutable_tensor_shape()); proto->set_dtype(dtype()); if (buf_) { - CASES(dtype(), ToProtoField(*buf_, shape_.num_elements(), proto)); + CASES(dtype(), return ToProtoField(*buf_, shape_.num_elements(), proto)); } + return true; } -void Tensor::AsProtoTensorContent(TensorProto* proto) const { +bool Tensor::AsProtoTensorContent(TensorProto* proto) const { proto->Clear(); proto->set_dtype(dtype()); shape_.AsProto(proto->mutable_tensor_shape()); if (buf_) { - CASES(dtype(), Helper::Encode(buf_, shape_.num_elements(), - proto->mutable_tensor_content())); + CASES(dtype(), return Helper::Encode(buf_, shape_.num_elements(), + proto->mutable_tensor_content())); } + return true; } size_t Tensor::TotalBytes() const { diff --git a/tensorflow/core/framework/tensor.h b/tensorflow/core/framework/tensor.h index 9da573259c432d..6c3e93f9d2dec5 100644 --- a/tensorflow/core/framework/tensor.h +++ b/tensorflow/core/framework/tensor.h @@ -404,8 +404,11 @@ class Tensor { /// `AsProtoField()` fills in the repeated field for `proto.dtype()`, while /// `AsProtoTensorContent()` encodes the content in `proto.tensor_content()` /// in a compact form. - void AsProtoField(TensorProto* proto) const; - void AsProtoTensorContent(TensorProto* proto) const; + /// + /// Both return `true` on success, or `false` if serialization fails (for + /// example, if a variant within the tensor fails to encode). + bool AsProtoField(TensorProto* proto) const; + bool AsProtoTensorContent(TensorProto* proto) const; /// \brief Return the tensor data as an `Eigen::Tensor` with the type and /// sizes of this `Tensor`. diff --git a/tensorflow/core/framework/variant.cc b/tensorflow/core/framework/variant.cc index c055cd4f0efbf5..e4dfdabb2f0cc5 100644 --- a/tensorflow/core/framework/variant.cc +++ b/tensorflow/core/framework/variant.cc @@ -52,20 +52,19 @@ std::string TypeNameVariant(const VariantTensorDataProto& value) { } template <> -void EncodeVariant(const VariantTensorDataProto& value, +bool EncodeVariant(const VariantTensorDataProto& value, VariantTensorData* data) { - data->FromConstProto(value); + return data->FromConstProto(value); } template <> bool DecodeVariant(VariantTensorData* data, VariantTensorDataProto* value) { - data->ToProto(value); - return true; + return data->ToProto(value); } template <> -void EncodeVariant(const VariantTensorDataProto& value, std::string* buf) { - value.SerializeToString(buf); +bool EncodeVariant(const VariantTensorDataProto& value, std::string* buf) { + return value.SerializeToString(buf); } template <> @@ -73,14 +72,15 @@ bool DecodeVariant(std::string* buf, VariantTensorDataProto* value) { return value->ParseFromString(*buf); } -void EncodeVariantList(const Variant* variant_array, int64_t n, +bool EncodeVariantList(const Variant* variant_array, int64_t n, std::unique_ptr e) { for (int i = 0; i < n; ++i) { std::string s; - variant_array[i].Encode(&s); + if (!variant_array[i].Encode(&s)) return false; e->Append(s); } e->Finalize(); + return true; } bool DecodeVariantList(std::unique_ptr d, diff --git a/tensorflow/core/framework/variant.h b/tensorflow/core/framework/variant.h index 152e0538f81bfe..07038180420b3b 100644 --- a/tensorflow/core/framework/variant.h +++ b/tensorflow/core/framework/variant.h @@ -47,10 +47,10 @@ template bool DecodeVariant(std::string* buf, T* value); template -void EncodeVariant(const T& value, VariantTensorData* data); +bool EncodeVariant(const T& value, VariantTensorData* data); template -void EncodeVariant(const T& value, std::string* buf); +bool EncodeVariant(const T& value, std::string* buf); // This is an implementation of a type-erased container that can store an // object of any type. The implementation is very similar to std::any, but has @@ -68,9 +68,12 @@ void EncodeVariant(const T& value, std::string* buf); // following functions: // // string TypeName() const; -// void Encode(VariantTensorData* data) const; +// bool Encode(VariantTensorData* data) const; // or: void Encode(...) // bool Decode(VariantTensorData data); // +// Note: `Encode` can return `bool` to indicate success/failure, or `void` +// (which is treated as always succeeding). +// // Simple POD types can elide the Encode/Decode functions, they are provided by // helper methods. // Here are some typical usage patterns: @@ -276,20 +279,24 @@ class Variant { } // Serialize the contents of the stored object into `data`. - void Encode(VariantTensorData* data) const { + // Returns `true` on success, or `false` if serialization fails. + bool Encode(VariantTensorData* data) const { if (!is_empty()) { - GetValue()->Encode(data); + return GetValue()->Encode(data); } + return true; } // Deserialize `data` and update the stored object. bool Decode(VariantTensorData data); // Helper methods to directly serialize/deserialize from strings. - void Encode(std::string* buf) const { + // Returns `true` on success, or `false` if serialization fails. + bool Encode(std::string* buf) const { if (!is_empty()) { - GetValue()->Encode(buf); + return GetValue()->Encode(buf); } + return true; } bool Decode(std::string buf) { if (!is_empty()) { @@ -319,9 +326,9 @@ class Variant { virtual void MoveInto(ValueInterface* memory) = 0; virtual std::string TypeName() const = 0; virtual std::string DebugString() const = 0; - virtual void Encode(VariantTensorData* data) const = 0; + virtual bool Encode(VariantTensorData* data) const = 0; virtual bool Decode(VariantTensorData data) = 0; - virtual void Encode(std::string* buf) const = 0; + virtual bool Encode(std::string* buf) const = 0; virtual bool Decode(std::string data) = 0; }; @@ -367,15 +374,17 @@ class Variant { std::string DebugString() const final { return DebugStringVariant(value); } - void Encode(VariantTensorData* data) const final { - EncodeVariant(value, data); + bool Encode(VariantTensorData* data) const final { + return EncodeVariant(value, data); } bool Decode(VariantTensorData data) final { return DecodeVariant(&data, &value); } - void Encode(std::string* buf) const final { EncodeVariant(value, buf); } + bool Encode(std::string* buf) const final { + return EncodeVariant(value, buf); + } bool Decode(std::string buf) final { return DecodeVariant(&buf, &value); } diff --git a/tensorflow/core/framework/variant_encode_decode.h b/tensorflow/core/framework/variant_encode_decode.h index 7e46f97e1d6631..794840afa1e1aa 100644 --- a/tensorflow/core/framework/variant_encode_decode.h +++ b/tensorflow/core/framework/variant_encode_decode.h @@ -52,52 +52,62 @@ struct TypeResolver {}; // Specialization for POD type template -void EncodeVariantImpl(const T& value, TypeResolver, +bool EncodeVariantImpl(const T& value, TypeResolver, VariantTensorData* data) { data->set_metadata(value); + return true; } // Specialization for tensorflow::Tensor template -void EncodeVariantImpl(const T& value, +bool EncodeVariantImpl(const T& value, TypeResolver, VariantTensorData* data) { data->tensors_.clear(); data->tensors_.push_back(value); + return true; } // Specialization for protobuf template -void EncodeVariantImpl(const T& value, +bool EncodeVariantImpl(const T& value, TypeResolver, VariantTensorData* data) { if (!value.SerializeToString(&data->metadata_)) { data->metadata_.clear(); LOG(ERROR) << "Failed to encode variant " << value.DebugString(); + return false; } + return true; } // Specialization for other types template -typename std::enable_if< - !std::is_pointer::type>::value>::type +typename std::enable_if::type>::value, + bool>::type EncodeVariantImpl(const T& value, TypeResolver, VariantTensorData* data) { - value.Encode(data); + if constexpr (std::is_void_v) { + value.Encode(data); + return true; + } else { + return static_cast(value.Encode(data)); + } } // Specialization for pointers template -typename std::enable_if< - std::is_pointer::type>::value>::type +typename std::enable_if::type>::value, + bool>::type EncodeVariantImpl(const T& value, TypeResolver, VariantTensorData* data) { // Pointers cannot be encoded. + return false; } // Specialization for POD type @@ -254,9 +264,10 @@ std::string DebugStringVariant(const T& value) { } template -void EncodeVariant(const T& value, VariantTensorData* data) { - EncodeVariantImpl(value, TypeResolver(), data); +bool EncodeVariant(const T& value, VariantTensorData* data) { + if (!EncodeVariantImpl(value, TypeResolver(), data)) return false; data->set_type_name(TypeNameVariant(value)); + return true; } template @@ -265,12 +276,12 @@ bool DecodeVariant(VariantTensorData* data, T* value) { } template -void EncodeVariant(const T& value, std::string* buf) { +bool EncodeVariant(const T& value, std::string* buf) { VariantTensorData data; - EncodeVariantImpl(value, TypeResolver(), &data); + if (!EncodeVariantImpl(value, TypeResolver(), &data)) return false; data.set_type_name(TypeNameVariant(value)); DCHECK(buf != nullptr); - data.SerializeToString(buf); + return data.SerializeToString(buf); } template @@ -288,21 +299,21 @@ template <> std::string TypeNameVariant(const VariantTensorDataProto& value); template <> -void EncodeVariant(const VariantTensorDataProto& value, +bool EncodeVariant(const VariantTensorDataProto& value, VariantTensorData* data); template <> bool DecodeVariant(VariantTensorData* data, VariantTensorDataProto* value); template <> -void EncodeVariant(const VariantTensorDataProto& value, std::string* buf); +bool EncodeVariant(const VariantTensorDataProto& value, std::string* buf); template <> bool DecodeVariant(std::string* buf, VariantTensorDataProto* value); // Encodes an array of Variant objects in to the given StringListEncoder. // `variant_array` is assumed to point to an array of `n` Variant objects. -void EncodeVariantList(const Variant* variant_array, int64_t n, +bool EncodeVariantList(const Variant* variant_array, int64_t n, std::unique_ptr e); // Decodes an array of Variant objects from the given StringListDecoder. diff --git a/tensorflow/core/framework/variant_tensor_data.cc b/tensorflow/core/framework/variant_tensor_data.cc index 906cfaa3d8e58a..2eaf866c9911fe 100644 --- a/tensorflow/core/framework/variant_tensor_data.cc +++ b/tensorflow/core/framework/variant_tensor_data.cc @@ -39,13 +39,14 @@ Tensor* VariantTensorData::add_tensors() { return &(tensors_[tensors_.size() - 1]); } -void VariantTensorData::ToProto(VariantTensorDataProto* proto) const { +bool VariantTensorData::ToProto(VariantTensorDataProto* proto) const { proto->set_type_name(type_name()); proto->set_metadata(metadata_); proto->clear_tensors(); for (const auto& tensor : tensors_) { - tensor.AsProtoField(proto->mutable_tensors()->Add()); + if (!tensor.AsProtoField(proto->mutable_tensors()->Add())) return false; } + return true; } bool VariantTensorData::FromProto(VariantTensorDataProto proto) { @@ -73,13 +74,13 @@ bool VariantTensorData::FromConstProto(const VariantTensorDataProto& proto) { std::string VariantTensorData::SerializeAsString() const { VariantTensorDataProto proto; - ToProto(&proto); + if (!ToProto(&proto)) return ""; return proto.SerializeAsString(); } bool VariantTensorData::SerializeToString(std::string* buf) { VariantTensorDataProto proto; - ToProto(&proto); + if (!ToProto(&proto)) return false; return proto.SerializeToString(buf); } diff --git a/tensorflow/core/framework/variant_tensor_data.h b/tensorflow/core/framework/variant_tensor_data.h index 0778cecaad9a13..e4fabce2994ba8 100644 --- a/tensorflow/core/framework/variant_tensor_data.h +++ b/tensorflow/core/framework/variant_tensor_data.h @@ -82,7 +82,7 @@ class VariantTensorData { Tensor* add_tensor(TensorConstructorArgs&&... args); // Conversion to and from VariantTensorDataProto - void ToProto(VariantTensorDataProto* proto) const; + bool ToProto(VariantTensorDataProto* proto) const; // This allows optimizations via std::move. bool FromProto(VariantTensorDataProto proto); bool FromConstProto(const VariantTensorDataProto& proto); diff --git a/tensorflow/core/framework/variant_test.cc b/tensorflow/core/framework/variant_test.cc index bf99cd721ad6de..270e08212751dd 100644 --- a/tensorflow/core/framework/variant_test.cc +++ b/tensorflow/core/framework/variant_test.cc @@ -689,4 +689,73 @@ TEST(BoolVariantTest, DecodeNonBool) { EXPECT_TRUE(parsed.flat()(0).get()); } +struct FallibleType { + bool fail_encode = false; + std::string value = "success"; + + explicit FallibleType(bool fail = false) : fail_encode(fail) {} + + bool Encode(VariantTensorData* data) const { + if (fail_encode) return false; + data->set_metadata(value); + return true; + } + + bool Decode(VariantTensorData data) { return true; } + + std::string TypeName() const { return "FallibleType"; } +}; + +struct InfallibleVoidType { + std::string value = "void_success"; + + void Encode(VariantTensorData* data) const { data->set_metadata(value); } + + bool Decode(VariantTensorData data) { return true; } + + std::string TypeName() const { return "InfallibleVoidType"; } +}; + +TEST(VariantTest, BoolEncodeSuccessAndFailure) { + Variant v_ok = FallibleType{/*fail_encode=*/false}; + VariantTensorData data; + EXPECT_TRUE(v_ok.Encode(&data)); + std::string buf; + EXPECT_TRUE(v_ok.Encode(&buf)); + + Variant v_fail = FallibleType{/*fail_encode=*/true}; + EXPECT_FALSE(v_fail.Encode(&data)); + EXPECT_FALSE(v_fail.Encode(&buf)); +} + +TEST(VariantTest, VoidEncodeTreatedAsTrue) { + Variant v = InfallibleVoidType{}; + VariantTensorData data; + EXPECT_TRUE(v.Encode(&data)); + std::string buf; + EXPECT_TRUE(v.Encode(&buf)); +} + +TEST(VariantTest, ErrorPropagationToAsProtoFieldAndTensorContent) { + Tensor t_ok(DT_VARIANT, TensorShape({2})); + t_ok.flat()(0) = FallibleType{/*fail_encode=*/false}; + t_ok.flat()(1) = InfallibleVoidType{}; + + TensorProto proto_field_ok; + EXPECT_TRUE(t_ok.AsProtoField(&proto_field_ok)); + + TensorProto proto_content_ok; + EXPECT_TRUE(t_ok.AsProtoTensorContent(&proto_content_ok)); + + Tensor t_fail(DT_VARIANT, TensorShape({2})); + t_fail.flat()(0) = FallibleType{/*fail_encode=*/false}; + t_fail.flat()(1) = FallibleType{/*fail_encode=*/true}; + + TensorProto proto_field_fail; + EXPECT_FALSE(t_fail.AsProtoField(&proto_field_fail)); + + TensorProto proto_content_fail; + EXPECT_FALSE(t_fail.AsProtoTensorContent(&proto_content_fail)); +} + } // end namespace tensorflow diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD index 39eb110d9c6d91..d513da9895a48c 100644 --- a/tensorflow/core/kernels/BUILD +++ b/tensorflow/core/kernels/BUILD @@ -2238,6 +2238,8 @@ tf_cc_test( "//tensorflow/core:test", "//tensorflow/core:test_main", "//tensorflow/core:testlib", + "//tensorflow/core/framework:types_proto_cc", + "@com_google_absl//absl/status", ], ) @@ -3349,7 +3351,10 @@ tf_kernel_library( tf_kernel_library( name = "identity_reader_op", prefix = "identity_reader_op", - deps = IO_DEPS + ["@com_google_absl//absl/strings"], + deps = IO_DEPS + [ + "@com_google_absl//absl/status", + "@com_google_absl//absl/strings", + ], ) tf_kernel_library( diff --git a/tensorflow/core/kernels/clustering_ops.cc b/tensorflow/core/kernels/clustering_ops.cc index fae5341f8434df..3fe17a432857a3 100644 --- a/tensorflow/core/kernels/clustering_ops.cc +++ b/tensorflow/core/kernels/clustering_ops.cc @@ -19,36 +19,32 @@ #include #include +#include "absl/container/flat_hash_set.h" #include "absl/log/check.h" #include "absl/status/status.h" #include "absl/strings/str_cat.h" +#include "Eigen/Core" // from @eigen_archive #include "tensorflow/core/framework/types.pb.h" #define EIGEN_USE_THREADS #include #include #include -#include #include #include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/op_requires.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/types.h" -#include "tensorflow/core/lib/core/errors.h" -#include "tensorflow/core/lib/core/threadpool.h" #include "tensorflow/core/lib/gtl/top_n.h" #include "tensorflow/core/lib/random/philox_random.h" #include "tensorflow/core/lib/random/simple_philox.h" #include "tensorflow/core/platform/blocking_counter.h" -#include "tensorflow/core/platform/byte_order.h" #include "tensorflow/core/platform/cpu_info.h" -#include "tensorflow/core/platform/logging.h" namespace tensorflow { namespace { -using errors::InvalidArgument; - template using RowMajorMatrix = Eigen::Matrix; @@ -141,7 +137,7 @@ class KmeansPlusPlusInitializationOp : public OpKernel { Eigen::Map sampled_points( output_sampled_points_tensor->matrix().data(), num_to_sample, point_dimensions); - std::unordered_set sampled_indices; + absl::flat_hash_set sampled_indices; random::PhiloxRandom random(seed); random::SimplePhilox rng(&random); @@ -329,8 +325,10 @@ class NearestNeighborsOp : public OpKernel { points_tensor.matrix().data(), num_points, point_dimensions); const Eigen::Map centers( centers_tensor.matrix().data(), num_centers, center_dimensions); - const int64_t k = - std::min(num_centers, k_tensor.scalar()()); + const int64_t k_tensor_val = k_tensor.scalar()(); + OP_REQUIRES(context, k_tensor_val >= 0, + absl::InvalidArgumentError("Expected k >= 0.")); + const int64_t k = std::min(num_centers, k_tensor_val); Tensor* output_nearest_center_indices_tensor; Tensor* output_nearest_center_distances_tensor; diff --git a/tensorflow/core/kernels/identity_op_test.cc b/tensorflow/core/kernels/identity_op_test.cc index 62da102905007d..b5db1d4b206641 100644 --- a/tensorflow/core/kernels/identity_op_test.cc +++ b/tensorflow/core/kernels/identity_op_test.cc @@ -13,11 +13,15 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include + +#include "absl/status/status.h" #include "tensorflow/core/framework/fake_input.h" #include "tensorflow/core/framework/node_def_builder.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_testutil.h" #include "tensorflow/core/framework/types.h" +#include "tensorflow/core/framework/types.pb.h" #include "tensorflow/core/kernels/ops_testutil.h" #include "tensorflow/core/kernels/ops_util.h" #include "tensorflow/core/lib/core/status_test_util.h" diff --git a/tensorflow/core/kernels/identity_reader_op.cc b/tensorflow/core/kernels/identity_reader_op.cc index 1668ea30f57942..2cfa032a41dd4b 100644 --- a/tensorflow/core/kernels/identity_reader_op.cc +++ b/tensorflow/core/kernels/identity_reader_op.cc @@ -15,9 +15,11 @@ limitations under the License. // See docs in ../ops/io_ops.cc. -#include +#include +#include "absl/status/status.h" #include "absl/strings/escaping.h" +#include "absl/strings/str_cat.h" #include "tensorflow/core/framework/reader_base.h" #include "tensorflow/core/framework/reader_base.pb.h" #include "tensorflow/core/framework/reader_op_kernel.h" diff --git a/tensorflow/core/kernels/inplace_ops.cc b/tensorflow/core/kernels/inplace_ops.cc index 6948cd86c1f8b1..fe855e5e268dd9 100644 --- a/tensorflow/core/kernels/inplace_ops.cc +++ b/tensorflow/core/kernels/inplace_ops.cc @@ -70,7 +70,7 @@ class ParallelConcatUpdate : public OpKernel { } void Compute(OpKernelContext* ctx) override { - auto value = ctx->input(0); + const Tensor& value = ctx->input(0); // Value should be at least rank 1. Also the 0th dimension should be // at least loc_. OP_REQUIRES(ctx, value.dims() >= 1, @@ -80,7 +80,7 @@ class ParallelConcatUpdate : public OpKernel { errors::InvalidArgument("0th dimension of value = ", value.dim_size(0), " must be greater than loc_ = ", loc_)); - auto update = ctx->input(1); + const Tensor& update = ctx->input(1); OP_REQUIRES( ctx, value.dims() == update.dims(), @@ -98,11 +98,25 @@ class ParallelConcatUpdate : public OpKernel { errors::InvalidArgument("update shape doesn't match: ", update.shape().DebugString())); - Tensor output = value; // This creates an alias intentionally. const auto& d = ctx->eigen_device(); - OP_REQUIRES_OK( - ctx, ::tensorflow::functor::DoParallelConcat(d, update, loc_, &output)); - ctx->set_output(0, output); + if (value.dtype() == DT_STRING) { + int forwarded_input = -1; + Tensor* output_ptr = nullptr; + OP_REQUIRES_OK(ctx, + ctx->forward_input_or_allocate_output( + {0}, 0, value.shape(), &output_ptr, &forwarded_input)); + if (forwarded_input == -1) { + OP_REQUIRES_OK(ctx, + ::tensorflow::functor::DoCopy(d, value, output_ptr)); + } + OP_REQUIRES_OK(ctx, ::tensorflow::functor::DoParallelConcat( + d, update, loc_, output_ptr)); + } else { + Tensor output = value; // This creates an alias intentionally. + OP_REQUIRES_OK(ctx, ::tensorflow::functor::DoParallelConcat( + d, update, loc_, &output)); + ctx->set_output(0, output); + } } private: @@ -208,9 +222,9 @@ class InplaceOpBase : public OpKernel { explicit InplaceOpBase(OpKernelConstruction* ctx) : OpKernel(ctx) {} void Compute(OpKernelContext* ctx) override { - auto x = ctx->input(0); - auto i = ctx->input(1); - auto v = ctx->input(2); + const Tensor& x = ctx->input(0); + const Tensor& i = ctx->input(1); + const Tensor& v = ctx->input(2); OP_REQUIRES(ctx, TensorShapeUtils::IsVector(i.shape()), errors::InvalidArgument("i must be a vector. ", @@ -231,17 +245,32 @@ class InplaceOpBase : public OpKernel { "i and x shape doesn't match at index 0: ", i.shape().DebugString(), " vs. ", v.shape().DebugString())); - Tensor y = x; // This creates an alias intentionally. + Tensor* y_ptr = nullptr; + Tensor aliased_y; + if (x.dtype() == DT_STRING) { + int forwarded_input = -1; + OP_REQUIRES_OK(ctx, ctx->forward_input_or_allocate_output( + {0}, 0, x.shape(), &y_ptr, &forwarded_input)); + if (forwarded_input == -1) { + OP_REQUIRES_OK(ctx, DoCopy(ctx, x, y_ptr)); + } + } else { + aliased_y = x; // This creates an alias intentionally. + ctx->set_output(0, aliased_y); + y_ptr = &aliased_y; + } + // Skip processing if tensors are empty. if (x.NumElements() > 0 && v.NumElements() > 0) { - OP_REQUIRES_OK(ctx, DoCompute(ctx, i, v, &y)); + OP_REQUIRES_OK(ctx, DoCompute(ctx, i, v, y_ptr)); } - ctx->set_output(0, y); } protected: virtual absl::Status DoCompute(OpKernelContext* ctx, const Tensor& i, const Tensor& v, Tensor* y) = 0; + virtual absl::Status DoCopy(OpKernelContext* ctx, const Tensor& x, + Tensor* y) = 0; }; } // end namespace @@ -325,6 +354,11 @@ class InplaceOp : public InplaceOpBase { const auto& d = ctx->eigen_device(); return ::tensorflow::functor::DoInplace(d, op, i, v, y); } + absl::Status DoCopy(OpKernelContext* ctx, const Tensor& x, + Tensor* y) override { + const auto& d = ctx->eigen_device(); + return ::tensorflow::functor::DoCopy(d, x, y); + } }; class CopyOpBase : public OpKernel { @@ -332,7 +366,7 @@ class CopyOpBase : public OpKernel { explicit CopyOpBase(OpKernelConstruction* ctx) : OpKernel(ctx) {} void Compute(OpKernelContext* ctx) override { - auto x = ctx->input(0); + const Tensor& x = ctx->input(0); Tensor* y; OP_REQUIRES_OK(ctx, ctx->allocate_output(0, x.shape(), &y)); OP_REQUIRES_OK(ctx, DoCompute(ctx, x, y)); diff --git a/tensorflow/core/kernels/stateless_random_ops_v2_util.h b/tensorflow/core/kernels/stateless_random_ops_v2_util.h index 95a7bf2a7645bd..a4f268cbe6cf7b 100644 --- a/tensorflow/core/kernels/stateless_random_ops_v2_util.h +++ b/tensorflow/core/kernels/stateless_random_ops_v2_util.h @@ -21,6 +21,7 @@ limitations under the License. #include #include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/rng_alg.h" #include "tensorflow/core/kernels/random_op.h" #include "tensorflow/core/kernels/stateless_random_ops_v2.h" #include "tensorflow/core/lib/random/random_distributions.h" @@ -57,9 +58,16 @@ GetKeyCounterAlgFromInputs(OpKernelContext* ctx, int key_input_idx, if (alg == RNG_ALG_AUTO_SELECT) { alg = RNG_ALG_PHILOX; } + if (alg != RNG_ALG_PHILOX && alg != RNG_ALG_THREEFRY) { + return absl::InvalidArgumentError( + absl::StrCat("Unsupported algorithm id: ", alg)); + } - TF_RETURN_IF_ERROR( - CheckKeyCounterShape(alg, key_t.shape(), counter_t.shape())); + // Use the algorithm's counter size, not the algorithm enum value itself + // (THREEFRY == 2), which incorrectly rejected valid length-1 counters. + TF_RETURN_IF_ERROR(CheckKeyCounterShape( + GetCounterSize(static_cast(alg)), key_t.shape(), + counter_t.shape())); return std::make_tuple(key_t, counter_t, alg); } @@ -76,6 +84,12 @@ void FillRandomTensor(OpKernelContext* ctx, Algorithm alg, const Tensor& key, functor::FillPhiloxRandom()( ctx, ctx->eigen_device(), key_data, counter_data, random::PhiloxRandom() /*dummy*/, flat.data(), flat.size(), dist); + } else if (alg == RNG_ALG_THREEFRY) { + OP_REQUIRES(ctx, false, + absl::UnimplementedError( + "The ThreeFry algorithm is only supported under XLA. Use " + "tf.function(jit_compile=True), or set alg to 'philox' or " + "'auto_select' for eager / non-XLA execution.")); } else { OP_REQUIRES(ctx, false, absl::InvalidArgumentError( diff --git a/tensorflow/core/kernels/stochastic_cast_op_test.cc b/tensorflow/core/kernels/stochastic_cast_op_test.cc index 4176b945e575d3..e9e63b80d08012 100644 --- a/tensorflow/core/kernels/stochastic_cast_op_test.cc +++ b/tensorflow/core/kernels/stochastic_cast_op_test.cc @@ -156,7 +156,8 @@ class StochasticCastOpToIntTest : public OpsTestBase { const InType value = InType(0.625); AddInput(TensorShape({kDim, 1}), [&value](int i) { return value; }); AddInput(TensorShape({1}), [](int i) { return kRngKey; }); - AddInput(TensorShape({1}), [](int i) { return kRngCounter; }); + AddInput(TensorShape({2}), + [](int i) { return i == 0 ? kRngCounter : 0; }); AddInput(TensorShape({}), [](int i) { return kAlgorithm; }); TF_ASSERT_OK(RunOpKernel()); @@ -202,7 +203,8 @@ class StochasticCastOpToIntTest : public OpsTestBase { TensorShape shape({static_cast(dim)}); Tensor* input = AddInput(in_type, shape); AddInput(TensorShape({1}), [](int i) { return kRngKey; }); - AddInput(TensorShape({1}), [](int i) { return kRngCounter; }); + AddInput(TensorShape({2}), + [](int i) { return i == 0 ? kRngCounter : 0; }); AddInput(TensorShape({}), [](int i) { return kAlgorithm; }); auto in = input->flat(); diff --git a/tensorflow/lite/core/subgraph.h b/tensorflow/lite/core/subgraph.h index e1f4f0e3c445cc..dc39a72b9b1b6c 100644 --- a/tensorflow/lite/core/subgraph.h +++ b/tensorflow/lite/core/subgraph.h @@ -695,7 +695,7 @@ class Subgraph { // Constructor should be called with the non-nullptr profiler argument. SubgraphAwareProfiler(Profiler* profiler, int64_t subgraph_index) : profiler_(profiler), subgraph_index_(subgraph_index) {} - ~SubgraphAwareProfiler() override {} + ~SubgraphAwareProfiler() override = default; uint32_t BeginEvent(const char* tag, EventType event_type, int64_t event_metadata1, diff --git a/tensorflow/lite/delegates/coreml/BUILD b/tensorflow/lite/delegates/coreml/BUILD index f7d1b373d170c7..6b16e624b31135 100644 --- a/tensorflow/lite/delegates/coreml/BUILD +++ b/tensorflow/lite/delegates/coreml/BUILD @@ -83,14 +83,13 @@ objc_library( copts = CXX17_BAZEL_ONLY_COPTS, deps = [ ":coreml_executor", - ":mlmodel_proto_cc", "//tensorflow/lite:kernel_api", + "//tensorflow/lite:minimal_logging", "//tensorflow/lite/core/c:common", "//tensorflow/lite/delegates/coreml/builders:op_builder", "//tensorflow/lite/kernels:kernel_util", "//tensorflow/lite/kernels/internal:optimized_base", "//tensorflow/lite/kernels/internal:types", "//tensorflow/lite/types:half", - "@FP16", ], ) diff --git a/tensorflow/lite/delegates/coreml/coreml_delegate_kernel.mm b/tensorflow/lite/delegates/coreml/coreml_delegate_kernel.mm index a545638bafd18a..6ec9fe8a3123e2 100644 --- a/tensorflow/lite/delegates/coreml/coreml_delegate_kernel.mm +++ b/tensorflow/lite/delegates/coreml/coreml_delegate_kernel.mm @@ -19,6 +19,7 @@ #include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h" #include "tensorflow/lite/kernels/internal/types.h" #include "tensorflow/lite/kernels/kernel_util.h" +#include "tensorflow/lite/minimal_logging.h" #import "tensorflow/lite/delegates/coreml/coreml_executor.h" @@ -262,7 +263,16 @@ TfLiteStatus TransposeToHWC(const float* chw, float* hwc, const TfLiteIntArray* } } -CoreMlDelegateKernel::~CoreMlDelegateKernel() { [executor_ cleanup]; } +CoreMlDelegateKernel::~CoreMlDelegateKernel() { + @try { + [executor_ cleanup]; + } @catch (NSException* exception) { + const char* reason = [exception.reason UTF8String]; + TFLITE_LOG_PROD(tflite::TFLITE_LOG_ERROR, + "Exception during CoreML cleanup: %s", + reason ? reason : "Unknown reason"); + } +} } // namespace coreml } // namespace delegates diff --git a/tensorflow/lite/delegates/coreml/coreml_executor.h b/tensorflow/lite/delegates/coreml/coreml_executor.h index 9a13984a876579..ce3b4945537eb2 100644 --- a/tensorflow/lite/delegates/coreml/coreml_executor.h +++ b/tensorflow/lite/delegates/coreml/coreml_executor.h @@ -41,8 +41,8 @@ struct TensorData { - (bool)cleanup; -@property MLModel* model API_AVAILABLE(ios(11)); -@property NSString* mlModelFilePath; -@property NSString* compiledModelFilePath; +@property(nonatomic, strong) MLModel* model API_AVAILABLE(ios(11)); +@property(nonatomic, copy) NSString* mlModelFilePath; +@property(nonatomic, copy) NSString* compiledModelFilePath; @property(nonatomic, readonly) int coreMlVersion; @end diff --git a/tensorflow/lite/delegates/coreml/coreml_executor.mm b/tensorflow/lite/delegates/coreml/coreml_executor.mm index 0e4e1a7053588b..8cc7b716d26b9f 100644 --- a/tensorflow/lite/delegates/coreml/coreml_executor.mm +++ b/tensorflow/lite/delegates/coreml/coreml_executor.mm @@ -165,22 +165,37 @@ - (bool)invokeWithInputs:(const std::vector&)inputs - (bool)cleanup { NSError* error = nil; - [[NSFileManager defaultManager] removeItemAtPath:_mlModelFilePath error:&error]; - if (error != nil) { - NSLog(@"Failed cleaning up model: %@", [error localizedDescription]); - return NO; + NSFileManager* fileManager = [NSFileManager defaultManager]; + bool success = true; + if (_mlModelFilePath.length > 0 && [fileManager fileExistsAtPath:_mlModelFilePath]) { + if (![fileManager removeItemAtPath:_mlModelFilePath error:&error]) { + NSLog(@"Failed cleaning up model: %@", [error localizedDescription]); + success = false; + } else { + self.mlModelFilePath = nil; + } + } else { + self.mlModelFilePath = nil; } - [[NSFileManager defaultManager] removeItemAtPath:_compiledModelFilePath error:&error]; - if (error != nil) { - NSLog(@"Failed cleaning up compiled model: %@", [error localizedDescription]); - return NO; + + error = nil; + if (_compiledModelFilePath.length > 0 && [fileManager fileExistsAtPath:_compiledModelFilePath]) { + if (![fileManager removeItemAtPath:_compiledModelFilePath error:&error]) { + NSLog(@"Failed cleaning up compiled model: %@", [error localizedDescription]); + success = false; + } else { + self.compiledModelFilePath = nil; + } + } else { + self.compiledModelFilePath = nil; } - return YES; + return success; } - (NSURL*)saveModel:(CoreML::Specification::Model*)model { NSURL* modelUrl = createTemporaryFile(); NSString* modelPath = [modelUrl path]; + self.mlModelFilePath = modelPath; if (model->specificationversion() == 3) { _coreMlVersion = 2; } else if (model->specificationversion() == 4) { @@ -203,8 +218,8 @@ - (bool)build:(NSURL*)modelUrl { NSLog(@"Error compiling model %@", [error localizedDescription]); return NO; } - _mlModelFilePath = [modelUrl path]; - _compiledModelFilePath = [compileUrl path]; + self.mlModelFilePath = [modelUrl path]; + self.compiledModelFilePath = [compileUrl path]; if (@available(iOS 12.0, *)) { MLModelConfiguration* config = [[MLModelConfiguration alloc] init]; diff --git a/tensorflow/lite/delegates/gpu/cl/BUILD b/tensorflow/lite/delegates/gpu/cl/BUILD index 7f4529e64b3f57..2b7a0503fce2db 100644 --- a/tensorflow/lite/delegates/gpu/cl/BUILD +++ b/tensorflow/lite/delegates/gpu/cl/BUILD @@ -235,6 +235,7 @@ cc_library( "//tensorflow/lite/delegates/gpu/common:status", "//tensorflow/lite/delegates/gpu/common:types", "//tensorflow/lite/experimental/acceleration/compatibility:android_info", + "@com_google_absl//absl/debugging:leak_check", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", ], diff --git a/tensorflow/lite/delegates/gpu/cl/cl_device.cc b/tensorflow/lite/delegates/gpu/cl/cl_device.cc index 90c106bb06e891..f3bf1729e04f69 100644 --- a/tensorflow/lite/delegates/gpu/cl/cl_device.cc +++ b/tensorflow/lite/delegates/gpu/cl/cl_device.cc @@ -20,6 +20,7 @@ limitations under the License. #include #include +#include "absl/debugging/leak_check.h" #include "absl/strings/ascii.h" #include "absl/strings/numbers.h" #include "absl/strings/str_cat.h" @@ -438,7 +439,14 @@ void CLDevice::DisableOneLayerTextureArray() { absl::Status CreateDefaultGPUDevice(CLDevice* result) { // Get num. platforms cl_uint num_platforms; - cl_int status = clGetPlatformIDs(0, nullptr, &num_platforms); + cl_int status; + { + // Some OpenCL drivers leak memory during initial platform discovery. + // Don't report those leaks, as they are out of our control, in third party + // proprietary code. + absl::LeakCheckDisabler disabler; + status = clGetPlatformIDs(0, nullptr, &num_platforms); + } if (status != CL_SUCCESS) { return absl::UnknownError( absl::StrFormat("clGetPlatformIDs returned %d", status)); diff --git a/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc b/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc index 012c1d46c8ec48..75f78b09a0b108 100644 --- a/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc +++ b/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc @@ -2062,17 +2062,17 @@ class Subgraph { TF_LITE_MAYBE_KERNEL_LOG( context, "unsupported fused activation (Relu) in node #%d", node_index); - return kTfLiteOk; + return kTfLiteError; case kTfLiteActReluN1To1: TF_LITE_MAYBE_KERNEL_LOG( context, "unsupported fused activation (ReluMinus1To1) in node #%d", node_index); - return kTfLiteOk; + return kTfLiteError; case kTfLiteActRelu6: TF_LITE_MAYBE_KERNEL_LOG( context, "unsupported fused activation (Relu6) in node #%d", node_index); - return kTfLiteOk; + return kTfLiteError; case kTfLiteActTanh: TF_LITE_MAYBE_KERNEL_LOG( context, "unsupported fused activation (Tanh) in node #%d", @@ -2636,6 +2636,7 @@ class Subgraph { context, "unsupported quantization type %d in tensor #%d in node #%d", tensor.quantization.type, tensor_index, node_index); + return kTfLiteError; } return kTfLiteOk; } @@ -5586,6 +5587,7 @@ class Subgraph { TF_LITE_MAYBE_KERNEL_LOG( logging_context, "invalid padding mode (%d) in node #%d", static_cast(pool_params->padding), node_index); + return kTfLiteError; } if (subgraph != nullptr) { @@ -5877,6 +5879,7 @@ class Subgraph { logging_context, "unexpected number of dimensions %d in the output shape in node %d", SizeOfDimension(&shape_tensor, 0), node_index); + return kTfLiteError; } TF_LITE_ENSURE_STATUS(CheckTensorStaticAllocation( logging_context, shape_tensor, node->inputs->data[1], @@ -5966,6 +5969,7 @@ class Subgraph { logging_context, "number of dimensions %d must be less than %d in SLICE node #%d", num_dims, XNN_MAX_TENSOR_DIMS, node_index); + return kTfLiteError; } TF_LITE_ENSURE_STATUS( CheckTensorFloatOrQUInt8Type(delegate, logging_context, input_tensor, @@ -5987,6 +5991,7 @@ class Subgraph { "begin %" PRId64 " must be greater than 0 in SLICE node #%d", begin[i], node_index); + return kTfLiteError; } if (size[i] <= 0) { // TODO(b/329228576): Add support for negative begin. @@ -6215,6 +6220,7 @@ class Subgraph { "number of dimensions %d must be less than %d " "in TRANSPOSE node #%d", dims_count, XNN_MAX_TENSOR_DIMS, node_index); + return kTfLiteError; } std::array perm; for (int i = 0; i < dims_count; ++i) { @@ -6272,6 +6278,7 @@ class Subgraph { "number of dimensions %d must be less than %d " "in STRIDED_SLICE node #%d", num_dims, XNN_MAX_TENSOR_DIMS, node_index); + return kTfLiteError; } // Only support strides = 1. diff --git a/tensorflow/lite/kernels/conv.cc b/tensorflow/lite/kernels/conv.cc index 6ffb218a269735..bc2562d9a7520a 100644 --- a/tensorflow/lite/kernels/conv.cc +++ b/tensorflow/lite/kernels/conv.cc @@ -101,13 +101,13 @@ struct OpData { int32_t output_activation_max; // Indexes are the offset to the memory buffer in the array used to keep track // of the allocated temporaries. - int32_t im2col_index; - int32_t hwcn_weights_index; - int32_t input_quantized_index; - int32_t scaling_factors_index; - int32_t accum_scratch_index; - int32_t input_offset_index; - int32_t row_sums_index; + int32_t im2col_index = kTensorNotAllocated; + int32_t hwcn_weights_index = kTensorNotAllocated; + int32_t input_quantized_index = kTensorNotAllocated; + int32_t scaling_factors_index = kTensorNotAllocated; + int32_t accum_scratch_index = kTensorNotAllocated; + int32_t input_offset_index = kTensorNotAllocated; + int32_t row_sums_index = kTensorNotAllocated; bool need_hwcn_weights = false; bool have_weights_been_transposed = false; @@ -469,6 +469,22 @@ TfLiteStatus Prepare(KernelType kernel_type, TfLiteContext* context, } } + if (is_hybrid && data->groups != 1) { + TF_LITE_ENSURE(context, + filter->type == kTfLiteInt8 || filter->type == kTfLiteInt4); + TF_LITE_ENSURE_EQ(context, filter->quantization.type, + kTfLiteAffineQuantization); + const auto* affine_quantization = + reinterpret_cast( + filter->quantization.params); + TF_LITE_ENSURE(context, affine_quantization != nullptr); + TF_LITE_ENSURE(context, affine_quantization->scale != nullptr); + TF_LITE_ENSURE_EQ(context, affine_quantization->quantized_dimension, 0); + TF_LITE_ENSURE_EQ(context, affine_quantization->scale->size, + filter->dims->data[0]); + data->is_hybrid_per_channel = true; + } + // The multi-threaded kernel supports neither dilation nor hybrid kernels, and // is incompatible with mutable input filters that might change between evals. data->supports_multithreaded_kernel = @@ -732,9 +748,9 @@ TfLiteStatus Prepare(KernelType kernel_type, TfLiteContext* context, filter->quantization.params); TF_LITE_ENSURE(context, affine_quantization); TF_LITE_ENSURE(context, affine_quantization->scale); - TF_LITE_ENSURE_EQ( - context, affine_quantization->scale->size, - filter->dims->data[affine_quantization->quantized_dimension]); + TF_LITE_ENSURE_EQ(context, affine_quantization->quantized_dimension, 0); + TF_LITE_ENSURE_EQ(context, affine_quantization->scale->size, + filter->dims->data[0]); node->temporaries->data[data->input_offset_index] = data->input_offset_id; TfLiteTensor* input_offsets; TF_LITE_ENSURE_OK( diff --git a/tensorflow/lite/kernels/conv3d.cc b/tensorflow/lite/kernels/conv3d.cc index 2bdf44a37987f4..78218017ce8907 100644 --- a/tensorflow/lite/kernels/conv3d.cc +++ b/tensorflow/lite/kernels/conv3d.cc @@ -127,6 +127,7 @@ TfLiteStatus Prepare(KernelType kernel_type, TfLiteContext* context, // Check input channels matching filter. TF_LITE_ENSURE_EQ(context, input->dims->data[4], filter->dims->data[3]); + TF_LITE_ENSURE(context, input->dims->data[0] >= 0); TF_LITE_ENSURE(context, input->dims->data[1] > 0); TF_LITE_ENSURE(context, input->dims->data[2] > 0); TF_LITE_ENSURE(context, input->dims->data[3] > 0); @@ -174,18 +175,24 @@ TfLiteStatus Prepare(KernelType kernel_type, TfLiteContext* context, // Matching GetWindowedOutputSize in TensorFlow. int out_width, out_height, out_depth; - opdata->padding = ComputePadding3DValues( + TF_LITE_ENSURE_STATUS(ComputePadding3DValuesChecked( params->stride_height, params->stride_width, params->stride_depth, params->dilation_height_factor, params->dilation_width_factor, params->dilation_depth_factor, height, width, depth, filter_height, filter_width, filter_depth, params->padding, &out_height, &out_width, - &out_depth); + &out_depth, &opdata->padding)); int output_spatial_elements = 0; TF_LITE_ENSURE_MSG(context, CheckedNumElements({out_depth, out_height, out_width}, output_spatial_elements) == kTfLiteOk, "%s", "Conv3D output spatial dimensions overflow."); + int total_output_elements = 0; + TF_LITE_ENSURE_MSG(context, + CheckedNumElements({batches, out_depth, out_height, + out_width, channels_out}, + total_output_elements) == kTfLiteOk, + "%s", "Conv3D output dimensions overflow."); std::unique_ptr output_size( TfLiteIntArrayCreate(5), TfLiteIntArrayFree); @@ -306,6 +313,9 @@ TfLiteStatus Eval(KernelType kernel_type, TfLiteContext* context, TfLiteTensor* output; TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output)); + if (NumElements(output) == 0) { + return kTfLiteOk; + } const TfLiteTensor* input; TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &input)); const TfLiteTensor* filter; diff --git a/tensorflow/lite/kernels/conv3d_test.cc b/tensorflow/lite/kernels/conv3d_test.cc index cd92d41754ddb7..feebf80e348159 100644 --- a/tensorflow/lite/kernels/conv3d_test.cc +++ b/tensorflow/lite/kernels/conv3d_test.cc @@ -451,6 +451,44 @@ TEST(Conv3dOpModel, NoIm2ColTensorTest) { 708, 794, 632, 734, 836, 938, 728, 846, 964, 1082})); } +TEST(Conv3dOpModel, HandlesZeroBatch) { + Conv3dOpModel m({TensorType_FLOAT32, {0, 1, 1, 1, 1}}, + {TensorType_FLOAT32, {1, 1, 1, 1, 1}}, + {TensorType_FLOAT32, {}}, Padding_VALID); + + m.SetFilter({1.0f}); + ASSERT_EQ(m.Invoke(), kTfLiteOk); + EXPECT_THAT(m.GetOutputShape(), ElementsAre(0, 1, 1, 1, 1)); +} + +TEST(Conv3dOpModel, HandlesEmptyOutputSpatialDimension) { + Conv3dOpModel m({TensorType_FLOAT32, {1, 1, 8, 4, 1}}, + {TensorType_FLOAT32, {2, 3, 2, 1, 1}}, + {TensorType_FLOAT32, {}}, Padding_VALID, + /*stride_depth=*/1, /*stride_width=*/1, /*stride_height=*/1, + ActivationFunctionType_NONE, + /*dilation_depth=*/1, /*dilation_width=*/3, + /*dilation_height=*/1); + + m.SetInput(CreateRangeVector(32)); + m.SetFilter(CreateRangeVector(12)); + ASSERT_EQ(m.Invoke(), kTfLiteOk); + EXPECT_THAT(m.GetOutputShape(), ElementsAre(1, 0, 6, 1, 1)); +} + +TEST(Conv3dPrepareSecurityTest, RejectsTotalOutputDimensionsOverflow) { + if (sizeof(void*) <= 4) { + GTEST_SKIP() << "Interpreter construction overflows before kernel Prepare " + "on 32-bit."; + } + constexpr int kHugeDim = 46341; + PrepareOnlyConv3dOpModel m({TensorType_FLOAT32, {kHugeDim, 1, 1, 1, 1}}, + {TensorType_FLOAT32, {1, 1, 1, 1, kHugeDim}}, + {TensorType_FLOAT32, {}}, Padding_SAME); + + EXPECT_EQ(m.AllocateTensors(), kTfLiteError); +} + TEST(Conv3dPrepareSecurityTest, RejectsShapeOverflow) { constexpr int kHugeDim = 46341; diff --git a/tensorflow/lite/kernels/conv3d_transpose.cc b/tensorflow/lite/kernels/conv3d_transpose.cc index 01350f2c3436f9..869aaea96435b4 100644 --- a/tensorflow/lite/kernels/conv3d_transpose.cc +++ b/tensorflow/lite/kernels/conv3d_transpose.cc @@ -27,6 +27,7 @@ limitations under the License. #include "tensorflow/lite/kernels/internal/types.h" #include "tensorflow/lite/kernels/kernel_util.h" #include "tensorflow/lite/kernels/padding.h" +#include "tensorflow/lite/util.h" namespace tflite { namespace ops { @@ -106,16 +107,29 @@ TfLiteStatus ResizeOutputAndTemporaryTensors( const int depth = shape_data[1]; const int height = shape_data[2]; const int width = shape_data[3]; + + int output_spatial_elements = 0; + TF_LITE_ENSURE_MSG(context, + CheckedNumElements({depth, height, width}, + output_spatial_elements) == kTfLiteOk, + "%s", + "Conv3DTranspose output spatial dimensions overflow."); + int total_output_elements = 0; + TF_LITE_ENSURE_MSG( + context, + CheckedNumElements({shape_data[0], depth, height, width, shape_data[4]}, + total_output_elements) == kTfLiteOk, + "%s", "Conv3DTranspose output dimensions overflow."); const int filter_depth = filter_shape.Dims(0); const int filter_height = filter_shape.Dims(1); const int filter_width = filter_shape.Dims(2); int unused_out_width, unused_out_height, unused_out_depth; - opdata->padding = ComputePadding3DValues( + TF_LITE_ENSURE_STATUS(ComputePadding3DValuesChecked( params->stride_height, params->stride_width, params->stride_depth, params->dilation_height_factor, params->dilation_width_factor, params->dilation_depth_factor, height, width, depth, filter_height, filter_width, filter_depth, params->padding, &unused_out_height, - &unused_out_width, &unused_out_depth); + &unused_out_width, &unused_out_depth, &opdata->padding)); // Computed shape must match the shape of the input tensor. TF_LITE_ENSURE_EQ(context, unused_out_depth, SizeOfDimension(input, 1)); TF_LITE_ENSURE_EQ(context, unused_out_height, SizeOfDimension(input, 2)); @@ -183,6 +197,25 @@ TfLiteStatus Prepare(KernelType kernel_type, TfLiteContext* context, // Input and filter must have the same number of channels. TF_LITE_ENSURE_EQ(context, SizeOfDimension(input, 4), SizeOfDimension(filter, 4)); + TF_LITE_ENSURE(context, input->dims->data[1] > 0); + TF_LITE_ENSURE(context, input->dims->data[2] > 0); + TF_LITE_ENSURE(context, input->dims->data[3] > 0); + TF_LITE_ENSURE(context, input->dims->data[4] > 0); + TF_LITE_ENSURE(context, filter->dims->data[0] > 0); + TF_LITE_ENSURE(context, filter->dims->data[1] > 0); + TF_LITE_ENSURE(context, filter->dims->data[2] > 0); + TF_LITE_ENSURE(context, filter->dims->data[3] > 0); + TF_LITE_ENSURE(context, filter->dims->data[4] > 0); + + // Validate stride values. + TF_LITE_ENSURE(context, params->stride_depth > 0); + TF_LITE_ENSURE(context, params->stride_height > 0); + TF_LITE_ENSURE(context, params->stride_width > 0); + + // Validate dilation values. + TF_LITE_ENSURE(context, params->dilation_depth_factor > 0); + TF_LITE_ENSURE(context, params->dilation_height_factor > 0); + TF_LITE_ENSURE(context, params->dilation_width_factor > 0); // Check types. TF_LITE_ENSURE_TYPES_EQ(context, input->type, kTfLiteFloat32); @@ -305,6 +338,10 @@ TfLiteStatus Eval(KernelType kernel_type, TfLiteContext* context, kernel_type = kReference; } + if (NumElements(output) == 0) { + return kTfLiteOk; + } + switch (input->type) { case kTfLiteFloat32: EvalFloat(kernel_type, context, node, params, opdata, input, filter, bias, diff --git a/tensorflow/lite/kernels/conv3d_transpose_test.cc b/tensorflow/lite/kernels/conv3d_transpose_test.cc index feb92571864bb4..9c80bc76bf54e0 100644 --- a/tensorflow/lite/kernels/conv3d_transpose_test.cc +++ b/tensorflow/lite/kernels/conv3d_transpose_test.cc @@ -211,6 +211,22 @@ TEST(Conv3dTransposePrepareSecurityTest, RejectsCol2ImOverflow) { EXPECT_EQ(m.AllocateTensors(), kTfLiteError); } +TEST(Conv3dTransposePrepareSecurityTest, + RejectsSpatialOutputDimensionsOverflow) { + if (sizeof(void*) <= 4) { + GTEST_SKIP() << "Interpreter construction overflows before kernel Prepare " + "on 32-bit."; + } + constexpr int kHugeDim = 46341; + PrepareOnlyConv3dTransposeOpModel m( + {1, kHugeDim, kHugeDim, 1, 1}, {TensorType_FLOAT32, {1, 1, 1, 1, 1}}, + {TensorType_FLOAT32, {1, 1, 1, 1, 1}}, {TensorType_FLOAT32, {}}, + Padding_SAME, /*stride_depth=*/kHugeDim, /*stride_width=*/1, + /*stride_height=*/kHugeDim); + + EXPECT_EQ(m.AllocateTensors(), kTfLiteError); +} + TEST(Conv3dTransposePrepareSecurityTest, RejectsZeroFilterOutputChannels) { PrepareOnlyConv3dTransposeOpModel m({1, 1, 1, 1, 1}, {TensorType_FLOAT32, {1, 1, 1, 0, 1}}, @@ -229,6 +245,63 @@ TEST(Conv3dTransposePrepareSecurityTest, RejectsMismatchedOutputChannels) { EXPECT_EQ(m.AllocateTensors(), kTfLiteError); } +TEST(Conv3dTransposePrepareSecurityTest, + RejectsZeroFilterOutputChannelsEvenIfOutputShapeMatches) { + PrepareOnlyConv3dTransposeOpModel m({1, 1, 1, 1, 0}, + {TensorType_FLOAT32, {1, 1, 1, 0, 1}}, + {TensorType_FLOAT32, {1, 1, 1, 1, 1}}, + {TensorType_FLOAT32, {}}, Padding_SAME); + + EXPECT_EQ(m.AllocateTensors(), kTfLiteError); +} + +TEST(Conv3dTransposePrepareSecurityTest, RejectsInvalidStrides) { + PrepareOnlyConv3dTransposeOpModel m( + {1, 1, 1, 1, 1}, {TensorType_FLOAT32, {1, 1, 1, 1, 1}}, + {TensorType_FLOAT32, {1, 1, 1, 1, 1}}, {TensorType_FLOAT32, {}}, + Padding_SAME, /*stride_depth=*/0, /*stride_width=*/1, + /*stride_height=*/1); + + EXPECT_EQ(m.AllocateTensors(), kTfLiteError); +} + +TEST(Conv3dTransposePrepareSecurityTest, RejectsInvalidDilations) { + PrepareOnlyConv3dTransposeOpModel m( + {1, 1, 1, 1, 1}, {TensorType_FLOAT32, {1, 1, 1, 1, 1}}, + {TensorType_FLOAT32, {1, 1, 1, 1, 1}}, {TensorType_FLOAT32, {}}, + Padding_SAME, /*stride_depth=*/1, /*stride_width=*/1, + /*stride_height=*/1, ActivationFunctionType_NONE, + /*dilation_depth=*/0, /*dilation_width=*/1, /*dilation_height=*/1); + + EXPECT_EQ(m.AllocateTensors(), kTfLiteError); +} + +TEST(Conv3dTransposePrepareSecurityTest, RejectsTotalOutputDimensionsOverflow) { + if (sizeof(void*) <= 4) { + GTEST_SKIP() << "Interpreter construction overflows before kernel Prepare " + "on 32-bit."; + } + constexpr int kHugeDim = 46341; + PrepareOnlyConv3dTransposeOpModel m( + {kHugeDim, 1, 1, 1, kHugeDim}, + {TensorType_FLOAT32, {1, 1, 1, kHugeDim, 1}}, + {TensorType_FLOAT32, {kHugeDim, 1, 1, 1, 1}}, {TensorType_FLOAT32, {}}, + Padding_SAME); + + EXPECT_EQ(m.AllocateTensors(), kTfLiteError); +} + +TEST_P(Conv3dTransposeOpTest, HandlesZeroElementsTest) { + Conv3dTransposeOpModel m( + {0, 1, 1, 1, 1}, {TensorType_FLOAT32, {1, 1, 1, 1, 1}}, + {TensorType_FLOAT32, {0, 1, 1, 1, 1}}, {TensorType_FLOAT32, {}}, + Conv3dTransposeOpTest::GetParam()); + + m.SetFilter({1.0f}); + ASSERT_EQ(m.Invoke(), kTfLiteOk); + EXPECT_THAT(m.GetOutputShape(), ElementsAre(0, 1, 1, 1, 1)); +} + TEST_P(Conv3dTransposeOpTest, SimpleFloat32Test) { Conv3dTransposeOpModel m( {1, 3, 3, 5, 2}, {TensorType_FLOAT32, {2, 2, 2, 2, 2}}, diff --git a/tensorflow/lite/kernels/conv_test.cc b/tensorflow/lite/kernels/conv_test.cc index 15722e5c9a36e7..4f71c0f987a2e7 100644 --- a/tensorflow/lite/kernels/conv_test.cc +++ b/tensorflow/lite/kernels/conv_test.cc @@ -350,6 +350,26 @@ TEST(ConvolutionPrepareSecurityTest, RejectsInvalidGroupedOutputChannels) { EXPECT_EQ(m.AllocateTensors(), kTfLiteError); } +TEST(ConvolutionPrepareSecurityTest, RejectsGroupedHybridUint8Filter) { + PrepareOnlyConvolutionOpModel m( + ops::builtin::Register_CONVOLUTION_GENERIC_OPT(), + {TensorType_FLOAT32, {1, 2, 2, 4}}, + {TensorType_UINT8, {2, 1, 1, 2}, 0, 0, 1.0f, 0}, {TensorType_FLOAT32, {}}, + /*stride_width=*/1, /*stride_height=*/1); + + EXPECT_EQ(m.AllocateTensors(), kTfLiteError); +} + +TEST(ConvolutionPrepareSecurityTest, RejectsGroupedHybridPerTensorFilter) { + PrepareOnlyConvolutionOpModel m( + ops::builtin::Register_CONVOLUTION_GENERIC_OPT(), + {TensorType_FLOAT32, {1, 2, 2, 4}}, + {TensorType_INT8, {2, 1, 1, 2}, 0, 0, 1.0f, 0}, {TensorType_FLOAT32, {}}, + /*stride_width=*/1, /*stride_height=*/1); + + EXPECT_EQ(m.AllocateTensors(), kTfLiteError); +} + TEST(ConvolutionPrepareSecurityTest, RejectsPaddingOverflow) { PrepareOnlyConvolutionOpModel m( ops::builtin::Register_CONVOLUTION_GENERIC_OPT(), @@ -2488,6 +2508,65 @@ TEST_P(ConvolutionOpTest, SimpleTestHybridPerChannelGrouped) { 0.16))); } +TEST_P(ConvolutionOpTest, HybridGroupedWithEqualChannelScales) { + float scale = 4.0 / 127.0; + HybridPerChannelConvolutionOpModel m( + GetRegistration(), {TensorType_FLOAT32, {1, 2, 2, 4}}, + {TensorType_INT8, + {2, 1, 1, 2}, + 0, + 0, + 0, + 0, + /*per_channel_quantization=*/true, + /*per_channel_quantization_scales=*/{scale, scale}, + /*per_channel_quantization_offsets=*/{0, 0}, + /*channel_index=*/0}, + {TensorType_FLOAT32, {}}, + /*stride_width=*/1, /*stride_height=*/1); + + m.SetInput({ + 1, + 2, + 3, + 4, // y=0, x=0 + 2, + 3, + 4, + 5, // y=0, x=1 + 3, + 4, + 5, + 6, // y=1, x=0 + 4, + 5, + 6, + 7, // y=1, x=1 + }); + m.SetSignedFilter({ + 1, + 2, // out_channel 0 + 3, + 4, // out_channel 1 + }); + m.SetBias({0, 0}); + + ASSERT_EQ(m.Invoke(), kTfLiteOk); + + EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear( + { + 5, + 25, // + 8, + 32, // + 11, + 39, // + 14, + 46, // + }, + 0.2))); +} + TEST_P(ConvolutionOpTest, SimpleTestHybridWithPaddingPerChannel) { // Test uses the right zero points for padding if needed. const int stride_width = 1; diff --git a/tensorflow/lite/kernels/padding.h b/tensorflow/lite/kernels/padding.h index 1c63665291aa8c..73cea8cc2dcf79 100644 --- a/tensorflow/lite/kernels/padding.h +++ b/tensorflow/lite/kernels/padding.h @@ -24,12 +24,14 @@ limitations under the License. namespace tflite { -inline TfLiteStatus CheckedNarrowPaddingValue(int64_t value, int* result) { - if (result == nullptr || value > std::numeric_limits::max() || - value < std::numeric_limits::min()) { +template +inline TfLiteStatus CheckedNarrowPaddingValue(int64_t value, + TargetInt* result) { + if (result == nullptr || value > std::numeric_limits::max() || + value < std::numeric_limits::min()) { return kTfLiteError; } - *result = static_cast(value); + *result = static_cast(value); return kTfLiteOk; } @@ -77,6 +79,19 @@ inline TfLiteStatus ComputePaddingWithOffsetChecked( return CheckedNarrowPaddingValue(total_padding / 2, padding); } +inline TfLiteStatus ComputePaddingWithOffsetChecked( + int stride, int dilation_rate, int in_size, int filter_size, int out_size, + int16_t* offset, int16_t* padding) { + int offset_int = 0; + int padding_int = 0; + TF_LITE_ENSURE_STATUS(ComputePaddingWithOffsetChecked( + stride, dilation_rate, in_size, filter_size, out_size, &offset_int, + &padding_int)); + TF_LITE_ENSURE_STATUS(CheckedNarrowPaddingValue(padding_int, padding)); + TF_LITE_ENSURE_STATUS(CheckedNarrowPaddingValue(offset_int, offset)); + return kTfLiteOk; +} + inline int ComputePadding(int stride, int dilation_rate, int in_size, int filter_size, int out_size) { int offset = 0; @@ -187,33 +202,64 @@ inline TfLitePaddingValues ComputePaddingHeightWidth( return padding_values; } +inline TfLiteStatus ComputePadding3DValuesChecked( + int stride_height, int stride_width, int stride_depth, + int dilation_rate_height, int dilation_rate_width, int dilation_rate_depth, + int in_height, int in_width, int in_depth, int filter_height, + int filter_width, int filter_depth, TfLitePadding padding, int* out_height, + int* out_width, int* out_depth, Padding3DValues* padding_values) { + if (out_height == nullptr || out_width == nullptr || out_depth == nullptr || + padding_values == nullptr) { + return kTfLiteError; + } + TF_LITE_ENSURE_STATUS(ComputeOutSizeChecked(padding, in_depth, filter_depth, + stride_depth, dilation_rate_depth, + out_depth)); + TF_LITE_ENSURE_STATUS( + ComputeOutSizeChecked(padding, in_height, filter_height, stride_height, + dilation_rate_height, out_height)); + TF_LITE_ENSURE_STATUS(ComputeOutSizeChecked(padding, in_width, filter_width, + stride_width, dilation_rate_width, + out_width)); + + TF_LITE_ENSURE_STATUS(ComputePaddingWithOffsetChecked( + stride_depth, dilation_rate_depth, in_depth, filter_depth, *out_depth, + &padding_values->depth_offset, &padding_values->depth)); + TF_LITE_ENSURE_STATUS(ComputePaddingWithOffsetChecked( + stride_height, dilation_rate_height, in_height, filter_height, + *out_height, &padding_values->height_offset, &padding_values->height)); + TF_LITE_ENSURE_STATUS(ComputePaddingWithOffsetChecked( + stride_width, dilation_rate_width, in_width, filter_width, *out_width, + &padding_values->width_offset, &padding_values->width)); + return kTfLiteOk; +} + inline Padding3DValues ComputePadding3DValues( int stride_height, int stride_width, int stride_depth, int dilation_rate_height, int dilation_rate_width, int dilation_rate_depth, int in_height, int in_width, int in_depth, int filter_height, int filter_width, int filter_depth, TfLitePadding padding, int* out_height, int* out_width, int* out_depth) { - *out_width = ComputeOutSize(padding, in_width, filter_width, stride_width, - dilation_rate_width); - *out_height = ComputeOutSize(padding, in_height, filter_height, stride_height, - dilation_rate_height); - *out_depth = ComputeOutSize(padding, in_depth, filter_depth, stride_depth, - dilation_rate_depth); - Padding3DValues padding_values; - int offset = 0; - padding_values.depth = - ComputePaddingWithOffset(stride_depth, dilation_rate_depth, in_depth, - filter_depth, *out_depth, &offset); - padding_values.depth_offset = offset; - padding_values.height = - ComputePaddingWithOffset(stride_height, dilation_rate_height, in_height, - filter_height, *out_height, &offset); - padding_values.height_offset = offset; - padding_values.width = - ComputePaddingWithOffset(stride_width, dilation_rate_width, in_width, - filter_width, *out_width, &offset); - padding_values.width_offset = offset; + if (out_height != nullptr) *out_height = 0; + if (out_width != nullptr) *out_width = 0; + if (out_depth != nullptr) *out_depth = 0; + padding_values.depth = 0; + padding_values.depth_offset = 0; + padding_values.height = 0; + padding_values.height_offset = 0; + padding_values.width = 0; + padding_values.width_offset = 0; + if (ComputePadding3DValuesChecked( + stride_height, stride_width, stride_depth, dilation_rate_height, + dilation_rate_width, dilation_rate_depth, in_height, in_width, + in_depth, filter_height, filter_width, filter_depth, padding, + out_height, out_width, out_depth, &padding_values) != kTfLiteOk) { + if (out_height != nullptr) *out_height = 0; + if (out_width != nullptr) *out_width = 0; + if (out_depth != nullptr) *out_depth = 0; + return {}; + } return padding_values; } } // namespace tflite diff --git a/tensorflow/lite/kernels/pooling3d.cc b/tensorflow/lite/kernels/pooling3d.cc index 3ab8701f247ec5..800cd280510b72 100644 --- a/tensorflow/lite/kernels/pooling3d.cc +++ b/tensorflow/lite/kernels/pooling3d.cc @@ -315,11 +315,11 @@ TfLiteStatus GenericPrepare(TfLiteContext* context, TfLiteNode* node) { // Matching GetWindowedOutputSize in TensorFlow. int out_width, out_height, out_depth; - params.padding_values = ComputePadding3DValues( + TF_LITE_ENSURE_STATUS(ComputePadding3DValuesChecked( params.stride_height, params.stride_width, params.stride_depth, 1, 1, 1, height, width, depth, params.filter_height, params.filter_width, params.filter_depth, params.padding_type, &out_height, &out_width, - &out_depth); + &out_depth, ¶ms.padding_values)); if (input->type == kTfLiteInt8) { TF_LITE_ENSURE_NEAR(context, input->params.scale, output->params.scale, diff --git a/tensorflow/lite/special_rules.bzl b/tensorflow/lite/special_rules.bzl index a91e80dda9d788..814ba66e69b78b 100644 --- a/tensorflow/lite/special_rules.bzl +++ b/tensorflow/lite/special_rules.bzl @@ -200,6 +200,7 @@ def flex_portable_tensorflow_deps(): "@com_google_absl//absl/synchronization", "@com_google_absl//absl/time", "@com_google_absl//absl/types:optional", + "@com_google_absl//absl/types:span", "@eigen_archive//:eigen3", "@gemmlowp", "@icu//:common", diff --git a/tensorflow/python/distribute/failure_handling/BUILD b/tensorflow/python/distribute/failure_handling/BUILD index b355fb317f7ce5..604dc1c3c06099 100644 --- a/tensorflow/python/distribute/failure_handling/BUILD +++ b/tensorflow/python/distribute/failure_handling/BUILD @@ -122,6 +122,7 @@ tf_py_strict_test( tf_py_strict_test( name = "gce_failure_handler_test", srcs = ["gce_failure_handler_test.py"], + flaky = 1, shard_count = 32, tags = [ "no_mac", # Fails on CI but works fine locally. diff --git a/tensorflow/python/distribute/failure_handling/gce_failure_handler_test.py b/tensorflow/python/distribute/failure_handling/gce_failure_handler_test.py index 5bb6232e51efc7..dac4b089bd04aa 100644 --- a/tensorflow/python/distribute/failure_handling/gce_failure_handler_test.py +++ b/tensorflow/python/distribute/failure_handling/gce_failure_handler_test.py @@ -105,6 +105,9 @@ def raise_if_not_all_exit(grace_period, mpr): raise RuntimeError('Waited long but at least one worker still exist. ' 'Considering size of our model, this should not' ' happen.') + # Allow OS sockets from terminated worker processes to fully clear TIME_WAIT + # before restarting workers on the same ports. + time.sleep(2) class GceFailureHandlingTest(test.TestCase, parameterized.TestCase): diff --git a/tensorflow/python/kernel_tests/array_ops/inplace_ops_test.py b/tensorflow/python/kernel_tests/array_ops/inplace_ops_test.py index 1a7ec203c9467e..960779a712f940 100644 --- a/tensorflow/python/kernel_tests/array_ops/inplace_ops_test.py +++ b/tensorflow/python/kernel_tests/array_ops/inplace_ops_test.py @@ -13,6 +13,7 @@ # limitations under the License. # ============================================================================== """Tests for inplace_ops.""" + import numpy as np from tensorflow.python.framework import dtypes @@ -23,7 +24,6 @@ from tensorflow.python.ops import inplace_ops from tensorflow.python.platform import test as test_lib - BASIC_TYPES = [ dtypes.float32, dtypes.int8, @@ -46,8 +46,9 @@ def testBasicUpdate(self): x = inplace_ops.inplace_update(x, [3], array_ops.ones([1, 3], dtype)) y[3, :] = 1 self.assertAllClose(x, y) - x = inplace_ops.inplace_update(x, [-1], - array_ops.ones([1, 3], dtype) * 2) + x = inplace_ops.inplace_update( + x, [-1], array_ops.ones([1, 3], dtype) * 2 + ) y[-1, :] = 2 self.assertAllClose(x, y) x = inplace_ops.inplace_update(x, 5, array_ops.ones([3], dtype) * 7) @@ -59,12 +60,14 @@ def testBasicUpdateBool(self): x = array_ops.ones([7, 3], dtypes.bool) y = np.ones([7, 3], dtypes.bool.as_numpy_dtype) self.assertAllClose(x, y) - x = inplace_ops.inplace_update(x, [3], array_ops.ones([1, 3], - dtypes.bool)) + x = inplace_ops.inplace_update( + x, [3], array_ops.ones([1, 3], dtypes.bool) + ) y[3, :] = True self.assertAllClose(x, y) - x = inplace_ops.inplace_update(x, [-1], - array_ops.zeros([1, 3], dtypes.bool)) + x = inplace_ops.inplace_update( + x, [-1], array_ops.zeros([1, 3], dtypes.bool) + ) y[-1, :] = False self.assertAllClose(x, y) x = inplace_ops.inplace_update(x, 5, array_ops.zeros([3], dtypes.bool)) @@ -159,20 +162,28 @@ def testAlias(self): self.assertAllClose(vy, vz) def testError(self): - with self.assertRaisesRegex(errors.InvalidArgumentError, - "must be a vector"): - _ = self.evaluate(inplace_ops.inplace_update([[1.]], [[0]], [[10]])) - with self.assertRaisesRegex(errors.InvalidArgumentError, - "x and v shape doesn't match"): - _ = self.evaluate(inplace_ops.inplace_update([[1.]], [0], [10])) - with self.assertRaisesRegex(errors.InvalidArgumentError, - "i and x shape doesn't match"): - _ = self.evaluate(inplace_ops.inplace_update([[1.]], [0, 1], [[10]])) + with self.assertRaisesRegex( + errors.InvalidArgumentError, "must be a vector" + ): + _ = self.evaluate(inplace_ops.inplace_update([[1.0]], [[0]], [[10]])) + with self.assertRaisesRegex( + errors.InvalidArgumentError, "x and v shape doesn't match" + ): + _ = self.evaluate(inplace_ops.inplace_update([[1.0]], [0], [10])) + with self.assertRaisesRegex( + errors.InvalidArgumentError, "i and x shape doesn't match" + ): + _ = self.evaluate(inplace_ops.inplace_update([[1.0]], [0, 1], [[10]])) def testEmpty(self): for dtype in [ - dtypes.float32, dtypes.float64, dtypes.int32, dtypes.int64, dtypes.bool, - dtypes.uint8, dtypes.bfloat16 + dtypes.float32, + dtypes.float64, + dtypes.int32, + dtypes.int64, + dtypes.bool, + dtypes.uint8, + dtypes.bfloat16, ]: with test_util.use_gpu(): test_shapes = [(), (1,), (2, 3), (0, 2), (2, 3, 5), (2, 0, 5)] @@ -185,11 +196,13 @@ def testEmpty(self): self.assertEqual(val.dtype, dtype.as_numpy_dtype) self.assertAllEqual(val, np.zeros(shape, dtype.as_numpy_dtype)) val = self.evaluate( - inplace_ops.empty_like(array_ops.zeros(shape, dtype))) + inplace_ops.empty_like(array_ops.zeros(shape, dtype)) + ) self.assertEqual(val.shape, shape) self.assertEqual(val.dtype, dtype.as_numpy_dtype) - val = self.evaluate(inplace_ops.empty_like( - array_ops.zeros(shape, dtype), init=True)) + val = self.evaluate( + inplace_ops.empty_like(array_ops.zeros(shape, dtype), init=True) + ) self.assertEqual(val.shape, shape) self.assertEqual(val.dtype, dtype.as_numpy_dtype) self.assertAllEqual(val, np.zeros(shape, dtype.as_numpy_dtype)) diff --git a/tensorflow/python/kernel_tests/linalg/BUILD b/tensorflow/python/kernel_tests/linalg/BUILD index f1e9acc8477b09..39d4158f62a388 100644 --- a/tensorflow/python/kernel_tests/linalg/BUILD +++ b/tensorflow/python/kernel_tests/linalg/BUILD @@ -750,8 +750,9 @@ cuda_py_strict_test( cuda_py_strict_test( name = "matrix_triangular_solve_op_test", size = "medium", + timeout = "long", srcs = ["matrix_triangular_solve_op_test.py"], - shard_count = 3, + shard_count = 8, deps = [ "//tensorflow/python/framework:test_lib", "//tensorflow/python/ops:array_ops", diff --git a/tensorflow/python/kernel_tests/random/stateless_random_ops_test.py b/tensorflow/python/kernel_tests/random/stateless_random_ops_test.py index f8a950a3fefc06..fad01517b67d2c 100644 --- a/tensorflow/python/kernel_tests/random/stateless_random_ops_test.py +++ b/tensorflow/python/kernel_tests/random/stateless_random_ops_test.py @@ -18,6 +18,7 @@ from absl.testing import parameterized import numpy as np + from tensorflow.python.compat import compat from tensorflow.python.eager import context from tensorflow.python.eager import def_function @@ -32,6 +33,7 @@ from tensorflow.python.ops import gen_stateless_random_ops_v2 from tensorflow.python.ops import math_ops from tensorflow.python.ops import random_ops +from tensorflow.python.ops import random_ops_util from tensorflow.python.ops import stateless_random_ops as stateless from tensorflow.python.platform import test @@ -511,6 +513,15 @@ def testGetKeyCounterAlg(self): alg = gen_stateless_random_ops_v2.stateless_random_get_alg() self.assertAllEqual(alg.shape, []) + @test_util.run_v2_only + def testThreefryKeyCounterShape(self): + """ThreeFry key/counter must satisfy non-XLA shape checks (see #100252).""" + seed = constant_op.constant([1, 2], dtype=dtypes.int32) + key, counter, alg = random_ops_util.get_key_counter_alg(seed, 'threefry') + self.assertAllEqual(key.shape, [1]) + self.assertAllEqual(counter.shape, [2]) + self.assertEqual(int(alg), random_ops_util.Algorithm.THREEFRY.value) + def assertDTypeEqual(self, a, b): self.assertEqual(dtypes.as_dtype(a), dtypes.as_dtype(b)) diff --git a/tensorflow/python/ops/clustering_ops_test.py b/tensorflow/python/ops/clustering_ops_test.py index e2a3fcb3b75fe1..ab5139106445fa 100644 --- a/tensorflow/python/ops/clustering_ops_test.py +++ b/tensorflow/python/ops/clustering_ops_test.py @@ -16,6 +16,7 @@ import numpy as np +from tensorflow.python.framework import errors_impl from tensorflow.python.framework import test_util from tensorflow.python.ops import clustering_ops from tensorflow.python.platform import test @@ -88,9 +89,9 @@ def testBasic(self): sample = self.evaluate( clustering_ops.kmc2_chain_initialization(self._distances, seed + i)) counts[sample] = counts.get(sample, 0) + 1 - self.assertEqual(len(counts), 2) - self.assertTrue(500 in counts) - self.assertTrue(1000 in counts) + self.assertLen(counts, 2) + self.assertIn(500, counts) + self.assertIn(1000, counts) self.assertGreaterEqual(counts[500], 5) self.assertGreaterEqual(counts[1000], 5) @@ -142,6 +143,31 @@ def testNearest2(self): self.assertAllClose(indices, [[0, 1], [0, 1], [1, 0], [4, 3]]) self.assertAllClose(distances, [[0., 2.], [5., 5.], [1., 5.], [0., 2.]]) + def testNearestNegativeK(self): + with self.cached_session(): + with self.assertRaisesRegex( + errors_impl.InvalidArgumentError, "Expected k >= 0." + ): + self.evaluate( + clustering_ops.nearest_neighbors(self._points, self._centers, -1) + ) + + def testNearestZeroK(self): + with self.cached_session(): + [indices, distances] = clustering_ops.nearest_neighbors( + self._points, self._centers, 0 + ) + self.assertEqual(self.evaluate(indices).shape, (4, 0)) + self.assertEqual(self.evaluate(distances).shape, (4, 0)) + + def testNearestLargerThanNumCenters(self): + with self.cached_session(): + [indices, distances] = clustering_ops.nearest_neighbors( + self._points, self._centers, 10 + ) + self.assertEqual(self.evaluate(indices).shape, (4, 5)) + self.assertEqual(self.evaluate(distances).shape, (4, 5)) + @test_util.run_all_in_graph_and_eager_modes # A test with large inputs. diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops.py b/tensorflow/python/ops/numpy_ops/np_array_ops.py index 5b7d87b83138af..93a1ad4b80eff6 100644 --- a/tensorflow/python/ops/numpy_ops/np_array_ops.py +++ b/tensorflow/python/ops/numpy_ops/np_array_ops.py @@ -1660,12 +1660,14 @@ def broadcast_arrays(*args, **kwargs): # pylint: disable=missing-docstring @tf_export.tf_export('experimental.numpy.sign', v1=[]) @np_utils.np_doc_only('sign') def sign(x, out=None, where=None, **kwargs): # pylint: disable=missing-docstring,redefined-outer-name - if out: - raise ValueError('tf.numpy doesnt support setting out.') - if where: - raise ValueError('tf.numpy doesnt support setting where.') + if out is not None: + raise ValueError("tf.numpy doesn't support setting out.") + if where is not None: + raise ValueError("tf.numpy doesn't support setting where.") if kwargs: - raise ValueError('tf.numpy doesnt support setting {}'.format(kwargs.keys())) + raise ValueError( + "tf.numpy doesn't support setting {}".format(kwargs.keys()) + ) x = asarray(x) diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py index 21b055b2e701e7..6f006cf0c3144c 100644 --- a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py +++ b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py @@ -1398,6 +1398,11 @@ def testSign(self): arr = np.asarray(state.randn(*shape) * 100, dtype=dtype) self.match(np_array_ops.sign(arr), np.sign(arr)) + with self.assertRaisesRegex(ValueError, "doesn't support setting out"): + np_array_ops.sign([1], out=[]) + with self.assertRaisesRegex(ValueError, "doesn't support setting where"): + np_array_ops.sign([1], where=False) + class ArrayManipulationTest(test.TestCase): diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops.py b/tensorflow/python/ops/numpy_ops/np_math_ops.py index cdad24086e2d69..e900e8cb0f420b 100644 --- a/tensorflow/python/ops/numpy_ops/np_math_ops.py +++ b/tensorflow/python/ops/numpy_ops/np_math_ops.py @@ -1180,7 +1180,8 @@ def f(a): 'Function `diff` currently requires a known rank for input `a`. ' f'Received: a={a} (unknown rank)' ) - if (axis + nd if axis < 0 else axis) >= nd: + axis_normalized = axis + nd if axis < 0 else axis + if axis_normalized < 0 or axis_normalized >= nd: raise ValueError( f'Argument `axis` (received axis={axis}) is out of bounds ' f'for input {a} of rank {nd}.' @@ -1607,6 +1608,10 @@ def rank_equal_case(): [array_ops.shape(a), array_ops.shape(weights)], ) weights_sum = math_ops.reduce_sum(weights, axis=axis) + control_flow_assert.Assert( + math_ops.reduce_all(math_ops.not_equal(weights_sum, 0)), + ['Weights sum to zero, cannot be normalized.'], + ) avg = math_ops.reduce_sum(a * weights, axis=axis) / weights_sum return avg, weights_sum @@ -1619,6 +1624,10 @@ def rank_not_equal_case(): array_ops.rank(weights) == 1, [array_ops.rank(weights)] ) weights_sum = math_ops.reduce_sum(weights) + control_flow_assert.Assert( + math_ops.reduce_all(math_ops.not_equal(weights_sum, 0)), + ['Weights sum to zero, cannot be normalized.'], + ) axes = ops.convert_to_tensor([[axis], [0]]) avg = math_ops.tensordot(a, weights, axes) / weights_sum return avg, weights_sum diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops_test.py b/tensorflow/python/ops/numpy_ops/np_math_ops_test.py index 54b71aa89f8c66..bba3ddd46d0963 100644 --- a/tensorflow/python/ops/numpy_ops/np_math_ops_test.py +++ b/tensorflow/python/ops/numpy_ops/np_math_ops_test.py @@ -466,6 +466,18 @@ def testAverageWrongShape(self): with self.assertRaisesWithPredicateMatch(errors.InvalidArgumentError, r''): np_math_ops.average(np.ones([2, 3]), axis=0, weights=np.ones([5])) + def testAverageZeroWeights(self): + # NumPy raises ZeroDivisionError when weights sum to zero. + x = np_array_ops.array([1, 2, 3]) + with self.assertRaises(errors.InvalidArgumentError): + np_math_ops.average(x, weights=np_array_ops.array([0, 0, 0])) + + x2 = np_array_ops.ones([2, 2]) + with self.assertRaises(errors.InvalidArgumentError): + np_math_ops.average( + x2, axis=0, weights=np_array_ops.array([[0, 1], [0, 1]]) + ) + def testClip(self): def run_test(arr, *args, **kwargs): @@ -669,6 +681,37 @@ def testFlatten(self): with self.assertRaises(ValueError): a2.flatten('invalid') + def testDiff(self): + a = np_array_ops.array([[1, 2, 3], [4, 6, 8]]) + self.match(np_math_ops.diff(a), np.diff(a)) + self.match(np_math_ops.diff(a, axis=0), np.diff(a, axis=0)) + self.match(np_math_ops.diff(a, axis=-1), np.diff(a, axis=-1)) + self.match(np_math_ops.diff(a, n=2, axis=1), np.diff(a, n=2, axis=1)) + self.match( + np_math_ops.diff(np_array_ops.array([1, 3, 6], dtype=np.int32)), + np.diff(np.array([1, 3, 6], dtype=np.int32)), + ) + self.match( + np_math_ops.diff(np_array_ops.array([True, False, True])), + np.diff(np.array([True, False, True])), + ) + # Dtype and value parity for float inputs and n=0. + self.match( + np_math_ops.diff(np_array_ops.array([1.5, 2.5, 4.0], np.float32)), + np.diff(np.array([1.5, 2.5, 4.0], dtype=np.float32)), + ) + self.match( + np_math_ops.diff(np_array_ops.array([1, 3, 6], dtype=np.int32), n=0), + np.diff(np.array([1, 3, 6], dtype=np.int32), n=0), + ) + # NumPy raises ValueError for 0-d inputs too. + with self.assertRaisesRegex(ValueError, 'out of bounds'): + np_math_ops.diff(np_array_ops.array(5)) + with self.assertRaisesRegex(ValueError, 'out of bounds'): + np_math_ops.diff(a, axis=2) + with self.assertRaisesRegex(ValueError, 'out of bounds'): + np_math_ops.diff(a, axis=-3) + def testIsInf(self): x1 = ops.convert_to_tensor(-2147483648) x2 = ops.convert_to_tensor(2147483647) diff --git a/tensorflow/python/ops/numpy_ops/tests/np_test.py b/tensorflow/python/ops/numpy_ops/tests/np_test.py index 37c869db58fc48..4c4594b912577f 100644 --- a/tensorflow/python/ops/numpy_ops/tests/np_test.py +++ b/tensorflow/python/ops/numpy_ops/tests/np_test.py @@ -26,9 +26,8 @@ import numpy as onp import six +from tensorflow.python.framework import errors from tensorflow.python.framework import errors_impl -from tensorflow.python.framework import ops -from tensorflow.python.ops.numpy_ops import np_config from tensorflow.python.ops.numpy_ops.tests.config import config from tensorflow.python.ops.numpy_ops.tests.config import FLAGS import tensorflow.python.ops.numpy_ops.tests.extensions as nje @@ -65,6 +64,7 @@ python_scalar_dtypes = [tnp.bool_, tnp.int_, tnp.float64, tnp.complex128] # pylint: disable=unnecessary-lambda,g-long-lambda,expression-not-assigned +# pylint: disable=line-too-long,used-before-assignment,function-redefined def _valid_dtypes_for_shape(shape, dtypes): # Not all (shape, dtype) pairs are valid. In particular, Python scalars only @@ -1785,10 +1785,16 @@ def testAverage(self, shape, dtype, axis, weights_shape, returned, rng_factory): try: self._CheckAgainstNumpy( onp_fun, lnp_fun, args_maker, check_dtypes=check_dtypes, tol=tol) - except ZeroDivisionError: + self._CompileAndCheck( + lnp_fun, + args_maker, + check_dtypes=check_dtypes, + rtol=tol, + atol=tol, + check_incomplete_shape=True, + ) + except (ZeroDivisionError, errors.InvalidArgumentError): self.skipTest("don't support checking for ZeroDivisionError") - self._CompileAndCheck(lnp_fun, args_maker, check_dtypes=check_dtypes, - rtol=tol, atol=tol, check_incomplete_shape=True) @named_parameters(jtu.cases_from_list( {"testcase_name": "_arg{}_ndmin={}".format(i, ndmin), @@ -2823,7 +2829,6 @@ def testLogspace(self, start_shape, stop_shape, num, check_dtypes=False, atol=atol, rtol=tol, check_incomplete_shape=True) - @named_parameters( jtu.cases_from_list( { @@ -3163,4 +3168,3 @@ def f(x): if __name__ == "__main__": absltest.main() - \ No newline at end of file diff --git a/tensorflow/python/ops/random_ops_util.py b/tensorflow/python/ops/random_ops_util.py index 4f9eefcc920e31..2bd0ee4af559ad 100644 --- a/tensorflow/python/ops/random_ops_util.py +++ b/tensorflow/python/ops/random_ops_util.py @@ -26,6 +26,12 @@ from tensorflow.python.ops import math_ops from tensorflow.python.util.tf_export import tf_export +# Matches RNG_MAX_COUNTER_SIZE in tensorflow/core/framework/rng_alg.h. +# ThreeFry only consumes the first uint64; the second element is padding so +# non-XLA kernels (which historically checked against the algorithm enum id) +# and StatelessRandomGetKeyCounter accept the counter shape. +_RNG_MAX_COUNTER_SIZE = 2 + @tf_export("random.Algorithm", "random.experimental.Algorithm") class Algorithm(enum.Enum): @@ -118,7 +124,7 @@ def _get_key_counter(seed, alg): key = array_ops.reshape( _uint32s_to_uint64(math_ops.cast(seed, dtypes.uint32)), [1] ) - counter = array_ops.zeros([1], dtypes.uint64) + counter = array_ops.zeros([_RNG_MAX_COUNTER_SIZE], dtypes.uint64) else: raise ValueError(unsupported_alg_error_msg(alg)) return key, counter diff --git a/tensorflow/python/ops/stateless_random_ops.py b/tensorflow/python/ops/stateless_random_ops.py index 5211b12b13a6f1..f19fd2059eba23 100644 --- a/tensorflow/python/ops/stateless_random_ops.py +++ b/tensorflow/python/ops/stateless_random_ops.py @@ -350,11 +350,13 @@ def stateless_random_uniform( are `"philox"` for [the Philox algorithm](https://www.thesalmons.org/john/random123/papers/random123sc11.pdf), `"threefry"` for [the ThreeFry - algorithm](https://www.thesalmons.org/john/random123/papers/random123sc11.pdf), - and `"auto_select"` (default) for the system to automatically select an - algorithm based the device type. Values of `tf.random.Algorithm` can also - be used. Note that with `"auto_select"`, the outputs of this function may - change when it is running on a different device. + algorithm](https://www.thesalmons.org/john/random123/papers/random123sc11.pdf) + (on CPU/GPU, ThreeFry requires XLA via `tf.function(jit_compile=True)`; + eager non-XLA kernels do not implement ThreeFry), and `"auto_select"` + (default) for the system to automatically select an algorithm based the + device type. Values of `tf.random.Algorithm` can also be used. Note that + with `"auto_select"`, the outputs of this function may change when it is + running on a different device. Returns: A tensor of the specified shape filled with random uniform values. diff --git a/tensorflow/python/util/tfprof_wrapper.cc b/tensorflow/python/util/tfprof_wrapper.cc index 00417ec0291d54..a0dddd039902af 100644 --- a/tensorflow/python/util/tfprof_wrapper.cc +++ b/tensorflow/python/util/tfprof_wrapper.cc @@ -13,6 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include #include #include diff --git a/tensorflow/tools/ci_build/gpu_build/parallel_gpu_execute.sh b/tensorflow/tools/ci_build/gpu_build/parallel_gpu_execute.sh index a00dcbc3f3404a..137897dd7c99c2 100755 --- a/tensorflow/tools/ci_build/gpu_build/parallel_gpu_execute.sh +++ b/tensorflow/tools/ci_build/gpu_build/parallel_gpu_execute.sh @@ -53,31 +53,38 @@ TEST_BINARY="$(rlocation $TEST_WORKSPACE/${1#./})" shift # ******************************************************************* -mkdir -p /var/lock +LOCK_DIR="${TF_LOCK_DIR:-/var/lock}" +mkdir -p "$LOCK_DIR" # Try to acquire any of the TF_GPU_COUNT * TF_TESTS_PER_GPU # slots to run a test at. # # Prefer to allocate 1 test per GPU over 4 tests on 1 GPU. # So, we iterate over TF_TESTS_PER_GPU first. -for j in `seq 0 $((TF_TESTS_PER_GPU-1))`; do - for i in `seq 0 $((TF_GPU_COUNT-1))`; do - exec {lock_fd}>/var/lock/gpulock${i}_${j} || exit 1 - if flock -n "$lock_fd"; - then - ( - # This export only works within the brackets, so it is isolated to one - # single command. - export CUDA_VISIBLE_DEVICES=$i - export HIP_VISIBLE_DEVICES=$i - echo "Running test $TEST_BINARY $* on GPU $CUDA_VISIBLE_DEVICES" - "$TEST_BINARY" $@ - ) - return_code=$? - flock -u "$lock_fd" - exit $return_code - fi +MAX_ATTEMPTS=30 +for attempt in $(seq 1 $MAX_ATTEMPTS); do + for j in `seq 0 $((TF_TESTS_PER_GPU-1))`; do + for i in `seq 0 $((TF_GPU_COUNT-1))`; do + exec {lock_fd}>"${LOCK_DIR}/gpulock${i}_${j}" || exit 1 + if flock -n "$lock_fd"; + then + ( + # This export only works within the brackets, so it is isolated to one + # single command. + export CUDA_VISIBLE_DEVICES=$i + export HIP_VISIBLE_DEVICES=$i + echo "Running test $TEST_BINARY $* on GPU $CUDA_VISIBLE_DEVICES" + "$TEST_BINARY" $@ + ) + return_code=$? + flock -u "$lock_fd" + exec {lock_fd}>&- + exit $return_code + fi + exec {lock_fd}>&- + done done + sleep 1 done -echo "Cannot find a free GPU to run the test $* on, exiting with failure..." +echo "Cannot find a free GPU to run the test $* on after ${MAX_ATTEMPTS} attempts, exiting with failure..." exit 1 diff --git a/tensorflow/tools/tf_sig_build_dockerfiles/Dockerfile b/tensorflow/tools/tf_sig_build_dockerfiles/Dockerfile index 69d1001cb1b512..170f417a8feb00 100644 --- a/tensorflow/tools/tf_sig_build_dockerfiles/Dockerfile +++ b/tensorflow/tools/tf_sig_build_dockerfiles/Dockerfile @@ -14,7 +14,7 @@ # ============================================================================== ################################################################################ -FROM ubuntu:22.04@sha256:c7eb020043d8fc2ae0793fb35a37bff1cf33f156d4d4b12ccc7f3ef8706c38b1 as builder +FROM ubuntu:22.04@sha256:2edbbc5dc405e9612ba3584ce95480277e3eb374407b5505fe26f17df77c7dbc as builder ################################################################################ # Install devtoolset build dependencies diff --git a/third_party/py/rules_cc_protobuf.patch b/third_party/py/rules_cc_protobuf.patch deleted file mode 100644 index fa9278867f7161..00000000000000 --- a/third_party/py/rules_cc_protobuf.patch +++ /dev/null @@ -1,29 +0,0 @@ ---- a/cc/defs.bzl -+++ b/cc/defs.bzl -@@ -18,3 +18,4 @@ - """Starlark rules for building C++ projects.""" - -+load("@com_google_protobuf//bazel:cc_proto_library.bzl", _protobuf_cc_proto_library = "cc_proto_library") - load("//cc:cc_binary.bzl", _cc_binary = "cc_binary") -@@ -48,20 +49,3 @@ - objc_import = _objc_import - --# DEPRECATED: use rule from com_google_protobuf repository --def cc_proto_library(**kwargs): -- """Deprecated redirection macro for cc_proto_library. -- -- Use cc_proto_library from com_google_protobuf. -- -- On Bazel <8, redirects to native.cc_proto_library. -- On Bazel >=8, redirects to a mock rule, that fails when analyzed. -- This allows for a gradual migration away from this macro. -- -- Args: -- **kwargs: passed directly into cc_proto_library -- """ -- __cc_proto_library = getattr(native, "cc_proto_library", _cc_proto_library) -- if "deprecation" not in kwargs: -- __cc_proto_library(deprecation = CC_PROTO_LIBRARY_DEPRECATION, **kwargs) -- else: -- __cc_proto_library(**kwargs) -+cc_proto_library = _protobuf_cc_proto_library \ No newline at end of file diff --git a/third_party/xla/.bazelrc b/third_party/xla/.bazelrc index cefcd34e1290cb..858a58563620ce 100644 --- a/third_party/xla/.bazelrc +++ b/third_party/xla/.bazelrc @@ -14,13 +14,6 @@ common:workspace --noenable_bzlmod --enable_workspace common --incompatible_enable_cc_toolchain_resolution common --repo_env USE_HERMETIC_CC_TOOLCHAIN=1 -# external_include_paths feature forces passing external repositories headers -# via -isystem alters how Clang resolves their file paths. This feature with -# old versions of rules_cc triggering errors. -# Following lines can be removed after rules_cc will be updated to >= 0.2.20 -common --features=-external_include_paths -common --host_features=-external_include_paths - # TODO: Migrate for https://github.com/bazelbuild/bazel/issues/7260 common:clang_local --noincompatible_enable_cc_toolchain_resolution common:clang_local --@rules_ml_toolchain//common:enable_hermetic_cc=False diff --git a/third_party/xla/.github/workflows/rocm_ci.yml b/third_party/xla/.github/workflows/rocm_ci.yml index a722371d30ade1..d06c4a3f49360b 100644 --- a/third_party/xla/.github/workflows/rocm_ci.yml +++ b/third_party/xla/.github/workflows/rocm_ci.yml @@ -180,7 +180,6 @@ jobs: ROCM_DISTRO_URL: ${{ needs.rocm-config.outputs.rocm-distro-url }} ROCM_DISTRO_HASH: ${{ needs.rocm-config.outputs.rocm-distro-hash }} EXECUTE_CI_BUILD_URL: https://raw.githubusercontent.com/ROCm/xla/refs/heads/${{ inputs.rocm_xla_branch || 'rocm-dev-infra' }}/build_tools/rocm/execute_ci_build_upstream.sh - GLOBAL_SYMBOL_VERSION_LDS_URL: https://raw.githubusercontent.com/ROCm/xla/refs/heads/${{ inputs.rocm_xla_branch || 'rocm-dev-infra' }}/build_tools/rocm/global_symbol_version.lds # Unique per run so a crashed run cannot collide with this one on a pooled node. CONTAINER_NAME: xla-${{ github.run_id }}-${{ github.run_attempt }} # Stable label so leftover containers from crashed runs can be swept on a pooled node. @@ -199,18 +198,6 @@ jobs: run: | wget -O build_tools/rocm/execute_ci_build_upstream.sh "${EXECUTE_CI_BUILD_URL}" chmod +x build_tools/rocm/execute_ci_build_upstream.sh - wget -O build_tools/rocm/global_symbol_version.lds "${GLOBAL_SYMBOL_VERSION_LDS_URL}" - LDS_SHA=$(sha256sum build_tools/rocm/global_symbol_version.lds | cut -d' ' -f1) - # Workaround for the symbol clash with ROCm's libLLVM.so - { - echo "LLVM_SYMBOL_CLASH_WAR<> "$GITHUB_ENV" - - *start_container - name: CPU Info @@ -219,7 +206,6 @@ jobs: - name: Test XLA [single_gpu] timeout-minutes: 120 run: | - # LLVM_SYMBOL_CLASH_WAR is intentionally unquoted so it splits into one flag per line. # shellcheck disable=SC2086 docker exec "${CONTAINER_NAME}" build_tools/rocm/execute_ci_build_upstream.sh \ --config=rocm_ci_hermetic \ @@ -232,7 +218,6 @@ jobs: --local_test_jobs=1 \ --internal_spawn_scheduler \ --strategy=TestRunner=dynamic \ - ${LLVM_SYMBOL_CLASH_WAR} \ --bes_keywords=xla \ --bes_keywords=upstream \ --bes_keywords=gpu \ @@ -241,7 +226,6 @@ jobs: - name: Test XLA [rocm_cpu] timeout-minutes: 80 run: | - # LLVM_SYMBOL_CLASH_WAR is intentionally unquoted so it splits into one flag per line. # shellcheck disable=SC2086 docker exec "${CONTAINER_NAME}" build_tools/rocm/execute_ci_build_upstream.sh \ --config=rocm_ci_hermetic \ @@ -252,7 +236,6 @@ jobs: --repo_env=ROCM_PATH="" \ --repo_env=ROCM_DISTRO_URL="${ROCM_DISTRO_URL}" \ --repo_env=ROCM_DISTRO_HASH="${ROCM_DISTRO_HASH}" \ - ${LLVM_SYMBOL_CLASH_WAR} \ --bes_keywords=xla \ --bes_keywords=upstream \ --bes_keywords=cpu \ diff --git a/third_party/xla/MODULE.bazel b/third_party/xla/MODULE.bazel index ae3b0122f1d23f..e6cad514526627 100644 --- a/third_party/xla/MODULE.bazel +++ b/third_party/xla/MODULE.bazel @@ -46,9 +46,9 @@ bazel_dep(name = "rules_ml_toolchain") # echo "sha256-${HASH}" archive_override( module_name = "rules_ml_toolchain", - integrity = "sha256-dcy3xNxpk0PwLQLwtbqm2sCcAT1GuStg8YVjRI6lnnc=", - strip_prefix = "rules_ml_toolchain-73cb731fed3ccf7551beac710bf1c5dbeb8be298", - urls = ["https://github.com/google-ml-infra/rules_ml_toolchain/archive/73cb731fed3ccf7551beac710bf1c5dbeb8be298.tar.gz"], + integrity = "sha256-Ztgjtsa06m8J1572/knhgeUwF1nuB9F2psgqAeEifWs=", + strip_prefix = "rules_ml_toolchain-e8709f15382e4da1de5ca15672cb4412da3d989d", + urls = ["https://github.com/google-ml-infra/rules_ml_toolchain/archive/e8709f15382e4da1de5ca15672cb4412da3d989d.tar.gz"], ) # TODO: Upstream the patch? diff --git a/third_party/xla/tensorflow.bazelrc b/third_party/xla/tensorflow.bazelrc index 12006c287621e5..8a3fd444b601bf 100644 --- a/third_party/xla/tensorflow.bazelrc +++ b/third_party/xla/tensorflow.bazelrc @@ -300,7 +300,8 @@ common:rocm --config=rocm_clang_hermetic common:rocm_ci --config=rocm common:rocm_ci --@local_config_rocm//rocm:rocm_path_type=hermetic -common:rocm_ci_hermetic --dynamic_mode=off +common:rocm_ci_hermetic --dynamic_mode=fully +common:rocm_ci_hermetic --@rules_ml_toolchain//common:enable_xla_test_global_symbol_version=True common:rocm_ci_hermetic --config=rocm_clang_hermetic common:rocm_ci_hermetic --repo_env=TF_ROCM_AMDGPU_TARGETS="gfx908,gfx90a" common:rocm_ci_hermetic --repo_env=ROCM_DISTRO_VERSION="rocm_7.13.0_gfx90a" diff --git a/third_party/xla/third_party/cuda_tile/BUILD.bazel b/third_party/xla/third_party/cuda_tile/BUILD.bazel index e69de29bb2d1d6..637d57529cb687 100644 --- a/third_party/xla/third_party/cuda_tile/BUILD.bazel +++ b/third_party/xla/third_party/cuda_tile/BUILD.bazel @@ -0,0 +1,4 @@ +exports_files( + glob(["patches/**"]), + visibility = ["//visibility:public"], +) diff --git a/third_party/xla/third_party/cuda_tile/patches/constructor.patch b/third_party/xla/third_party/cuda_tile/patches/constructor.patch new file mode 100644 index 00000000000000..98b074b4be7358 --- /dev/null +++ b/third_party/xla/third_party/cuda_tile/patches/constructor.patch @@ -0,0 +1,17 @@ +Remove redundant defaulted copy/move constructors to preserve aggregate status. + +diff --git a/tools/cuda-tile-tblgen/CudaTileOp.h b/tools/cuda-tile-tblgen/CudaTileOp.h +--- a/tools/cuda-tile-tblgen/CudaTileOp.h ++++ b/tools/cuda-tile-tblgen/CudaTileOp.h +@@ -55,11 +55,6 @@ + std::optional> selectedVariants; + ParameterType type; + +- OperationParameter(const OperationParameter &other) = default; +- OperationParameter(OperationParameter &&other) = default; +- OperationParameter &operator=(const OperationParameter &other) = default; +- OperationParameter &operator=(OperationParameter &&other) = default; +- + std::string getDescription() const; + TileIRType getTypeDescription() const; + }; diff --git a/third_party/xla/third_party/cuda_tile/patches/series b/third_party/xla/third_party/cuda_tile/patches/series new file mode 100644 index 00000000000000..b495651479c217 --- /dev/null +++ b/third_party/xla/third_party/cuda_tile/patches/series @@ -0,0 +1 @@ +constructor.patch diff --git a/third_party/xla/third_party/cuda_tile/workspace.bzl b/third_party/xla/third_party/cuda_tile/workspace.bzl index 7dcb4e0c479335..3be6b6cfe6aca8 100644 --- a/third_party/xla/third_party/cuda_tile/workspace.bzl +++ b/third_party/xla/third_party/cuda_tile/workspace.bzl @@ -13,4 +13,5 @@ def repo(): sha256 = CUDA_TILE_SHA256, strip_prefix = "cuda-tile-{}".format(CUDA_TILE_COMMIT), urls = tf_mirror_urls("https://github.com/NVIDIA/cuda-tile/archive/{}.tar.gz".format(CUDA_TILE_COMMIT)), + patch_file = ["//third_party/cuda_tile:patches/constructor.patch"], ) diff --git a/third_party/xla/third_party/gpus/rocm/build_defs.bzl.tpl b/third_party/xla/third_party/gpus/rocm/build_defs.bzl.tpl index 76a5c04c130de1..4d93b1768b7725 100644 --- a/third_party/xla/third_party/gpus/rocm/build_defs.bzl.tpl +++ b/third_party/xla/third_party/gpus/rocm/build_defs.bzl.tpl @@ -76,6 +76,7 @@ def if_rocm_hipblaslt(x): def rocm_library(copts = [], deps = [], **kwargs): """Wrapper over cc_library which adds default ROCm options.""" + deps = list(deps) if "@local_config_rocm//rocm:rocm_headers" not in deps: deps.append("@local_config_rocm//rocm:rocm_headers") cc_library(copts = rocm_default_copts() + copts, deps = deps, **kwargs) diff --git a/third_party/xla/third_party/py/rules_cc_protobuf.patch b/third_party/xla/third_party/py/rules_cc_protobuf.patch deleted file mode 100644 index fa9278867f7161..00000000000000 --- a/third_party/xla/third_party/py/rules_cc_protobuf.patch +++ /dev/null @@ -1,29 +0,0 @@ ---- a/cc/defs.bzl -+++ b/cc/defs.bzl -@@ -18,3 +18,4 @@ - """Starlark rules for building C++ projects.""" - -+load("@com_google_protobuf//bazel:cc_proto_library.bzl", _protobuf_cc_proto_library = "cc_proto_library") - load("//cc:cc_binary.bzl", _cc_binary = "cc_binary") -@@ -48,20 +49,3 @@ - objc_import = _objc_import - --# DEPRECATED: use rule from com_google_protobuf repository --def cc_proto_library(**kwargs): -- """Deprecated redirection macro for cc_proto_library. -- -- Use cc_proto_library from com_google_protobuf. -- -- On Bazel <8, redirects to native.cc_proto_library. -- On Bazel >=8, redirects to a mock rule, that fails when analyzed. -- This allows for a gradual migration away from this macro. -- -- Args: -- **kwargs: passed directly into cc_proto_library -- """ -- __cc_proto_library = getattr(native, "cc_proto_library", _cc_proto_library) -- if "deprecation" not in kwargs: -- __cc_proto_library(deprecation = CC_PROTO_LIBRARY_DEPRECATION, **kwargs) -- else: -- __cc_proto_library(**kwargs) -+cc_proto_library = _protobuf_cc_proto_library \ No newline at end of file diff --git a/third_party/xla/third_party/tensor_ir/BUILD.bazel b/third_party/xla/third_party/tensor_ir/BUILD.bazel index e69de29bb2d1d6..637d57529cb687 100644 --- a/third_party/xla/third_party/tensor_ir/BUILD.bazel +++ b/third_party/xla/third_party/tensor_ir/BUILD.bazel @@ -0,0 +1,4 @@ +exports_files( + glob(["patches/**"]), + visibility = ["//visibility:public"], +) diff --git a/third_party/xla/third_party/tensor_ir/patches/series b/third_party/xla/third_party/tensor_ir/patches/series new file mode 100644 index 00000000000000..e3b672772362dd --- /dev/null +++ b/third_party/xla/third_party/tensor_ir/patches/series @@ -0,0 +1 @@ +unused_variable.patch diff --git a/third_party/xla/third_party/tensor_ir/patches/unused_variable.patch b/third_party/xla/third_party/tensor_ir/patches/unused_variable.patch new file mode 100644 index 00000000000000..d38d9f8924952c --- /dev/null +++ b/third_party/xla/third_party/tensor_ir/patches/unused_variable.patch @@ -0,0 +1,57 @@ +Mark unused variables in tensor_ir as [[maybe_unused]]. + +diff --git a/lib/Conversion/TensorToCudaTile/AffineMapImpl.cpp b/lib/Conversion/TensorToCudaTile/AffineMapImpl.cpp +--- a/lib/Conversion/TensorToCudaTile/AffineMapImpl.cpp ++++ b/lib/Conversion/TensorToCudaTile/AffineMapImpl.cpp +@@ -1138,7 +1138,7 @@ TensorToCudaTileConversionState::compute + getTensorDescriptorStaticMetadata(descriptor); + + size_t ptrIdx = mapping->inputNo; +- size_t numArgsConsumed = processTensorDescriptor( ++ [[maybe_unused]] size_t numArgsConsumed = processTensorDescriptor( + tensorTy, tensorIRArg, cudaTileArgs, ptrIdx, metadata, + CudaTileTensorDescriptor::TensorKind::Input, iterSpaceMap, iterSpaceIds, + descriptor.allowTma, descriptor.alignment, descriptor.cost, +diff --git a/lib/Dialect/Canonicalization.cpp b/lib/Dialect/Canonicalization.cpp +--- a/lib/Dialect/Canonicalization.cpp ++++ b/lib/Dialect/Canonicalization.cpp +@@ -181,7 +181,7 @@ static DenseI64ArrayAttr composeSliceSta + llvm::zip_equal(innerStarts.asArrayRef(), innerStrides.asArrayRef(), + outerStarts.asArrayRef())) { + int64_t start; +- bool overflow = ++ [[maybe_unused]] bool overflow = + !checkedMultiplyAdd(outerStart, innerStride, innerStart, start); + assert(!overflow && "slice composition was not checked"); + starts.push_back(start); +@@ -206,7 +206,7 @@ static DenseI64ArrayAttr composeSliceLim + int64_t start; + int64_t limit; + int64_t stride; +- bool composed = computeComposedSliceDimension( ++ [[maybe_unused]] bool composed = computeComposedSliceDimension( + innerStart, innerLimit, innerStride, outerStart, outerLimit, + outerStride, start, limit, stride); + assert(composed && "slice composition was not checked"); +@@ -223,7 +223,8 @@ static DenseI64ArrayAttr composeSliceStr + for (auto [innerStride, outerStride] : + llvm::zip_equal(innerStrides.asArrayRef(), outerStrides.asArrayRef())) { + int64_t stride; +- bool overflow = llvm::MulOverflow(innerStride, outerStride, stride); ++ [[maybe_unused]] bool overflow = ++ llvm::MulOverflow(innerStride, outerStride, stride); + assert(!overflow && "slice composition was not checked"); + strides.push_back(stride); + } +diff --git a/lib/Dialect/TensorAttrs.cpp b/lib/Dialect/TensorAttrs.cpp +--- a/lib/Dialect/TensorAttrs.cpp ++++ b/lib/Dialect/TensorAttrs.cpp +@@ -1662,7 +1662,7 @@ MatmulSourceAttr::slice(ArrayRef/var/lock/gpulock${i}_${j} || exit 1 - if flock -n "$lock_fd"; - then - ( - # This export only works within the brackets, so it is isolated to one - # single command. - export CUDA_VISIBLE_DEVICES=$i - export HIP_VISIBLE_DEVICES=$i - echo "Running test $TEST_BINARY $* on GPU $CUDA_VISIBLE_DEVICES" - "$TEST_BINARY" $@ - ) - return_code=$? - flock -u "$lock_fd" - exit $return_code - fi +MAX_ATTEMPTS=30 +for attempt in $(seq 1 $MAX_ATTEMPTS); do + for j in `seq 0 $((TF_TESTS_PER_GPU-1))`; do + for i in `seq 0 $((TF_GPU_COUNT-1))`; do + exec {lock_fd}>"${LOCK_DIR}/gpulock${i}_${j}" || exit 1 + if flock -n "$lock_fd"; + then + ( + # This export only works within the brackets, so it is isolated to one + # single command. + export CUDA_VISIBLE_DEVICES=$i + export HIP_VISIBLE_DEVICES=$i + echo "Running test $TEST_BINARY $* on GPU $CUDA_VISIBLE_DEVICES" + "$TEST_BINARY" $@ + ) + return_code=$? + flock -u "$lock_fd" + exec {lock_fd}>&- + exit $return_code + fi + exec {lock_fd}>&- + done done + sleep 1 done -echo "Cannot find a free GPU to run the test $* on, exiting with failure..." +echo "Cannot find a free GPU to run the test $* on after ${MAX_ATTEMPTS} attempts, exiting with failure..." exit 1 diff --git a/third_party/xla/workspace3.bzl b/third_party/xla/workspace3.bzl index a52a078a5f3840..cb6c2ca5880b15 100644 --- a/third_party/xla/workspace3.bzl +++ b/third_party/xla/workspace3.bzl @@ -74,10 +74,10 @@ def workspace(): # Details: https://github.com/google-ml-infra/rules_ml_toolchain tf_http_archive( name = "rules_ml_toolchain", - sha256 = "75ccb7c4dc699343f02d02f0b5baa6dac09c013d46b92b60f18563448ea59e77", - strip_prefix = "rules_ml_toolchain-73cb731fed3ccf7551beac710bf1c5dbeb8be298", + sha256 = "66d823b6c6b4ea6f09d79ef6fe49e181e5301759ee07d176a6c82a01e1227d6b", + strip_prefix = "rules_ml_toolchain-e8709f15382e4da1de5ca15672cb4412da3d989d", urls = tf_mirror_urls( - "https://github.com/google-ml-infra/rules_ml_toolchain/archive/73cb731fed3ccf7551beac710bf1c5dbeb8be298.tar.gz", + "https://github.com/google-ml-infra/rules_ml_toolchain/archive/e8709f15382e4da1de5ca15672cb4412da3d989d.tar.gz", ), ) diff --git a/third_party/xla/xla/BUILD b/third_party/xla/xla/BUILD index 3f55250582219f..b0f03243321a23 100644 --- a/third_party/xla/xla/BUILD +++ b/third_party/xla/xla/BUILD @@ -522,6 +522,7 @@ cc_library( "@com_google_absl//absl/functional:function_ref", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", + "@com_google_absl//absl/numeric:bits", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", diff --git a/third_party/xla/xla/autotune_cache.proto b/third_party/xla/xla/autotune_cache.proto index 9dd1ab53e13c79..94f4bba186d0bc 100644 --- a/third_party/xla/xla/autotune_cache.proto +++ b/third_party/xla/xla/autotune_cache.proto @@ -28,6 +28,10 @@ message Config { BackendConfig backend_config = 2; } +message CandidateConfigs { + repeated Config configs = 1; +} + message AutotuneTargetKey { string device = 1; // Optional and user specified version, example: xla/jax release version. diff --git a/third_party/xla/xla/backends/autotuner/BUILD b/third_party/xla/xla/backends/autotuner/BUILD index 23945aaf4543f1..34586e5765eb8e 100644 --- a/third_party/xla/xla/backends/autotuner/BUILD +++ b/third_party/xla/xla/backends/autotuner/BUILD @@ -125,6 +125,7 @@ cc_library( deps = [ ":backends_proto_cc", ":codegen_backend", + "//xla:autotune_cache_proto_cc", "//xla/hlo/ir:hlo", "//xla/service:executable", "//xla/stream_executor:kernel_stats", @@ -137,11 +138,13 @@ cc_library( "@com_google_absl//absl/log:check", "@com_google_absl//absl/memory", "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", "@com_google_absl//absl/time", "@com_google_absl//absl/types:span", + "@tsl//tsl/platform:protobuf", ], ) @@ -168,7 +171,9 @@ xla_cc_test( "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_matchers", "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings:str_format", "@com_google_googletest//:gtest_main", + "@tsl//tsl/platform:path", ], ) @@ -220,6 +225,7 @@ cc_library( ":codegen_backend", ":codegen_orchestrator", ":hlo_extractor", + "//xla:autotune_cache_proto_cc", "//xla:autotuning_proto_cc", "//xla:status_macros", "//xla/hlo/ir:hlo", @@ -245,6 +251,7 @@ cc_library( "@com_google_absl//absl/time", "@com_google_absl//absl/types:span", "@tsl//tsl/platform:fingerprint", + "@tsl//tsl/platform:protobuf", ], ) @@ -277,6 +284,7 @@ xla_cc_test( "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_matchers", "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:string_view", "@com_google_absl//absl/time", "@com_google_absl//absl/types:span", diff --git a/third_party/xla/xla/backends/autotuner/codegen_orchestrator.cc b/third_party/xla/xla/backends/autotuner/codegen_orchestrator.cc index 27c084b24fdd17..ea0ff4e3376479 100644 --- a/third_party/xla/xla/backends/autotuner/codegen_orchestrator.cc +++ b/third_party/xla/xla/backends/autotuner/codegen_orchestrator.cc @@ -16,6 +16,7 @@ limitations under the License. #include "xla/backends/autotuner/codegen_orchestrator.h" #include +#include #include #include #include @@ -25,10 +26,13 @@ limitations under the License. #include "absl/log/log.h" #include "absl/memory/memory.h" #include "absl/status/status.h" +#include "absl/status/status_macros.h" #include "absl/status/statusor.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" +#include "absl/strings/string_view.h" #include "absl/types/span.h" +#include "xla/autotune_cache.pb.h" #include "xla/backends/autotuner/backends.pb.h" #include "xla/backends/autotuner/codegen_backend.h" #include "xla/hlo/ir/hlo_instruction.h" @@ -36,8 +40,10 @@ limitations under the License. #include "xla/stream_executor/kernel_stats.h" #include "xla/tsl/concurrency/executor.h" #include "xla/tsl/concurrency/future.h" +#include "xla/tsl/platform/env.h" #include "xla/tsl/platform/errors.h" #include "xla/tsl/platform/threadpool.h" +#include "tsl/platform/protobuf.h" namespace xla { namespace { @@ -50,6 +56,49 @@ absl::Status MakeCombinedConfigError(absl::Span errors) { return absl::InternalError(combined_error); } +absl::StatusOr> LoadCandidateConfigs( + absl::string_view candidate_configs_file, + absl::Span> codegen_backends) { + if (candidate_configs_file.empty()) { + return std::vector{}; + } + std::string content; + absl::Status read_status = tsl::ReadFileToString( + tsl::Env::Default(), std::string(candidate_configs_file), &content); + if (!read_status.ok()) { + return absl::InvalidArgumentError( + absl::StrCat("Failed to read candidate configs file '", + candidate_configs_file, "': ", read_status.message())); + } + autotuner::CandidateConfigs candidate_configs_proto; + bool parsed = tsl::protobuf::TextFormat::ParseFromString( + content, &candidate_configs_proto); + if (!parsed) { + return absl::InvalidArgumentError(absl::StrCat( + "Failed to parse candidate configs file '", candidate_configs_file, + "' as textproto or binary proto.")); + } + std::vector candidate_configs; + candidate_configs.reserve(candidate_configs_proto.configs_size()); + for (const auto& config : candidate_configs_proto.configs()) { + bool backend_found = false; + for (const auto& backend : codegen_backends) { + if (backend->backend() == config.backend()) { + backend_found = true; + break; + } + } + if (!backend_found) { + return absl::InvalidArgumentError( + absl::StrCat("Backend ", Backend_Name(config.backend()), + " in candidate configs is not registered with the " + "orchestrator.")); + } + candidate_configs.push_back(config); + } + return candidate_configs; +} + } // namespace absl::StatusOr> @@ -61,12 +110,31 @@ CodegenOrchestrator::Create( "CodegenOrchestrator initialization failed. No codegen backends " "provided."); } + ABSL_ASSIGN_OR_RETURN( + std::vector candidate_configs, + LoadCandidateConfigs(options.candidate_configs_file, codegen_backends)); return absl::WrapUnique( - new CodegenOrchestrator(std::move(codegen_backends), std::move(options))); + new CodegenOrchestrator(std::move(codegen_backends), std::move(options), + std::move(candidate_configs))); } absl::StatusOr> CodegenOrchestrator::GetSupportedConfigs(const HloInstruction& instr) const { + if (!candidate_configs_.empty()) { + std::vector configs; + configs.reserve(candidate_configs_.size()); + for (const auto& candidate : candidate_configs_) { + for (const auto& codegen_backend : codegen_backends_) { + if (codegen_backend->backend() == candidate.backend()) { + configs.push_back(Config{ + codegen_backend.get(), + std::make_unique(candidate.backend_config())}); + break; + } + } + } + return configs; + } std::vector configs; std::vector errors; for (auto& codegen_backend : codegen_backends_) { @@ -94,6 +162,16 @@ CodegenOrchestrator::GetSupportedConfigs(const HloInstruction& instr) const { absl::StatusOr> CodegenOrchestrator::GetSupportedConfigsWithEstimates( const HloInstruction& instr) const { + if (!candidate_configs_.empty()) { + ABSL_ASSIGN_OR_RETURN(std::vector configs, GetSupportedConfigs(instr)); + std::vector estimated_configs; + estimated_configs.reserve(configs.size()); + for (auto& config : configs) { + estimated_configs.push_back( + EstimatedConfig{std::move(config), std::nullopt}); + } + return estimated_configs; + } std::vector configs; std::vector errors; for (auto& codegen_backend : codegen_backends_) { diff --git a/third_party/xla/xla/backends/autotuner/codegen_orchestrator.h b/third_party/xla/xla/backends/autotuner/codegen_orchestrator.h index 14108ca2ada57b..68e6b57a7589c5 100644 --- a/third_party/xla/xla/backends/autotuner/codegen_orchestrator.h +++ b/third_party/xla/xla/backends/autotuner/codegen_orchestrator.h @@ -26,6 +26,7 @@ limitations under the License. #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/time/time.h" +#include "xla/autotune_cache.pb.h" #include "xla/backends/autotuner/codegen_backend.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/service/executable.h" @@ -44,6 +45,9 @@ class CodegenOrchestrator { std::function allow_reg_spills_fn = [](const HloInstruction&, autotuner::Backend) { return false; }; + // File containing a list of serialized configs to override supported + // configs for all instructions. + std::string candidate_configs_file = ""; }; // TODO(b/444398084): Unify Cache::Config and CodegenOrchestrator::Config @@ -102,9 +106,10 @@ class CodegenOrchestrator { private: CodegenOrchestrator( std::vector> codegen_backends, - Options options) + Options options, std::vector candidate_configs = {}) : codegen_backends_(std::move(codegen_backends)), - options_(std::move(options)) {} + options_(std::move(options)), + candidate_configs_(std::move(candidate_configs)) {} absl::Status IsValidExecutable( const absl::StatusOr>& executable, @@ -112,6 +117,7 @@ class CodegenOrchestrator { std::vector> codegen_backends_; Options options_; + std::vector candidate_configs_; }; } // namespace xla diff --git a/third_party/xla/xla/backends/autotuner/codegen_orchestrator_test.cc b/third_party/xla/xla/backends/autotuner/codegen_orchestrator_test.cc index 225d2ce9b67d66..d9bb8161e47018 100644 --- a/third_party/xla/xla/backends/autotuner/codegen_orchestrator_test.cc +++ b/third_party/xla/xla/backends/autotuner/codegen_orchestrator_test.cc @@ -16,6 +16,8 @@ limitations under the License. #include "xla/backends/autotuner/codegen_orchestrator.h" #include +#include +#include #include #include @@ -24,6 +26,7 @@ limitations under the License. #include "absl/status/status.h" #include "absl/status/status_matchers.h" #include "absl/status/statusor.h" +#include "absl/strings/str_format.h" #include "xla/backends/autotuner/backends.pb.h" #include "xla/backends/autotuner/codegen_backend.h" #include "xla/backends/autotuner/mock_codegen_backend.h" @@ -40,6 +43,7 @@ limitations under the License. #include "xla/tsl/platform/env.h" #include "xla/tsl/platform/test.h" #include "xla/tsl/platform/threadpool.h" +#include "tsl/platform/path.h" namespace xla { namespace { @@ -102,6 +106,108 @@ TEST_F(CodegenOrchestratorTest, GetSupportedConfigsAggregatesFromAllBackends) { EXPECT_THAT(*supported[1].backend_config, ConfigMatcher("test_config_2")); } +TEST_F(CodegenOrchestratorTest, CandidateConfigsFileOverridesSupportedConfigs) { + std::string temp_file = + tsl::io::JoinPath(tsl::testing::TmpDir(), "candidate_configs.pbtxt"); + std::string file_content = absl::StrFormat( + R"pb( + configs { + backend: UNSPECIFIED_BACKEND + backend_config { gemm { algorithm: %d } } + } + configs { + backend: UNSPECIFIED_BACKEND + backend_config { gemm { algorithm: %d } } + } + )pb", + GetAlgorithmId("test_candidate_1"), GetAlgorithmId("test_candidate_2")); + ASSERT_OK( + tsl::WriteStringToFile(tsl::Env::Default(), temp_file, file_content)); + + auto backend = std::make_unique(); + EXPECT_CALL(*backend, name()).WillRepeatedly(Return("mock_backend")); + EXPECT_CALL(*backend, backend()) + .WillRepeatedly(Return(autotuner::Backend::UNSPECIFIED_BACKEND)); + EXPECT_CALL(*backend, GetSupportedConfigs(_)).Times(0); + EXPECT_CALL(*backend, GetSupportedConfigsWithEstimates(_)).Times(0); + + std::vector> backends; + backends.push_back(std::move(backend)); + + CodegenOrchestrator::Options options; + options.candidate_configs_file = temp_file; + + ASSERT_OK_AND_ASSIGN(auto orchestrator, CodegenOrchestrator::Create( + std::move(backends), options)); + + auto dummy_instr = HloInstruction::CreateConstant(LiteralUtil::CreateR0(1)); + ASSERT_OK_AND_ASSIGN(auto supported, + orchestrator->GetSupportedConfigs(*dummy_instr)); + ASSERT_THAT(supported, SizeIs(2)); + EXPECT_THAT(*supported[0].backend_config, ConfigMatcher("test_candidate_1")); + EXPECT_THAT(*supported[1].backend_config, ConfigMatcher("test_candidate_2")); + + ASSERT_OK_AND_ASSIGN( + auto supported_with_estimates, + orchestrator->GetSupportedConfigsWithEstimates(*dummy_instr)); + ASSERT_THAT(supported_with_estimates, SizeIs(2)); + EXPECT_THAT(*supported_with_estimates[0].config.backend_config, + ConfigMatcher("test_candidate_1")); + EXPECT_EQ(supported_with_estimates[0].estimated_runtime, std::nullopt); + EXPECT_THAT(*supported_with_estimates[1].config.backend_config, + ConfigMatcher("test_candidate_2")); + EXPECT_EQ(supported_with_estimates[1].estimated_runtime, std::nullopt); +} + +TEST_F(CodegenOrchestratorTest, + CandidateConfigsFileNotFoundReturnsInvalidArgument) { + auto backend = std::make_unique(); + std::vector> backends; + backends.push_back(std::move(backend)); + + CodegenOrchestrator::Options options; + options.candidate_configs_file = "/nonexistent/path/configs.pbtxt"; + EXPECT_THAT(CodegenOrchestrator::Create(std::move(backends), options), + StatusIs(absl::StatusCode::kInvalidArgument)); +} + +TEST_F(CodegenOrchestratorTest, + CandidateConfigsFileInvalidProtoReturnsInvalidArgument) { + std::string temp_file = + tsl::io::JoinPath(tsl::testing::TmpDir(), "invalid_configs.pbtxt"); + ASSERT_OK(tsl::WriteStringToFile(tsl::Env::Default(), temp_file, + "invalid { proto content")); + + auto backend = std::make_unique(); + std::vector> backends; + backends.push_back(std::move(backend)); + + CodegenOrchestrator::Options options; + options.candidate_configs_file = temp_file; + EXPECT_THAT(CodegenOrchestrator::Create(std::move(backends), options), + StatusIs(absl::StatusCode::kInvalidArgument)); +} + +TEST_F(CodegenOrchestratorTest, + CandidateConfigsFileUnregisteredBackendReturnsInvalidArgument) { + std::string temp_file = + tsl::io::JoinPath(tsl::testing::TmpDir(), "unregistered_backend.pbtxt"); + ASSERT_OK(tsl::WriteStringToFile( + tsl::Env::Default(), temp_file, + "configs { backend: TRITON backend_config { gemm { algorithm: 1 } } }")); + + auto backend = std::make_unique(); + EXPECT_CALL(*backend, backend()) + .WillRepeatedly(Return(autotuner::Backend::UNSPECIFIED_BACKEND)); + std::vector> backends; + backends.push_back(std::move(backend)); + + CodegenOrchestrator::Options options; + options.candidate_configs_file = temp_file; + EXPECT_THAT(CodegenOrchestrator::Create(std::move(backends), options), + StatusIs(absl::StatusCode::kInvalidArgument)); +} + TEST_F(CodegenOrchestratorTest, GetSupportedConfigsToleratesPartialBackendFailure) { std::vector> configs; diff --git a/third_party/xla/xla/backends/autotuner/config_assigner.cc b/third_party/xla/xla/backends/autotuner/config_assigner.cc index ead6d437065497..4d21fa148cc882 100644 --- a/third_party/xla/xla/backends/autotuner/config_assigner.cc +++ b/third_party/xla/xla/backends/autotuner/config_assigner.cc @@ -39,6 +39,7 @@ limitations under the License. #include "absl/strings/string_view.h" #include "absl/time/time.h" #include "absl/types/span.h" +#include "xla/autotune_cache.pb.h" #include "xla/autotuning.pb.h" #include "xla/backends/autotuner/autotuner.h" #include "xla/backends/autotuner/autotuner_cache_interface.h" @@ -60,6 +61,7 @@ limitations under the License. #include "xla/tsl/platform/errors.h" #include "xla/tsl/platform/threadpool.h" #include "tsl/platform/fingerprint.h" +#include "tsl/platform/protobuf.h" namespace xla { namespace { @@ -109,6 +111,35 @@ std::string GetKvStoreKey( backend_fingerprint, "_", shard_index); } +absl::StatusOr> ParseForcedConfig( + absl::string_view force_config, const CodegenOrchestrator& orchestrator) { + if (force_config.empty()) { + return std::nullopt; + } + autotuner::Config config_proto; + bool parsed = + tsl::protobuf::TextFormat::ParseFromString(force_config, &config_proto); + if (!parsed) { + return absl::InvalidArgumentError(absl::StrCat( + "Failed to parse force_config as textproto: ", force_config)); + } + CodegenBackend* matched_backend = nullptr; + for (const auto& backend : orchestrator.codegen_backends()) { + if (backend->backend() == config_proto.backend()) { + matched_backend = backend.get(); + break; + } + } + if (matched_backend == nullptr) { + return absl::InvalidArgumentError(absl::StrCat( + "Backend ", Backend_Name(config_proto.backend()), + " in force_config is not registered with the orchestrator.")); + } + return ConfigAssigner::Config{ + matched_backend, + std::make_unique(config_proto.backend_config())}; +} + } // namespace absl::StatusOr> ConfigAssigner::Create( @@ -117,9 +148,11 @@ absl::StatusOr> ConfigAssigner::Create( std::unique_ptr absl_nonnull orchestrator, std::unique_ptr absl_nullable autotuner, tsl::thread::ThreadPool* thread_pool) { + ABSL_ASSIGN_OR_RETURN(std::optional forced_config, + ParseForcedConfig(options.force_config, *orchestrator)); return absl::WrapUnique(new ConfigAssigner( std::move(options), std::move(cache), std::move(orchestrator), - std::move(autotuner), thread_pool)); + std::move(autotuner), thread_pool, std::move(forced_config))); } absl::Status ConfigAssigner::AssignConfigs( @@ -298,6 +331,14 @@ tsl::Future ConfigAssigner::GetConfig( } VLOG(1) << "Getting config for HLO: " << instr->ToString(print_options); } + if (forced_config_.has_value()) { + VLOG(1) << "Using forced config: " + << forced_config_->codegen_backend->name() << " : " + << forced_config_->backend_config->ShortDebugString(); + return Config{ + forced_config_->codegen_backend, + std::make_unique(*forced_config_->backend_config)}; + } std::optional cached_config = LookUp(instr); if (cached_config.has_value()) { VLOG(1) << "Using cached config: " << cached_config->ToString(); @@ -547,11 +588,12 @@ std::string ConfigAssigner::Options::ToString() const { "prefer_estimated_configs": %v, "dump_hlos": %v, "use_new_cache_format": %v, - "compile_all_supported_configs": %v + "compile_all_supported_configs": %v, + "force_config": "%s" })json", expect_all_instructions_in_cache, allow_autotuning, prefer_estimated_configs, dump_hlos, use_new_cache_format, - compile_all_supported_configs); + compile_all_supported_configs, force_config); } AutotunerCacheInterface::CacheStats ConfigAssigner::GetCacheStats() const { diff --git a/third_party/xla/xla/backends/autotuner/config_assigner.h b/third_party/xla/xla/backends/autotuner/config_assigner.h index 14bc3ce7173f53..27ec7b25ff5064 100644 --- a/third_party/xla/xla/backends/autotuner/config_assigner.h +++ b/third_party/xla/xla/backends/autotuner/config_assigner.h @@ -61,6 +61,9 @@ class ConfigAssigner { // When autotuning is disabled, if true, compiles all supported configs in // parallel before returning the first successful one. bool compile_all_supported_configs = false; + // Single serialized config to override config of all instructions, + // bypassing cache and autotuning. + std::string force_config = ""; std::string ToString() const; }; @@ -94,12 +97,14 @@ class ConfigAssigner { absl_nonnull std::unique_ptr cache, absl_nonnull std::unique_ptr orchestrator, absl_nullable std::unique_ptr autotuner, - tsl::thread::ThreadPool* thread_pool) + tsl::thread::ThreadPool* thread_pool, + std::optional forced_config = std::nullopt) : options_(options), optimal_config_cache_(std::move(cache)), orchestrator_(std::move(orchestrator)), autotuner_(std::move(autotuner)), - thread_pool_(thread_pool) {} + thread_pool_(thread_pool), + forced_config_(std::move(forced_config)) {} using InstructionGroup = std::vector; @@ -151,6 +156,7 @@ class ConfigAssigner { absl_nullable std::unique_ptr autotuner_; tsl::thread::ThreadPool* thread_pool_ = nullptr; int dump_counter_ = 0; + std::optional forced_config_; }; } // namespace xla diff --git a/third_party/xla/xla/backends/autotuner/config_assigner_test.cc b/third_party/xla/xla/backends/autotuner/config_assigner_test.cc index a77a9633983ad8..6096857bf1dbb9 100644 --- a/third_party/xla/xla/backends/autotuner/config_assigner_test.cc +++ b/third_party/xla/xla/backends/autotuner/config_assigner_test.cc @@ -26,6 +26,7 @@ limitations under the License. #include "absl/status/status.h" #include "absl/status/status_matchers.h" #include "absl/status/statusor.h" +#include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" #include "absl/time/time.h" #include "absl/types/span.h" @@ -244,6 +245,62 @@ TEST_F(ConfigAssignerTest, CacheHit) { EXPECT_THAT(config_assigner->AssignConfig(dummy_instr.get()), IsOk()); } +TEST_F(ConfigAssignerTest, ForceConfigBypassesCacheAndAutotuning) { + auto cache_manager = std::make_unique(); + EXPECT_CALL(*cache_manager, Lookup(_)).Times(0); + EXPECT_CALL(*cache_manager, Insert(_, _)).Times(0); + + auto backend = std::make_unique(); + EXPECT_CALL(*backend, name()).WillRepeatedly(Return("mock_backend")); + EXPECT_CALL(*backend, backend()) + .WillRepeatedly(Return(autotuner::Backend::UNSPECIFIED_BACKEND)); + EXPECT_CALL(*backend, GetSupportedConfigs).Times(0); + EXPECT_CALL(*backend, ApplyConfig(_, ConfigMatcher("test_forced_config"))) + .Times(1) + .WillOnce(Return(absl::OkStatus())); + + std::vector> backends; + backends.push_back(std::move(backend)); + + config_.force_config = absl::StrCat( + "backend: UNSPECIFIED_BACKEND\nbackend_config { gemm { algorithm: ", + GetAlgorithmId("test_forced_config"), " } }"); + + auto profiler = std::make_unique(); + ASSERT_OK_AND_ASSIGN( + auto config_assigner, + CreateConfigAssigner(std::move(backends), std::move(profiler), config_, + std::move(cache_manager))); + + auto dummy_instr = HloInstruction::CreateConstant(LiteralUtil::CreateR0(1)); + EXPECT_THAT(config_assigner->AssignConfig(dummy_instr.get()), IsOk()); +} + +TEST_F(ConfigAssignerTest, ForceConfigInvalidProtoFails) { + auto backend = std::make_unique(); + std::vector> backends; + backends.push_back(std::move(backend)); + + config_.force_config = "invalid proto content {"; + EXPECT_THAT( + CreateConfigAssigner(std::move(backends), nullptr, config_, nullptr), + StatusIs(absl::StatusCode::kInvalidArgument)); +} + +TEST_F(ConfigAssignerTest, ForceConfigUnregisteredBackendFails) { + auto backend = std::make_unique(); + EXPECT_CALL(*backend, backend()) + .WillRepeatedly(Return(autotuner::Backend::UNSPECIFIED_BACKEND)); + std::vector> backends; + backends.push_back(std::move(backend)); + + config_.force_config = + "backend: TRITON\nbackend_config { gemm { algorithm: 1 } }"; + EXPECT_THAT( + CreateConfigAssigner(std::move(backends), nullptr, config_, nullptr), + StatusIs(absl::StatusCode::kInvalidArgument)); +} + TEST_F(ConfigAssignerTest, ExpectAllInstructionsInCache) { auto cache_manager = std::make_unique(); EXPECT_CALL(*cache_manager, Lookup(_)).WillOnce(Return(std::nullopt)); @@ -949,7 +1006,8 @@ TEST(ConfigAssignerOptionsTest, ToString) { " \"prefer_estimated_configs\": true,\n" " \"dump_hlos\": false,\n" " \"use_new_cache_format\": false,\n" - " \"compile_all_supported_configs\": false\n" + " \"compile_all_supported_configs\": false,\n" + " \"force_config\": \"\"\n" "}"; EXPECT_EQ(config.ToString(), expected); } diff --git a/third_party/xla/xla/backends/cpu/benchmarks/e2e/keras/benchmark.py b/third_party/xla/xla/backends/cpu/benchmarks/e2e/keras/benchmark.py index 53130c3574cbd5..ca1c0d2ebc9c33 100644 --- a/third_party/xla/xla/backends/cpu/benchmarks/e2e/keras/benchmark.py +++ b/third_party/xla/xla/backends/cpu/benchmarks/e2e/keras/benchmark.py @@ -115,8 +115,9 @@ def main(): default="gemma2_2b_en", help=( "Preset name of the model to benchmark. This was tested with" - " the following models: gemma2_2b_en, gemma3_1b, gemma4_2b" - " gpt_oss_20b_en, mixtral_8_7b_en, qwen3_14b_en, qwen2_1.5b_en." + " the following models: gemma2_2b_en, gemma3_1b, gemma4_2b," + " gpt_oss_20b_en, mistral_7b_en mixtral_8_7b_en, " + " qwen3_14b_en, qwen2_1.5b_en." " Llama models also should work." " See https://keras.io/keras_hub/presets/" " for the full list of presets." @@ -149,6 +150,9 @@ def main(): elif "llama" in model_name: lm = keras_hub.models.LlamaCausalLM.from_preset(model_name) tokenizer = keras_hub.models.LlamaTokenizer.from_preset(model_name) + elif "mistral" in model_name: + lm = keras_hub.models.MistralCausalLM.from_preset(model_name) + tokenizer = keras_hub.models.MistralTokenizer.from_preset(model_name) elif "mixtral" in model_name: lm = keras_hub.models.MixtralCausalLM.from_preset(model_name) tokenizer = keras_hub.models.MixtralTokenizer.from_preset(model_name) diff --git a/third_party/xla/xla/backends/cpu/benchmarks/huge_hlo/hlo_files.bzl b/third_party/xla/xla/backends/cpu/benchmarks/huge_hlo/hlo_files.bzl index 0aa16e84b02252..0a78d5a1b57134 100644 --- a/third_party/xla/xla/backends/cpu/benchmarks/huge_hlo/hlo_files.bzl +++ b/third_party/xla/xla/backends/cpu/benchmarks/huge_hlo/hlo_files.bzl @@ -27,5 +27,6 @@ HUGE_HLO_FILES = [ "gemma3_4b_text_keras_jax_batch1_in8_out100.hlo", "gemma4_2b_bf16.hlo", "gemma4_4b_text_keras_jax_bf16.hlo", + "mistral_7b_en_keras_jax_bf16.hlo", # go/keep-sorted end ] diff --git a/third_party/xla/xla/backends/cpu/benchmarks/huge_hlo/mistral_7b_en_keras_jax_bf16.hlo b/third_party/xla/xla/backends/cpu/benchmarks/huge_hlo/mistral_7b_en_keras_jax_bf16.hlo new file mode 100644 index 00000000000000..a2923a99bfe444 --- /dev/null +++ b/third_party/xla/xla/backends/cpu/benchmarks/huge_hlo/mistral_7b_en_keras_jax_bf16.hlo @@ -0,0 +1,13705 @@ +HloModule jit_compiled_generate_function, entry_computation_layout={(pred[1,41]{1,0}, s32[1,41]{1,0}, u32[2]{0}, bf16[32000,4096]{1,0}, bf16[4096,32000]{1,0}, /*index=5*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=10*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=15*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=20*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=25*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=30*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=35*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=40*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=45*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=50*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=55*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=60*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=65*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=70*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=75*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=80*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=85*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=90*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=95*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=100*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=105*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=110*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=115*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=120*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=125*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=130*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=135*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=140*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=145*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=150*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=155*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=160*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=165*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=170*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=175*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=180*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=185*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=190*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=195*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=200*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=205*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=210*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=215*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=220*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=225*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=230*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=235*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=240*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=245*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=250*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=255*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=260*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=265*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=270*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=275*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=280*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=285*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=290*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096]{0})->(pred[1,41]{1,0}, s32[1,41]{1,0}, u32[2]{0})}, allow_spmd_sharding_propagation_to_parameters={true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true,true}, allow_spmd_sharding_propagation_to_output={true,true,true} + +%_where.1 (Arg_0.2: pred[1,41], Arg_1.2: s32[1,41], Arg_2.1: s32[1,41]) -> s32[1,41] { + %Arg_0.2 = pred[1,41]{1,0} parameter(0) + %Arg_1.2 = s32[1,41]{1,0} parameter(1) + %Arg_2.1 = s32[1,41]{1,0} parameter(2) + ROOT %select_n.1 = s32[1,41]{1,0} select(%Arg_0.2, %Arg_1.2, %Arg_2.1), metadata={op_name="select_n"} +} + +%region_0.2 (reduce_and.3: pred[], reduce_and.4: pred[]) -> pred[] { + %reduce_and.3 = pred[] parameter(0), metadata={op_name="reduce_and"} + %reduce_and.4 = pred[] parameter(1), metadata={op_name="reduce_and"} + ROOT %reduce_and.5 = pred[] and(%reduce_and.3, %reduce_and.4), metadata={op_name="reduce_and"} +} + +%_take.3 (Arg_0.3: bf16[32000,4096], Arg_1.3: s32[1,41]) -> bf16[1,41,4096] { + %Arg_1.3 = s32[1,41]{1,0} parameter(1) + %constant.30 = s32[] constant(0) + %lt.2 = s32[1,41]{1,0} broadcast(%constant.30), dimensions={}, metadata={op_name="lt"} + %lt.3 = pred[1,41]{1,0} compare(%Arg_1.3, %lt.2), direction=LT, metadata={op_name="lt"} + %constant.28 = s32[] constant(32000) + %add.2 = s32[1,41]{1,0} broadcast(%constant.28), dimensions={}, metadata={op_name="add"} + %add.3 = s32[1,41]{1,0} add(%Arg_1.3, %add.2), metadata={op_name="add"} + %jit__where_.1 = s32[1,41]{1,0} call(%lt.3, %add.3, %Arg_1.3), to_apply=%_where.1, metadata={op_name="jit(_where)"} + %broadcast_in_dim.4 = s32[1,41,1]{2,1,0} reshape(%jit__where_.1), metadata={op_name="broadcast_in_dim"} + %constant.27 = s32[] constant(0) + %ge.2 = s32[1,41,1]{2,1,0} broadcast(%constant.27), dimensions={}, metadata={op_name="ge"} + %ge.3 = pred[1,41,1]{2,1,0} compare(%broadcast_in_dim.4, %ge.2), direction=GE, metadata={op_name="ge"} + %constant.26 = s32[] constant(31999) + %le.2 = s32[1,41,1]{2,1,0} broadcast(%constant.26), dimensions={}, metadata={op_name="le"} + %le.3 = pred[1,41,1]{2,1,0} compare(%broadcast_in_dim.4, %le.2), direction=LE, metadata={op_name="le"} + %and.1 = pred[1,41,1]{2,1,0} and(%ge.3, %le.3), metadata={op_name="and"} + %constant.29 = pred[] constant(true) + %reduce_and.7 = pred[1,41]{1,0} reduce(%and.1, %constant.29), dimensions={2}, to_apply=%region_0.2, metadata={op_name="reduce_and"} + %broadcast_in_dim.5 = pred[1,41,4096]{2,1,0} broadcast(%reduce_and.7), dimensions={0,1}, metadata={op_name="broadcast_in_dim"} + %Arg_0.3 = bf16[32000,4096]{1,0} parameter(0) + %gather.1 = bf16[1,41,4096]{2,1,0} gather(%Arg_0.3, %broadcast_in_dim.4), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,4096}, metadata={op_name="gather"} + %constant.25 = bf16[] constant(nan) + %broadcast_in_dim.3 = bf16[1,41,4096]{2,1,0} broadcast(%constant.25), dimensions={}, metadata={op_name="broadcast_in_dim"} + ROOT %select_n.3 = bf16[1,41,4096]{2,1,0} select(%broadcast_in_dim.5, %gather.1, %broadcast_in_dim.3), metadata={op_name="select_n"} +} + +%region_1.4 (reduce_sum.3: f32[], reduce_sum.4: f32[]) -> f32[] { + %reduce_sum.3 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.4 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.5 = f32[] add(%reduce_sum.3, %reduce_sum.4), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention_layernorm/reduce_sum" stack_frame_id=24} +} + +%_where_0.5 (Arg_0.5: pred[1,1,32,41,41], Arg_1.5: f32[1,32,41,1,41], Arg_2.3: f32[]) -> f32[1,1,32,41,41] { + %Arg_0.5 = pred[1,1,32,41,41]{4,3,2,1,0} parameter(0) + %Arg_1.5 = f32[1,32,41,1,41]{4,3,2,1,0} parameter(1) + %transpose.1 = f32[1,1,32,41,41]{4,3,2,1,0} reshape(%Arg_1.5), metadata={op_name="transpose"} + %Arg_2.3 = f32[] parameter(2) + %broadcast_in_dim.21 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%Arg_2.3), dimensions={}, metadata={op_name="broadcast_in_dim"} + ROOT %select_n.5 = f32[1,1,32,41,41]{4,3,2,1,0} select(%Arg_0.5, %transpose.1, %broadcast_in_dim.21), metadata={op_name="select_n"} +} + +%region_2.6 (reduce_max.3: f32[], reduce_max.4: f32[]) -> f32[] { + %reduce_max.3 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.4 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.5 = f32[] maximum(%reduce_max.3, %reduce_max.4), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/vmap()/reduce_max" stack_frame_id=34} +} + +%region_3.7 (reduce_sum.10: f32[], reduce_sum.11: f32[]) -> f32[] { + %reduce_sum.10 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.11 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.12 = f32[] add(%reduce_sum.10, %reduce_sum.11), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/vmap()/reduce_sum" stack_frame_id=34} +} + +%region_4.8 (reduce_sum.17: f32[], reduce_sum.18: f32[]) -> f32[] { + %reduce_sum.17 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.18 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.19 = f32[] add(%reduce_sum.17, %reduce_sum.18), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/feedforward_layernorm/reduce_sum" stack_frame_id=43} +} + +%silu.9 (Arg_0.7: f32[1,41,14336]) -> f32[1,41,14336] { + %Arg_0.7 = f32[1,41,14336]{2,1,0} parameter(0) + %constant.32 = f32[] constant(1) + %broadcast.13 = f32[1,41,14336]{2,1,0} broadcast(%constant.32), dimensions={} + %neg.5 = f32[1,41,14336]{2,1,0} negate(%Arg_0.7), metadata={op_name="neg"} + %exp.2 = f32[1,41,14336]{2,1,0} exponential(%neg.5), metadata={op_name="exp"} + %add.11 = f32[1,41,14336]{2,1,0} add(%exp.2, %broadcast.13), metadata={op_name="add"} + %div.11 = f32[1,41,14336]{2,1,0} divide(%broadcast.13, %add.11), metadata={op_name="div"} + ROOT %mul.36 = f32[1,41,14336]{2,1,0} multiply(%Arg_0.7, %div.11), metadata={op_name="mul"} +} + +%region_5.10 (reduce_sum.24: f32[], reduce_sum.25: f32[]) -> f32[] { + %reduce_sum.24 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.25 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.26 = f32[] add(%reduce_sum.24, %reduce_sum.25), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention_layernorm/reduce_sum" stack_frame_id=24} +} + +%region_6.11 (reduce_max.10: f32[], reduce_max.11: f32[]) -> f32[] { + %reduce_max.10 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.11 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.12 = f32[] maximum(%reduce_max.10, %reduce_max.11), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/vmap()/reduce_max" stack_frame_id=34} +} + +%region_7.12 (reduce_sum.31: f32[], reduce_sum.32: f32[]) -> f32[] { + %reduce_sum.31 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.32 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.33 = f32[] add(%reduce_sum.31, %reduce_sum.32), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/vmap()/reduce_sum" stack_frame_id=34} +} + +%region_8.13 (reduce_sum.38: f32[], reduce_sum.39: f32[]) -> f32[] { + %reduce_sum.38 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.39 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.40 = f32[] add(%reduce_sum.38, %reduce_sum.39), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/feedforward_layernorm/reduce_sum" stack_frame_id=43} +} + +%region_9.14 (reduce_sum.45: f32[], reduce_sum.46: f32[]) -> f32[] { + %reduce_sum.45 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.46 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.47 = f32[] add(%reduce_sum.45, %reduce_sum.46), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention_layernorm/reduce_sum" stack_frame_id=24} +} + +%region_10.15 (reduce_max.17: f32[], reduce_max.18: f32[]) -> f32[] { + %reduce_max.17 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.18 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.19 = f32[] maximum(%reduce_max.17, %reduce_max.18), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/vmap()/reduce_max" stack_frame_id=34} +} + +%region_11.16 (reduce_sum.52: f32[], reduce_sum.53: f32[]) -> f32[] { + %reduce_sum.52 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.53 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.54 = f32[] add(%reduce_sum.52, %reduce_sum.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/vmap()/reduce_sum" stack_frame_id=34} +} + +%region_12.17 (reduce_sum.59: f32[], reduce_sum.60: f32[]) -> f32[] { + %reduce_sum.59 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.60 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.61 = f32[] add(%reduce_sum.59, %reduce_sum.60), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/feedforward_layernorm/reduce_sum" stack_frame_id=43} +} + +%region_13.18 (reduce_sum.66: f32[], reduce_sum.67: f32[]) -> f32[] { + %reduce_sum.66 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.67 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.68 = f32[] add(%reduce_sum.66, %reduce_sum.67), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention_layernorm/reduce_sum" stack_frame_id=24} +} + +%region_14.19 (reduce_max.24: f32[], reduce_max.25: f32[]) -> f32[] { + %reduce_max.24 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.25 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.26 = f32[] maximum(%reduce_max.24, %reduce_max.25), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/vmap()/reduce_max" stack_frame_id=34} +} + +%region_15.20 (reduce_sum.73: f32[], reduce_sum.74: f32[]) -> f32[] { + %reduce_sum.73 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.74 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.75 = f32[] add(%reduce_sum.73, %reduce_sum.74), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/vmap()/reduce_sum" stack_frame_id=34} +} + +%region_16.21 (reduce_sum.80: f32[], reduce_sum.81: f32[]) -> f32[] { + %reduce_sum.80 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.81 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.82 = f32[] add(%reduce_sum.80, %reduce_sum.81), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/feedforward_layernorm/reduce_sum" stack_frame_id=43} +} + +%region_17.22 (reduce_sum.87: f32[], reduce_sum.88: f32[]) -> f32[] { + %reduce_sum.87 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.88 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.89 = f32[] add(%reduce_sum.87, %reduce_sum.88), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention_layernorm/reduce_sum" stack_frame_id=24} +} + +%region_18.23 (reduce_max.31: f32[], reduce_max.32: f32[]) -> f32[] { + %reduce_max.31 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.32 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.33 = f32[] maximum(%reduce_max.31, %reduce_max.32), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/vmap()/reduce_max" stack_frame_id=34} +} + +%region_19.24 (reduce_sum.94: f32[], reduce_sum.95: f32[]) -> f32[] { + %reduce_sum.94 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.95 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.96 = f32[] add(%reduce_sum.94, %reduce_sum.95), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/vmap()/reduce_sum" stack_frame_id=34} +} + +%region_20.25 (reduce_sum.101: f32[], reduce_sum.102: f32[]) -> f32[] { + %reduce_sum.101 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.102 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.103 = f32[] add(%reduce_sum.101, %reduce_sum.102), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/feedforward_layernorm/reduce_sum" stack_frame_id=43} +} + +%region_21.26 (reduce_sum.108: f32[], reduce_sum.109: f32[]) -> f32[] { + %reduce_sum.108 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.109 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.110 = f32[] add(%reduce_sum.108, %reduce_sum.109), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention_layernorm/reduce_sum" stack_frame_id=24} +} + +%region_22.27 (reduce_max.38: f32[], reduce_max.39: f32[]) -> f32[] { + %reduce_max.38 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.39 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.40 = f32[] maximum(%reduce_max.38, %reduce_max.39), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/vmap()/reduce_max" stack_frame_id=34} +} + +%region_23.28 (reduce_sum.115: f32[], reduce_sum.116: f32[]) -> f32[] { + %reduce_sum.115 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.116 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.117 = f32[] add(%reduce_sum.115, %reduce_sum.116), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/vmap()/reduce_sum" stack_frame_id=34} +} + +%region_24.29 (reduce_sum.122: f32[], reduce_sum.123: f32[]) -> f32[] { + %reduce_sum.122 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.123 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.124 = f32[] add(%reduce_sum.122, %reduce_sum.123), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/feedforward_layernorm/reduce_sum" stack_frame_id=43} +} + +%region_25.30 (reduce_sum.129: f32[], reduce_sum.130: f32[]) -> f32[] { + %reduce_sum.129 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.130 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.131 = f32[] add(%reduce_sum.129, %reduce_sum.130), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention_layernorm/reduce_sum" stack_frame_id=24} +} + +%region_26.31 (reduce_max.45: f32[], reduce_max.46: f32[]) -> f32[] { + %reduce_max.45 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.46 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.47 = f32[] maximum(%reduce_max.45, %reduce_max.46), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/vmap()/reduce_max" stack_frame_id=34} +} + +%region_27.32 (reduce_sum.136: f32[], reduce_sum.137: f32[]) -> f32[] { + %reduce_sum.136 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.137 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.138 = f32[] add(%reduce_sum.136, %reduce_sum.137), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/vmap()/reduce_sum" stack_frame_id=34} +} + +%region_28.33 (reduce_sum.143: f32[], reduce_sum.144: f32[]) -> f32[] { + %reduce_sum.143 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.144 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.145 = f32[] add(%reduce_sum.143, %reduce_sum.144), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/feedforward_layernorm/reduce_sum" stack_frame_id=43} +} + +%region_29.34 (reduce_sum.150: f32[], reduce_sum.151: f32[]) -> f32[] { + %reduce_sum.150 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.151 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.152 = f32[] add(%reduce_sum.150, %reduce_sum.151), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention_layernorm/reduce_sum" stack_frame_id=24} +} + +%region_30.35 (reduce_max.52: f32[], reduce_max.53: f32[]) -> f32[] { + %reduce_max.52 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.53 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.54 = f32[] maximum(%reduce_max.52, %reduce_max.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/vmap()/reduce_max" stack_frame_id=34} +} + +%region_31.36 (reduce_sum.157: f32[], reduce_sum.158: f32[]) -> f32[] { + %reduce_sum.157 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.158 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.159 = f32[] add(%reduce_sum.157, %reduce_sum.158), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/vmap()/reduce_sum" stack_frame_id=34} +} + +%region_32.37 (reduce_sum.164: f32[], reduce_sum.165: f32[]) -> f32[] { + %reduce_sum.164 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.165 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.166 = f32[] add(%reduce_sum.164, %reduce_sum.165), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/feedforward_layernorm/reduce_sum" stack_frame_id=43} +} + +%region_33.38 (reduce_sum.171: f32[], reduce_sum.172: f32[]) -> f32[] { + %reduce_sum.171 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.172 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.173 = f32[] add(%reduce_sum.171, %reduce_sum.172), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention_layernorm/reduce_sum" stack_frame_id=24} +} + +%region_34.39 (reduce_max.59: f32[], reduce_max.60: f32[]) -> f32[] { + %reduce_max.59 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.60 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.61 = f32[] maximum(%reduce_max.59, %reduce_max.60), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/vmap()/reduce_max" stack_frame_id=34} +} + +%region_35.40 (reduce_sum.178: f32[], reduce_sum.179: f32[]) -> f32[] { + %reduce_sum.178 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.179 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.180 = f32[] add(%reduce_sum.178, %reduce_sum.179), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/vmap()/reduce_sum" stack_frame_id=34} +} + +%region_36.41 (reduce_sum.185: f32[], reduce_sum.186: f32[]) -> f32[] { + %reduce_sum.185 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.186 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.187 = f32[] add(%reduce_sum.185, %reduce_sum.186), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/feedforward_layernorm/reduce_sum" stack_frame_id=43} +} + +%region_37.42 (reduce_sum.192: f32[], reduce_sum.193: f32[]) -> f32[] { + %reduce_sum.192 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.193 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.194 = f32[] add(%reduce_sum.192, %reduce_sum.193), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention_layernorm/reduce_sum" stack_frame_id=24} +} + +%region_38.43 (reduce_max.66: f32[], reduce_max.67: f32[]) -> f32[] { + %reduce_max.66 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.67 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.68 = f32[] maximum(%reduce_max.66, %reduce_max.67), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/vmap()/reduce_max" stack_frame_id=34} +} + +%region_39.44 (reduce_sum.199: f32[], reduce_sum.200: f32[]) -> f32[] { + %reduce_sum.199 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.200 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.201 = f32[] add(%reduce_sum.199, %reduce_sum.200), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/vmap()/reduce_sum" stack_frame_id=34} +} + +%region_40.45 (reduce_sum.206: f32[], reduce_sum.207: f32[]) -> f32[] { + %reduce_sum.206 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.207 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.208 = f32[] add(%reduce_sum.206, %reduce_sum.207), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/feedforward_layernorm/reduce_sum" stack_frame_id=43} +} + +%region_41.46 (reduce_sum.213: f32[], reduce_sum.214: f32[]) -> f32[] { + %reduce_sum.213 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.214 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.215 = f32[] add(%reduce_sum.213, %reduce_sum.214), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention_layernorm/reduce_sum" stack_frame_id=24} +} + +%region_42.47 (reduce_max.73: f32[], reduce_max.74: f32[]) -> f32[] { + %reduce_max.73 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.74 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.75 = f32[] maximum(%reduce_max.73, %reduce_max.74), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/vmap()/reduce_max" stack_frame_id=34} +} + +%region_43.48 (reduce_sum.220: f32[], reduce_sum.221: f32[]) -> f32[] { + %reduce_sum.220 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.221 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.222 = f32[] add(%reduce_sum.220, %reduce_sum.221), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/vmap()/reduce_sum" stack_frame_id=34} +} + +%region_44.49 (reduce_sum.227: f32[], reduce_sum.228: f32[]) -> f32[] { + %reduce_sum.227 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.228 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.229 = f32[] add(%reduce_sum.227, %reduce_sum.228), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/feedforward_layernorm/reduce_sum" stack_frame_id=43} +} + +%region_45.50 (reduce_sum.234: f32[], reduce_sum.235: f32[]) -> f32[] { + %reduce_sum.234 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.235 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.236 = f32[] add(%reduce_sum.234, %reduce_sum.235), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention_layernorm/reduce_sum" stack_frame_id=24} +} + +%region_46.51 (reduce_max.80: f32[], reduce_max.81: f32[]) -> f32[] { + %reduce_max.80 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.81 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.82 = f32[] maximum(%reduce_max.80, %reduce_max.81), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/vmap()/reduce_max" stack_frame_id=34} +} + +%region_47.52 (reduce_sum.241: f32[], reduce_sum.242: f32[]) -> f32[] { + %reduce_sum.241 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.242 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.243 = f32[] add(%reduce_sum.241, %reduce_sum.242), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/vmap()/reduce_sum" stack_frame_id=34} +} + +%region_48.53 (reduce_sum.248: f32[], reduce_sum.249: f32[]) -> f32[] { + %reduce_sum.248 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.249 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.250 = f32[] add(%reduce_sum.248, %reduce_sum.249), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/feedforward_layernorm/reduce_sum" stack_frame_id=43} +} + +%region_49.54 (reduce_sum.255: f32[], reduce_sum.256: f32[]) -> f32[] { + %reduce_sum.255 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.256 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.257 = f32[] add(%reduce_sum.255, %reduce_sum.256), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention_layernorm/reduce_sum" stack_frame_id=24} +} + +%region_50.55 (reduce_max.87: f32[], reduce_max.88: f32[]) -> f32[] { + %reduce_max.87 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.88 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.89 = f32[] maximum(%reduce_max.87, %reduce_max.88), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/vmap()/reduce_max" stack_frame_id=34} +} + +%region_51.56 (reduce_sum.262: f32[], reduce_sum.263: f32[]) -> f32[] { + %reduce_sum.262 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.263 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.264 = f32[] add(%reduce_sum.262, %reduce_sum.263), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/vmap()/reduce_sum" stack_frame_id=34} +} + +%region_52.57 (reduce_sum.269: f32[], reduce_sum.270: f32[]) -> f32[] { + %reduce_sum.269 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.270 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.271 = f32[] add(%reduce_sum.269, %reduce_sum.270), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/feedforward_layernorm/reduce_sum" stack_frame_id=43} +} + +%region_53.58 (reduce_sum.276: f32[], reduce_sum.277: f32[]) -> f32[] { + %reduce_sum.276 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.277 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.278 = f32[] add(%reduce_sum.276, %reduce_sum.277), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention_layernorm/reduce_sum" stack_frame_id=24} +} + +%region_54.59 (reduce_max.94: f32[], reduce_max.95: f32[]) -> f32[] { + %reduce_max.94 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.95 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.96 = f32[] maximum(%reduce_max.94, %reduce_max.95), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/vmap()/reduce_max" stack_frame_id=34} +} + +%region_55.60 (reduce_sum.283: f32[], reduce_sum.284: f32[]) -> f32[] { + %reduce_sum.283 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.284 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.285 = f32[] add(%reduce_sum.283, %reduce_sum.284), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/vmap()/reduce_sum" stack_frame_id=34} +} + +%region_56.61 (reduce_sum.290: f32[], reduce_sum.291: f32[]) -> f32[] { + %reduce_sum.290 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.291 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.292 = f32[] add(%reduce_sum.290, %reduce_sum.291), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/feedforward_layernorm/reduce_sum" stack_frame_id=43} +} + +%region_57.62 (reduce_sum.297: f32[], reduce_sum.298: f32[]) -> f32[] { + %reduce_sum.297 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.298 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.299 = f32[] add(%reduce_sum.297, %reduce_sum.298), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention_layernorm/reduce_sum" stack_frame_id=24} +} + +%region_58.63 (reduce_max.101: f32[], reduce_max.102: f32[]) -> f32[] { + %reduce_max.101 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.102 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.103 = f32[] maximum(%reduce_max.101, %reduce_max.102), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/vmap()/reduce_max" stack_frame_id=34} +} + +%region_59.64 (reduce_sum.304: f32[], reduce_sum.305: f32[]) -> f32[] { + %reduce_sum.304 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.305 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.306 = f32[] add(%reduce_sum.304, %reduce_sum.305), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/vmap()/reduce_sum" stack_frame_id=34} +} + +%region_60.65 (reduce_sum.311: f32[], reduce_sum.312: f32[]) -> f32[] { + %reduce_sum.311 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.312 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.313 = f32[] add(%reduce_sum.311, %reduce_sum.312), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/feedforward_layernorm/reduce_sum" stack_frame_id=43} +} + +%region_61.66 (reduce_sum.318: f32[], reduce_sum.319: f32[]) -> f32[] { + %reduce_sum.318 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.319 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.320 = f32[] add(%reduce_sum.318, %reduce_sum.319), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention_layernorm/reduce_sum" stack_frame_id=24} +} + +%region_62.67 (reduce_max.108: f32[], reduce_max.109: f32[]) -> f32[] { + %reduce_max.108 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.109 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.110 = f32[] maximum(%reduce_max.108, %reduce_max.109), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/vmap()/reduce_max" stack_frame_id=34} +} + +%region_63.68 (reduce_sum.325: f32[], reduce_sum.326: f32[]) -> f32[] { + %reduce_sum.325 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.326 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.327 = f32[] add(%reduce_sum.325, %reduce_sum.326), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/vmap()/reduce_sum" stack_frame_id=34} +} + +%region_64.69 (reduce_sum.332: f32[], reduce_sum.333: f32[]) -> f32[] { + %reduce_sum.332 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.333 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.334 = f32[] add(%reduce_sum.332, %reduce_sum.333), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/feedforward_layernorm/reduce_sum" stack_frame_id=43} +} + +%region_65.70 (reduce_sum.339: f32[], reduce_sum.340: f32[]) -> f32[] { + %reduce_sum.339 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.340 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.341 = f32[] add(%reduce_sum.339, %reduce_sum.340), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention_layernorm/reduce_sum" stack_frame_id=24} +} + +%region_66.71 (reduce_max.115: f32[], reduce_max.116: f32[]) -> f32[] { + %reduce_max.115 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.116 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.117 = f32[] maximum(%reduce_max.115, %reduce_max.116), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/vmap()/reduce_max" stack_frame_id=34} +} + +%region_67.72 (reduce_sum.346: f32[], reduce_sum.347: f32[]) -> f32[] { + %reduce_sum.346 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.347 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.348 = f32[] add(%reduce_sum.346, %reduce_sum.347), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/vmap()/reduce_sum" stack_frame_id=34} +} + +%region_68.73 (reduce_sum.353: f32[], reduce_sum.354: f32[]) -> f32[] { + %reduce_sum.353 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.354 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.355 = f32[] add(%reduce_sum.353, %reduce_sum.354), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/feedforward_layernorm/reduce_sum" stack_frame_id=43} +} + +%region_69.74 (reduce_sum.360: f32[], reduce_sum.361: f32[]) -> f32[] { + %reduce_sum.360 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.361 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.362 = f32[] add(%reduce_sum.360, %reduce_sum.361), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention_layernorm/reduce_sum" stack_frame_id=24} +} + +%region_70.75 (reduce_max.122: f32[], reduce_max.123: f32[]) -> f32[] { + %reduce_max.122 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.123 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.124 = f32[] maximum(%reduce_max.122, %reduce_max.123), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/vmap()/reduce_max" stack_frame_id=34} +} + +%region_71.76 (reduce_sum.367: f32[], reduce_sum.368: f32[]) -> f32[] { + %reduce_sum.367 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.368 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.369 = f32[] add(%reduce_sum.367, %reduce_sum.368), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/vmap()/reduce_sum" stack_frame_id=34} +} + +%region_72.77 (reduce_sum.374: f32[], reduce_sum.375: f32[]) -> f32[] { + %reduce_sum.374 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.375 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.376 = f32[] add(%reduce_sum.374, %reduce_sum.375), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/feedforward_layernorm/reduce_sum" stack_frame_id=43} +} + +%region_73.78 (reduce_sum.381: f32[], reduce_sum.382: f32[]) -> f32[] { + %reduce_sum.381 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.382 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.383 = f32[] add(%reduce_sum.381, %reduce_sum.382), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention_layernorm/reduce_sum" stack_frame_id=24} +} + +%region_74.79 (reduce_max.129: f32[], reduce_max.130: f32[]) -> f32[] { + %reduce_max.129 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.130 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.131 = f32[] maximum(%reduce_max.129, %reduce_max.130), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/vmap()/reduce_max" stack_frame_id=34} +} + +%region_75.80 (reduce_sum.388: f32[], reduce_sum.389: f32[]) -> f32[] { + %reduce_sum.388 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.389 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.390 = f32[] add(%reduce_sum.388, %reduce_sum.389), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/vmap()/reduce_sum" stack_frame_id=34} +} + +%region_76.81 (reduce_sum.395: f32[], reduce_sum.396: f32[]) -> f32[] { + %reduce_sum.395 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.396 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.397 = f32[] add(%reduce_sum.395, %reduce_sum.396), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/feedforward_layernorm/reduce_sum" stack_frame_id=43} +} + +%region_77.82 (reduce_sum.402: f32[], reduce_sum.403: f32[]) -> f32[] { + %reduce_sum.402 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.403 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.404 = f32[] add(%reduce_sum.402, %reduce_sum.403), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention_layernorm/reduce_sum" stack_frame_id=24} +} + +%region_78.83 (reduce_max.136: f32[], reduce_max.137: f32[]) -> f32[] { + %reduce_max.136 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.137 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.138 = f32[] maximum(%reduce_max.136, %reduce_max.137), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/vmap()/reduce_max" stack_frame_id=34} +} + +%region_79.84 (reduce_sum.409: f32[], reduce_sum.410: f32[]) -> f32[] { + %reduce_sum.409 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.410 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.411 = f32[] add(%reduce_sum.409, %reduce_sum.410), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/vmap()/reduce_sum" stack_frame_id=34} +} + +%region_80.85 (reduce_sum.416: f32[], reduce_sum.417: f32[]) -> f32[] { + %reduce_sum.416 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.417 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.418 = f32[] add(%reduce_sum.416, %reduce_sum.417), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/feedforward_layernorm/reduce_sum" stack_frame_id=43} +} + +%region_81.86 (reduce_sum.423: f32[], reduce_sum.424: f32[]) -> f32[] { + %reduce_sum.423 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.424 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.425 = f32[] add(%reduce_sum.423, %reduce_sum.424), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention_layernorm/reduce_sum" stack_frame_id=24} +} + +%region_82.87 (reduce_max.143: f32[], reduce_max.144: f32[]) -> f32[] { + %reduce_max.143 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.144 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.145 = f32[] maximum(%reduce_max.143, %reduce_max.144), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/vmap()/reduce_max" stack_frame_id=34} +} + +%region_83.88 (reduce_sum.430: f32[], reduce_sum.431: f32[]) -> f32[] { + %reduce_sum.430 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.431 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.432 = f32[] add(%reduce_sum.430, %reduce_sum.431), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/vmap()/reduce_sum" stack_frame_id=34} +} + +%region_84.89 (reduce_sum.437: f32[], reduce_sum.438: f32[]) -> f32[] { + %reduce_sum.437 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.438 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.439 = f32[] add(%reduce_sum.437, %reduce_sum.438), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/feedforward_layernorm/reduce_sum" stack_frame_id=43} +} + +%region_85.90 (reduce_sum.444: f32[], reduce_sum.445: f32[]) -> f32[] { + %reduce_sum.444 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.445 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.446 = f32[] add(%reduce_sum.444, %reduce_sum.445), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention_layernorm/reduce_sum" stack_frame_id=24} +} + +%region_86.91 (reduce_max.150: f32[], reduce_max.151: f32[]) -> f32[] { + %reduce_max.150 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.151 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.152 = f32[] maximum(%reduce_max.150, %reduce_max.151), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/vmap()/reduce_max" stack_frame_id=34} +} + +%region_87.92 (reduce_sum.451: f32[], reduce_sum.452: f32[]) -> f32[] { + %reduce_sum.451 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.452 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.453 = f32[] add(%reduce_sum.451, %reduce_sum.452), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/vmap()/reduce_sum" stack_frame_id=34} +} + +%region_88.93 (reduce_sum.458: f32[], reduce_sum.459: f32[]) -> f32[] { + %reduce_sum.458 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.459 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.460 = f32[] add(%reduce_sum.458, %reduce_sum.459), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/feedforward_layernorm/reduce_sum" stack_frame_id=43} +} + +%region_89.94 (reduce_sum.465: f32[], reduce_sum.466: f32[]) -> f32[] { + %reduce_sum.465 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.466 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.467 = f32[] add(%reduce_sum.465, %reduce_sum.466), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention_layernorm/reduce_sum" stack_frame_id=24} +} + +%region_90.95 (reduce_max.157: f32[], reduce_max.158: f32[]) -> f32[] { + %reduce_max.157 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.158 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.159 = f32[] maximum(%reduce_max.157, %reduce_max.158), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/vmap()/reduce_max" stack_frame_id=34} +} + +%region_91.96 (reduce_sum.472: f32[], reduce_sum.473: f32[]) -> f32[] { + %reduce_sum.472 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.473 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.474 = f32[] add(%reduce_sum.472, %reduce_sum.473), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/vmap()/reduce_sum" stack_frame_id=34} +} + +%region_92.97 (reduce_sum.479: f32[], reduce_sum.480: f32[]) -> f32[] { + %reduce_sum.479 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.480 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.481 = f32[] add(%reduce_sum.479, %reduce_sum.480), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/feedforward_layernorm/reduce_sum" stack_frame_id=43} +} + +%region_93.98 (reduce_sum.486: f32[], reduce_sum.487: f32[]) -> f32[] { + %reduce_sum.486 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.487 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.488 = f32[] add(%reduce_sum.486, %reduce_sum.487), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention_layernorm/reduce_sum" stack_frame_id=24} +} + +%region_94.99 (reduce_max.164: f32[], reduce_max.165: f32[]) -> f32[] { + %reduce_max.164 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.165 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.166 = f32[] maximum(%reduce_max.164, %reduce_max.165), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/vmap()/reduce_max" stack_frame_id=34} +} + +%region_95.100 (reduce_sum.493: f32[], reduce_sum.494: f32[]) -> f32[] { + %reduce_sum.493 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.494 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.495 = f32[] add(%reduce_sum.493, %reduce_sum.494), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/vmap()/reduce_sum" stack_frame_id=34} +} + +%region_96.101 (reduce_sum.500: f32[], reduce_sum.501: f32[]) -> f32[] { + %reduce_sum.500 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.501 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.502 = f32[] add(%reduce_sum.500, %reduce_sum.501), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/feedforward_layernorm/reduce_sum" stack_frame_id=43} +} + +%region_97.102 (reduce_sum.507: f32[], reduce_sum.508: f32[]) -> f32[] { + %reduce_sum.507 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.508 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.509 = f32[] add(%reduce_sum.507, %reduce_sum.508), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention_layernorm/reduce_sum" stack_frame_id=24} +} + +%region_98.103 (reduce_max.171: f32[], reduce_max.172: f32[]) -> f32[] { + %reduce_max.171 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.172 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.173 = f32[] maximum(%reduce_max.171, %reduce_max.172), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/vmap()/reduce_max" stack_frame_id=34} +} + +%region_99.104 (reduce_sum.514: f32[], reduce_sum.515: f32[]) -> f32[] { + %reduce_sum.514 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.515 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.516 = f32[] add(%reduce_sum.514, %reduce_sum.515), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/vmap()/reduce_sum" stack_frame_id=34} +} + +%region_100.105 (reduce_sum.521: f32[], reduce_sum.522: f32[]) -> f32[] { + %reduce_sum.521 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.522 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.523 = f32[] add(%reduce_sum.521, %reduce_sum.522), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/feedforward_layernorm/reduce_sum" stack_frame_id=43} +} + +%region_101.106 (reduce_sum.528: f32[], reduce_sum.529: f32[]) -> f32[] { + %reduce_sum.528 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.529 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.530 = f32[] add(%reduce_sum.528, %reduce_sum.529), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention_layernorm/reduce_sum" stack_frame_id=24} +} + +%region_102.107 (reduce_max.178: f32[], reduce_max.179: f32[]) -> f32[] { + %reduce_max.178 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.179 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.180 = f32[] maximum(%reduce_max.178, %reduce_max.179), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/vmap()/reduce_max" stack_frame_id=34} +} + +%region_103.108 (reduce_sum.535: f32[], reduce_sum.536: f32[]) -> f32[] { + %reduce_sum.535 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.536 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.537 = f32[] add(%reduce_sum.535, %reduce_sum.536), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/vmap()/reduce_sum" stack_frame_id=34} +} + +%region_104.109 (reduce_sum.542: f32[], reduce_sum.543: f32[]) -> f32[] { + %reduce_sum.542 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.543 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.544 = f32[] add(%reduce_sum.542, %reduce_sum.543), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/feedforward_layernorm/reduce_sum" stack_frame_id=43} +} + +%region_105.110 (reduce_sum.549: f32[], reduce_sum.550: f32[]) -> f32[] { + %reduce_sum.549 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.550 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.551 = f32[] add(%reduce_sum.549, %reduce_sum.550), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention_layernorm/reduce_sum" stack_frame_id=24} +} + +%region_106.111 (reduce_max.185: f32[], reduce_max.186: f32[]) -> f32[] { + %reduce_max.185 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.186 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.187 = f32[] maximum(%reduce_max.185, %reduce_max.186), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/vmap()/reduce_max" stack_frame_id=34} +} + +%region_107.112 (reduce_sum.556: f32[], reduce_sum.557: f32[]) -> f32[] { + %reduce_sum.556 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.557 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.558 = f32[] add(%reduce_sum.556, %reduce_sum.557), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/vmap()/reduce_sum" stack_frame_id=34} +} + +%region_108.113 (reduce_sum.563: f32[], reduce_sum.564: f32[]) -> f32[] { + %reduce_sum.563 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.564 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.565 = f32[] add(%reduce_sum.563, %reduce_sum.564), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/feedforward_layernorm/reduce_sum" stack_frame_id=43} +} + +%region_109.114 (reduce_sum.570: f32[], reduce_sum.571: f32[]) -> f32[] { + %reduce_sum.570 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.571 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.572 = f32[] add(%reduce_sum.570, %reduce_sum.571), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention_layernorm/reduce_sum" stack_frame_id=24} +} + +%region_110.115 (reduce_max.192: f32[], reduce_max.193: f32[]) -> f32[] { + %reduce_max.192 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.193 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.194 = f32[] maximum(%reduce_max.192, %reduce_max.193), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/vmap()/reduce_max" stack_frame_id=34} +} + +%region_111.116 (reduce_sum.577: f32[], reduce_sum.578: f32[]) -> f32[] { + %reduce_sum.577 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.578 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.579 = f32[] add(%reduce_sum.577, %reduce_sum.578), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/vmap()/reduce_sum" stack_frame_id=34} +} + +%region_112.117 (reduce_sum.584: f32[], reduce_sum.585: f32[]) -> f32[] { + %reduce_sum.584 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.585 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.586 = f32[] add(%reduce_sum.584, %reduce_sum.585), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/feedforward_layernorm/reduce_sum" stack_frame_id=43} +} + +%region_113.118 (reduce_sum.591: f32[], reduce_sum.592: f32[]) -> f32[] { + %reduce_sum.591 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.592 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.593 = f32[] add(%reduce_sum.591, %reduce_sum.592), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention_layernorm/reduce_sum" stack_frame_id=24} +} + +%region_114.119 (reduce_max.199: f32[], reduce_max.200: f32[]) -> f32[] { + %reduce_max.199 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.200 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.201 = f32[] maximum(%reduce_max.199, %reduce_max.200), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/vmap()/reduce_max" stack_frame_id=34} +} + +%region_115.120 (reduce_sum.598: f32[], reduce_sum.599: f32[]) -> f32[] { + %reduce_sum.598 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.599 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.600 = f32[] add(%reduce_sum.598, %reduce_sum.599), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/vmap()/reduce_sum" stack_frame_id=34} +} + +%region_116.121 (reduce_sum.605: f32[], reduce_sum.606: f32[]) -> f32[] { + %reduce_sum.605 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.606 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.607 = f32[] add(%reduce_sum.605, %reduce_sum.606), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/feedforward_layernorm/reduce_sum" stack_frame_id=43} +} + +%region_117.122 (reduce_sum.612: f32[], reduce_sum.613: f32[]) -> f32[] { + %reduce_sum.612 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.613 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.614 = f32[] add(%reduce_sum.612, %reduce_sum.613), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention_layernorm/reduce_sum" stack_frame_id=24} +} + +%region_118.123 (reduce_max.206: f32[], reduce_max.207: f32[]) -> f32[] { + %reduce_max.206 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.207 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.208 = f32[] maximum(%reduce_max.206, %reduce_max.207), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/vmap()/reduce_max" stack_frame_id=34} +} + +%region_119.124 (reduce_sum.619: f32[], reduce_sum.620: f32[]) -> f32[] { + %reduce_sum.619 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.620 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.621 = f32[] add(%reduce_sum.619, %reduce_sum.620), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/vmap()/reduce_sum" stack_frame_id=34} +} + +%region_120.125 (reduce_sum.626: f32[], reduce_sum.627: f32[]) -> f32[] { + %reduce_sum.626 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.627 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.628 = f32[] add(%reduce_sum.626, %reduce_sum.627), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/feedforward_layernorm/reduce_sum" stack_frame_id=43} +} + +%region_121.126 (reduce_sum.633: f32[], reduce_sum.634: f32[]) -> f32[] { + %reduce_sum.633 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.634 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.635 = f32[] add(%reduce_sum.633, %reduce_sum.634), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention_layernorm/reduce_sum" stack_frame_id=24} +} + +%region_122.127 (reduce_max.213: f32[], reduce_max.214: f32[]) -> f32[] { + %reduce_max.213 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.214 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.215 = f32[] maximum(%reduce_max.213, %reduce_max.214), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/vmap()/reduce_max" stack_frame_id=34} +} + +%region_123.128 (reduce_sum.640: f32[], reduce_sum.641: f32[]) -> f32[] { + %reduce_sum.640 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.641 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.642 = f32[] add(%reduce_sum.640, %reduce_sum.641), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/vmap()/reduce_sum" stack_frame_id=34} +} + +%region_124.129 (reduce_sum.647: f32[], reduce_sum.648: f32[]) -> f32[] { + %reduce_sum.647 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.648 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.649 = f32[] add(%reduce_sum.647, %reduce_sum.648), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/feedforward_layernorm/reduce_sum" stack_frame_id=43} +} + +%region_125.130 (reduce_sum.654: f32[], reduce_sum.655: f32[]) -> f32[] { + %reduce_sum.654 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.655 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.656 = f32[] add(%reduce_sum.654, %reduce_sum.655), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention_layernorm/reduce_sum" stack_frame_id=24} +} + +%region_126.131 (reduce_sum.661: s32[], reduce_sum.662: s32[]) -> s32[] { + %reduce_sum.661 = s32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.662 = s32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.663 = s32[] add(%reduce_sum.661, %reduce_sum.662), metadata={op_name="jit(compiled_generate_function)/reduce_sum" stack_frame_id=46} +} + +%region_127.132 (reduce_min.3: s32[], reduce_min.4: s32[]) -> s32[] { + %reduce_min.3 = s32[] parameter(0), metadata={op_name="reduce_min"} + %reduce_min.4 = s32[] parameter(1), metadata={op_name="reduce_min"} + ROOT %reduce_min.5 = s32[] minimum(%reduce_min.3, %reduce_min.4), metadata={op_name="jit(compiled_generate_function)/reduce_min" stack_frame_id=49} +} + +%_where_2.133 (Arg_0.10: pred[1,1], Arg_1.8: s32[1,1], Arg_2.5: s32[1,1]) -> s32[1,1] { + %Arg_0.10 = pred[1,1]{1,0} parameter(0) + %Arg_1.8 = s32[1,1]{1,0} parameter(1) + %Arg_2.5 = s32[1,1]{1,0} parameter(2) + ROOT %select_n.8 = s32[1,1]{1,0} select(%Arg_0.10, %Arg_1.8, %Arg_2.5), metadata={op_name="select_n"} +} + +%region_129.134 (reduce_and.11: pred[], reduce_and.12: pred[]) -> pred[] { + %reduce_and.11 = pred[] parameter(0), metadata={op_name="reduce_and"} + %reduce_and.12 = pred[] parameter(1), metadata={op_name="reduce_and"} + ROOT %reduce_and.13 = pred[] and(%reduce_and.11, %reduce_and.12), metadata={op_name="reduce_and"} +} + +%_take_1.135 (Arg_0.11: bf16[32000,4096], Arg_1.9: s32[1,1]) -> bf16[1,1,4096] { + %Arg_1.9 = s32[1,1]{1,0} parameter(1) + %constant.64 = s32[1,1]{1,0} constant({ {0} }) + %lt.223 = pred[1,1]{1,0} compare(%Arg_1.9, %constant.64), direction=LT, metadata={op_name="lt"} + %constant.62 = s32[1,1]{1,0} constant({ {32000} }) + %add.227 = s32[1,1]{1,0} add(%Arg_1.9, %constant.62), metadata={op_name="add"} + %jit__where_.3 = s32[1,1]{1,0} call(%lt.223, %add.227, %Arg_1.9), to_apply=%_where_2.133, metadata={op_name="jit(_where)"} + %broadcast_in_dim.573 = s32[1,1,1]{2,1,0} reshape(%jit__where_.3), metadata={op_name="broadcast_in_dim"} + %constant.61 = s32[1,1,1]{2,1,0} constant({ { {0} } }) + %ge.222 = pred[1,1,1]{2,1,0} compare(%broadcast_in_dim.573, %constant.61), direction=GE, metadata={op_name="ge"} + %constant.60 = s32[1,1,1]{2,1,0} constant({ { {31999} } }) + %le.5 = pred[1,1,1]{2,1,0} compare(%broadcast_in_dim.573, %constant.60), direction=LE, metadata={op_name="le"} + %and.96 = pred[1,1,1]{2,1,0} and(%ge.222, %le.5), metadata={op_name="and"} + %constant.63 = pred[] constant(true) + %reduce_and.15 = pred[1,1]{1,0} reduce(%and.96, %constant.63), dimensions={2}, to_apply=%region_129.134, metadata={op_name="reduce_and"} + %broadcast_in_dim.574 = pred[1,1,4096]{2,1,0} broadcast(%reduce_and.15), dimensions={0,1}, metadata={op_name="broadcast_in_dim"} + %Arg_0.11 = bf16[32000,4096]{1,0} parameter(0) + %gather.3 = bf16[1,1,4096]{2,1,0} gather(%Arg_0.11, %broadcast_in_dim.573), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,4096}, metadata={op_name="gather"} + %constant.59 = bf16[] constant(nan) + %broadcast_in_dim.572 = bf16[1,1,4096]{2,1,0} broadcast(%constant.59), dimensions={}, metadata={op_name="broadcast_in_dim"} + ROOT %select_n.10 = bf16[1,1,4096]{2,1,0} select(%broadcast_in_dim.574, %gather.3, %broadcast_in_dim.572), metadata={op_name="select_n"} +} + +%region_130.136 (reduce_sum.668: f32[], reduce_sum.669: f32[]) -> f32[] { + %reduce_sum.668 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.669 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.670 = f32[] add(%reduce_sum.668, %reduce_sum.669), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention_layernorm/reduce_sum" stack_frame_id=73} +} + +%_where_3.137 (Arg_0.13: pred[1,1,32,1,41], Arg_1.11: f32[1,32,1,1,41], Arg_2.7: f32[]) -> f32[1,1,32,1,41] { + %Arg_0.13 = pred[1,1,32,1,41]{4,3,2,1,0} parameter(0) + %Arg_1.11 = f32[1,32,1,1,41]{4,3,2,1,0} parameter(1) + %transpose.34 = f32[1,1,32,1,41]{4,3,2,1,0} reshape(%Arg_1.11), metadata={op_name="transpose"} + %Arg_2.7 = f32[] parameter(2) + %broadcast_in_dim.586 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%Arg_2.7), dimensions={}, metadata={op_name="broadcast_in_dim"} + ROOT %select_n.14 = f32[1,1,32,1,41]{4,3,2,1,0} select(%Arg_0.13, %transpose.34, %broadcast_in_dim.586), metadata={op_name="select_n"} +} + +%region_131.138 (reduce_max.220: f32[], reduce_max.221: f32[]) -> f32[] { + %reduce_max.220 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.221 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.222 = f32[] maximum(%reduce_max.220, %reduce_max.221), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/vmap()/reduce_max" stack_frame_id=83} +} + +%region_132.139 (reduce_sum.675: f32[], reduce_sum.676: f32[]) -> f32[] { + %reduce_sum.675 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.676 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.677 = f32[] add(%reduce_sum.675, %reduce_sum.676), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/vmap()/reduce_sum" stack_frame_id=83} +} + +%region_133.140 (reduce_sum.682: f32[], reduce_sum.683: f32[]) -> f32[] { + %reduce_sum.682 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.683 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.684 = f32[] add(%reduce_sum.682, %reduce_sum.683), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/feedforward_layernorm/reduce_sum" stack_frame_id=92} +} + +%silu_4.141 (Arg_0.15: f32[1,1,14336]) -> f32[1,1,14336] { + %Arg_0.15 = f32[1,1,14336]{2,1,0} parameter(0) + %constant.66 = f32[] constant(1) + %broadcast.25 = f32[1,1,14336]{2,1,0} broadcast(%constant.66), dimensions={} + %neg.133 = f32[1,1,14336]{2,1,0} negate(%Arg_0.15), metadata={op_name="neg"} + %exp.35 = f32[1,1,14336]{2,1,0} exponential(%neg.133), metadata={op_name="exp"} + %add.241 = f32[1,1,14336]{2,1,0} add(%exp.35, %broadcast.25), metadata={op_name="add"} + %div.326 = f32[1,1,14336]{2,1,0} divide(%broadcast.25, %add.241), metadata={op_name="div"} + ROOT %mul.1165 = f32[1,1,14336]{2,1,0} multiply(%Arg_0.15, %div.326), metadata={op_name="mul"} +} + +%region_134.142 (reduce_sum.689: f32[], reduce_sum.690: f32[]) -> f32[] { + %reduce_sum.689 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.690 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.691 = f32[] add(%reduce_sum.689, %reduce_sum.690), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention_layernorm/reduce_sum" stack_frame_id=73} +} + +%region_135.143 (reduce_max.227: f32[], reduce_max.228: f32[]) -> f32[] { + %reduce_max.227 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.228 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.229 = f32[] maximum(%reduce_max.227, %reduce_max.228), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/vmap()/reduce_max" stack_frame_id=83} +} + +%region_136.144 (reduce_sum.696: f32[], reduce_sum.697: f32[]) -> f32[] { + %reduce_sum.696 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.697 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.698 = f32[] add(%reduce_sum.696, %reduce_sum.697), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/vmap()/reduce_sum" stack_frame_id=83} +} + +%region_137.145 (reduce_sum.703: f32[], reduce_sum.704: f32[]) -> f32[] { + %reduce_sum.703 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.704 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.705 = f32[] add(%reduce_sum.703, %reduce_sum.704), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/feedforward_layernorm/reduce_sum" stack_frame_id=92} +} + +%region_138.146 (reduce_sum.710: f32[], reduce_sum.711: f32[]) -> f32[] { + %reduce_sum.710 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.711 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.712 = f32[] add(%reduce_sum.710, %reduce_sum.711), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention_layernorm/reduce_sum" stack_frame_id=73} +} + +%region_139.147 (reduce_max.234: f32[], reduce_max.235: f32[]) -> f32[] { + %reduce_max.234 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.235 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.236 = f32[] maximum(%reduce_max.234, %reduce_max.235), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/vmap()/reduce_max" stack_frame_id=83} +} + +%region_140.148 (reduce_sum.717: f32[], reduce_sum.718: f32[]) -> f32[] { + %reduce_sum.717 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.718 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.719 = f32[] add(%reduce_sum.717, %reduce_sum.718), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/vmap()/reduce_sum" stack_frame_id=83} +} + +%region_141.149 (reduce_sum.724: f32[], reduce_sum.725: f32[]) -> f32[] { + %reduce_sum.724 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.725 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.726 = f32[] add(%reduce_sum.724, %reduce_sum.725), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/feedforward_layernorm/reduce_sum" stack_frame_id=92} +} + +%region_142.150 (reduce_sum.731: f32[], reduce_sum.732: f32[]) -> f32[] { + %reduce_sum.731 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.732 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.733 = f32[] add(%reduce_sum.731, %reduce_sum.732), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention_layernorm/reduce_sum" stack_frame_id=73} +} + +%region_143.151 (reduce_max.241: f32[], reduce_max.242: f32[]) -> f32[] { + %reduce_max.241 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.242 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.243 = f32[] maximum(%reduce_max.241, %reduce_max.242), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/vmap()/reduce_max" stack_frame_id=83} +} + +%region_144.152 (reduce_sum.738: f32[], reduce_sum.739: f32[]) -> f32[] { + %reduce_sum.738 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.739 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.740 = f32[] add(%reduce_sum.738, %reduce_sum.739), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/vmap()/reduce_sum" stack_frame_id=83} +} + +%region_145.153 (reduce_sum.745: f32[], reduce_sum.746: f32[]) -> f32[] { + %reduce_sum.745 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.746 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.747 = f32[] add(%reduce_sum.745, %reduce_sum.746), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/feedforward_layernorm/reduce_sum" stack_frame_id=92} +} + +%region_146.154 (reduce_sum.752: f32[], reduce_sum.753: f32[]) -> f32[] { + %reduce_sum.752 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.753 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.754 = f32[] add(%reduce_sum.752, %reduce_sum.753), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention_layernorm/reduce_sum" stack_frame_id=73} +} + +%region_147.155 (reduce_max.248: f32[], reduce_max.249: f32[]) -> f32[] { + %reduce_max.248 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.249 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.250 = f32[] maximum(%reduce_max.248, %reduce_max.249), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/vmap()/reduce_max" stack_frame_id=83} +} + +%region_148.156 (reduce_sum.759: f32[], reduce_sum.760: f32[]) -> f32[] { + %reduce_sum.759 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.760 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.761 = f32[] add(%reduce_sum.759, %reduce_sum.760), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/vmap()/reduce_sum" stack_frame_id=83} +} + +%region_149.157 (reduce_sum.766: f32[], reduce_sum.767: f32[]) -> f32[] { + %reduce_sum.766 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.767 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.768 = f32[] add(%reduce_sum.766, %reduce_sum.767), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/feedforward_layernorm/reduce_sum" stack_frame_id=92} +} + +%region_150.158 (reduce_sum.773: f32[], reduce_sum.774: f32[]) -> f32[] { + %reduce_sum.773 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.774 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.775 = f32[] add(%reduce_sum.773, %reduce_sum.774), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention_layernorm/reduce_sum" stack_frame_id=73} +} + +%region_151.159 (reduce_max.255: f32[], reduce_max.256: f32[]) -> f32[] { + %reduce_max.255 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.256 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.257 = f32[] maximum(%reduce_max.255, %reduce_max.256), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/vmap()/reduce_max" stack_frame_id=83} +} + +%region_152.160 (reduce_sum.780: f32[], reduce_sum.781: f32[]) -> f32[] { + %reduce_sum.780 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.781 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.782 = f32[] add(%reduce_sum.780, %reduce_sum.781), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/vmap()/reduce_sum" stack_frame_id=83} +} + +%region_153.161 (reduce_sum.787: f32[], reduce_sum.788: f32[]) -> f32[] { + %reduce_sum.787 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.788 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.789 = f32[] add(%reduce_sum.787, %reduce_sum.788), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/feedforward_layernorm/reduce_sum" stack_frame_id=92} +} + +%region_154.162 (reduce_sum.794: f32[], reduce_sum.795: f32[]) -> f32[] { + %reduce_sum.794 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.795 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.796 = f32[] add(%reduce_sum.794, %reduce_sum.795), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention_layernorm/reduce_sum" stack_frame_id=73} +} + +%region_155.163 (reduce_max.262: f32[], reduce_max.263: f32[]) -> f32[] { + %reduce_max.262 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.263 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.264 = f32[] maximum(%reduce_max.262, %reduce_max.263), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/vmap()/reduce_max" stack_frame_id=83} +} + +%region_156.164 (reduce_sum.801: f32[], reduce_sum.802: f32[]) -> f32[] { + %reduce_sum.801 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.802 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.803 = f32[] add(%reduce_sum.801, %reduce_sum.802), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/vmap()/reduce_sum" stack_frame_id=83} +} + +%region_157.165 (reduce_sum.808: f32[], reduce_sum.809: f32[]) -> f32[] { + %reduce_sum.808 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.809 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.810 = f32[] add(%reduce_sum.808, %reduce_sum.809), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/feedforward_layernorm/reduce_sum" stack_frame_id=92} +} + +%region_158.166 (reduce_sum.815: f32[], reduce_sum.816: f32[]) -> f32[] { + %reduce_sum.815 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.816 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.817 = f32[] add(%reduce_sum.815, %reduce_sum.816), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention_layernorm/reduce_sum" stack_frame_id=73} +} + +%region_159.167 (reduce_max.269: f32[], reduce_max.270: f32[]) -> f32[] { + %reduce_max.269 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.270 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.271 = f32[] maximum(%reduce_max.269, %reduce_max.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/vmap()/reduce_max" stack_frame_id=83} +} + +%region_160.168 (reduce_sum.822: f32[], reduce_sum.823: f32[]) -> f32[] { + %reduce_sum.822 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.823 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.824 = f32[] add(%reduce_sum.822, %reduce_sum.823), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/vmap()/reduce_sum" stack_frame_id=83} +} + +%region_161.169 (reduce_sum.829: f32[], reduce_sum.830: f32[]) -> f32[] { + %reduce_sum.829 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.830 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.831 = f32[] add(%reduce_sum.829, %reduce_sum.830), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/feedforward_layernorm/reduce_sum" stack_frame_id=92} +} + +%region_162.170 (reduce_sum.836: f32[], reduce_sum.837: f32[]) -> f32[] { + %reduce_sum.836 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.837 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.838 = f32[] add(%reduce_sum.836, %reduce_sum.837), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention_layernorm/reduce_sum" stack_frame_id=73} +} + +%region_163.171 (reduce_max.276: f32[], reduce_max.277: f32[]) -> f32[] { + %reduce_max.276 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.277 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.278 = f32[] maximum(%reduce_max.276, %reduce_max.277), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/vmap()/reduce_max" stack_frame_id=83} +} + +%region_164.172 (reduce_sum.843: f32[], reduce_sum.844: f32[]) -> f32[] { + %reduce_sum.843 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.844 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.845 = f32[] add(%reduce_sum.843, %reduce_sum.844), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/vmap()/reduce_sum" stack_frame_id=83} +} + +%region_165.173 (reduce_sum.850: f32[], reduce_sum.851: f32[]) -> f32[] { + %reduce_sum.850 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.851 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.852 = f32[] add(%reduce_sum.850, %reduce_sum.851), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/feedforward_layernorm/reduce_sum" stack_frame_id=92} +} + +%region_166.174 (reduce_sum.857: f32[], reduce_sum.858: f32[]) -> f32[] { + %reduce_sum.857 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.858 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.859 = f32[] add(%reduce_sum.857, %reduce_sum.858), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention_layernorm/reduce_sum" stack_frame_id=73} +} + +%region_167.175 (reduce_max.283: f32[], reduce_max.284: f32[]) -> f32[] { + %reduce_max.283 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.284 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.285 = f32[] maximum(%reduce_max.283, %reduce_max.284), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/vmap()/reduce_max" stack_frame_id=83} +} + +%region_168.176 (reduce_sum.864: f32[], reduce_sum.865: f32[]) -> f32[] { + %reduce_sum.864 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.865 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.866 = f32[] add(%reduce_sum.864, %reduce_sum.865), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/vmap()/reduce_sum" stack_frame_id=83} +} + +%region_169.177 (reduce_sum.871: f32[], reduce_sum.872: f32[]) -> f32[] { + %reduce_sum.871 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.872 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.873 = f32[] add(%reduce_sum.871, %reduce_sum.872), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/feedforward_layernorm/reduce_sum" stack_frame_id=92} +} + +%region_170.178 (reduce_sum.878: f32[], reduce_sum.879: f32[]) -> f32[] { + %reduce_sum.878 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.879 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.880 = f32[] add(%reduce_sum.878, %reduce_sum.879), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention_layernorm/reduce_sum" stack_frame_id=73} +} + +%region_171.179 (reduce_max.290: f32[], reduce_max.291: f32[]) -> f32[] { + %reduce_max.290 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.291 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.292 = f32[] maximum(%reduce_max.290, %reduce_max.291), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/vmap()/reduce_max" stack_frame_id=83} +} + +%region_172.180 (reduce_sum.885: f32[], reduce_sum.886: f32[]) -> f32[] { + %reduce_sum.885 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.886 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.887 = f32[] add(%reduce_sum.885, %reduce_sum.886), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/vmap()/reduce_sum" stack_frame_id=83} +} + +%region_173.181 (reduce_sum.892: f32[], reduce_sum.893: f32[]) -> f32[] { + %reduce_sum.892 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.893 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.894 = f32[] add(%reduce_sum.892, %reduce_sum.893), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/feedforward_layernorm/reduce_sum" stack_frame_id=92} +} + +%region_174.182 (reduce_sum.899: f32[], reduce_sum.900: f32[]) -> f32[] { + %reduce_sum.899 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.900 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.901 = f32[] add(%reduce_sum.899, %reduce_sum.900), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention_layernorm/reduce_sum" stack_frame_id=73} +} + +%region_175.183 (reduce_max.297: f32[], reduce_max.298: f32[]) -> f32[] { + %reduce_max.297 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.298 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.299 = f32[] maximum(%reduce_max.297, %reduce_max.298), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/vmap()/reduce_max" stack_frame_id=83} +} + +%region_176.184 (reduce_sum.906: f32[], reduce_sum.907: f32[]) -> f32[] { + %reduce_sum.906 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.907 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.908 = f32[] add(%reduce_sum.906, %reduce_sum.907), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/vmap()/reduce_sum" stack_frame_id=83} +} + +%region_177.185 (reduce_sum.913: f32[], reduce_sum.914: f32[]) -> f32[] { + %reduce_sum.913 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.914 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.915 = f32[] add(%reduce_sum.913, %reduce_sum.914), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/feedforward_layernorm/reduce_sum" stack_frame_id=92} +} + +%region_178.186 (reduce_sum.920: f32[], reduce_sum.921: f32[]) -> f32[] { + %reduce_sum.920 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.921 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.922 = f32[] add(%reduce_sum.920, %reduce_sum.921), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention_layernorm/reduce_sum" stack_frame_id=73} +} + +%region_179.187 (reduce_max.304: f32[], reduce_max.305: f32[]) -> f32[] { + %reduce_max.304 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.305 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.306 = f32[] maximum(%reduce_max.304, %reduce_max.305), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/vmap()/reduce_max" stack_frame_id=83} +} + +%region_180.188 (reduce_sum.927: f32[], reduce_sum.928: f32[]) -> f32[] { + %reduce_sum.927 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.928 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.929 = f32[] add(%reduce_sum.927, %reduce_sum.928), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/vmap()/reduce_sum" stack_frame_id=83} +} + +%region_181.189 (reduce_sum.934: f32[], reduce_sum.935: f32[]) -> f32[] { + %reduce_sum.934 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.935 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.936 = f32[] add(%reduce_sum.934, %reduce_sum.935), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/feedforward_layernorm/reduce_sum" stack_frame_id=92} +} + +%region_182.190 (reduce_sum.941: f32[], reduce_sum.942: f32[]) -> f32[] { + %reduce_sum.941 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.942 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.943 = f32[] add(%reduce_sum.941, %reduce_sum.942), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention_layernorm/reduce_sum" stack_frame_id=73} +} + +%region_183.191 (reduce_max.311: f32[], reduce_max.312: f32[]) -> f32[] { + %reduce_max.311 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.312 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.313 = f32[] maximum(%reduce_max.311, %reduce_max.312), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/vmap()/reduce_max" stack_frame_id=83} +} + +%region_184.192 (reduce_sum.948: f32[], reduce_sum.949: f32[]) -> f32[] { + %reduce_sum.948 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.949 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.950 = f32[] add(%reduce_sum.948, %reduce_sum.949), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/vmap()/reduce_sum" stack_frame_id=83} +} + +%region_185.193 (reduce_sum.955: f32[], reduce_sum.956: f32[]) -> f32[] { + %reduce_sum.955 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.956 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.957 = f32[] add(%reduce_sum.955, %reduce_sum.956), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/feedforward_layernorm/reduce_sum" stack_frame_id=92} +} + +%region_186.194 (reduce_sum.962: f32[], reduce_sum.963: f32[]) -> f32[] { + %reduce_sum.962 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.963 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.964 = f32[] add(%reduce_sum.962, %reduce_sum.963), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention_layernorm/reduce_sum" stack_frame_id=73} +} + +%region_187.195 (reduce_max.318: f32[], reduce_max.319: f32[]) -> f32[] { + %reduce_max.318 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.319 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.320 = f32[] maximum(%reduce_max.318, %reduce_max.319), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/vmap()/reduce_max" stack_frame_id=83} +} + +%region_188.196 (reduce_sum.969: f32[], reduce_sum.970: f32[]) -> f32[] { + %reduce_sum.969 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.970 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.971 = f32[] add(%reduce_sum.969, %reduce_sum.970), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/vmap()/reduce_sum" stack_frame_id=83} +} + +%region_189.197 (reduce_sum.976: f32[], reduce_sum.977: f32[]) -> f32[] { + %reduce_sum.976 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.977 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.978 = f32[] add(%reduce_sum.976, %reduce_sum.977), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/feedforward_layernorm/reduce_sum" stack_frame_id=92} +} + +%region_190.198 (reduce_sum.983: f32[], reduce_sum.984: f32[]) -> f32[] { + %reduce_sum.983 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.984 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.985 = f32[] add(%reduce_sum.983, %reduce_sum.984), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention_layernorm/reduce_sum" stack_frame_id=73} +} + +%region_191.199 (reduce_max.325: f32[], reduce_max.326: f32[]) -> f32[] { + %reduce_max.325 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.326 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.327 = f32[] maximum(%reduce_max.325, %reduce_max.326), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/vmap()/reduce_max" stack_frame_id=83} +} + +%region_192.200 (reduce_sum.990: f32[], reduce_sum.991: f32[]) -> f32[] { + %reduce_sum.990 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.991 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.992 = f32[] add(%reduce_sum.990, %reduce_sum.991), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/vmap()/reduce_sum" stack_frame_id=83} +} + +%region_193.201 (reduce_sum.997: f32[], reduce_sum.998: f32[]) -> f32[] { + %reduce_sum.997 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.998 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.999 = f32[] add(%reduce_sum.997, %reduce_sum.998), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/feedforward_layernorm/reduce_sum" stack_frame_id=92} +} + +%region_194.202 (reduce_sum.1004: f32[], reduce_sum.1005: f32[]) -> f32[] { + %reduce_sum.1004 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1005 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1006 = f32[] add(%reduce_sum.1004, %reduce_sum.1005), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention_layernorm/reduce_sum" stack_frame_id=73} +} + +%region_195.203 (reduce_max.332: f32[], reduce_max.333: f32[]) -> f32[] { + %reduce_max.332 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.333 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.334 = f32[] maximum(%reduce_max.332, %reduce_max.333), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/vmap()/reduce_max" stack_frame_id=83} +} + +%region_196.204 (reduce_sum.1011: f32[], reduce_sum.1012: f32[]) -> f32[] { + %reduce_sum.1011 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1012 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1013 = f32[] add(%reduce_sum.1011, %reduce_sum.1012), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/vmap()/reduce_sum" stack_frame_id=83} +} + +%region_197.205 (reduce_sum.1018: f32[], reduce_sum.1019: f32[]) -> f32[] { + %reduce_sum.1018 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1019 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1020 = f32[] add(%reduce_sum.1018, %reduce_sum.1019), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/feedforward_layernorm/reduce_sum" stack_frame_id=92} +} + +%region_198.206 (reduce_sum.1025: f32[], reduce_sum.1026: f32[]) -> f32[] { + %reduce_sum.1025 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1026 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1027 = f32[] add(%reduce_sum.1025, %reduce_sum.1026), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention_layernorm/reduce_sum" stack_frame_id=73} +} + +%region_199.207 (reduce_max.339: f32[], reduce_max.340: f32[]) -> f32[] { + %reduce_max.339 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.340 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.341 = f32[] maximum(%reduce_max.339, %reduce_max.340), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/vmap()/reduce_max" stack_frame_id=83} +} + +%region_200.208 (reduce_sum.1032: f32[], reduce_sum.1033: f32[]) -> f32[] { + %reduce_sum.1032 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1033 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1034 = f32[] add(%reduce_sum.1032, %reduce_sum.1033), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/vmap()/reduce_sum" stack_frame_id=83} +} + +%region_201.209 (reduce_sum.1039: f32[], reduce_sum.1040: f32[]) -> f32[] { + %reduce_sum.1039 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1040 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1041 = f32[] add(%reduce_sum.1039, %reduce_sum.1040), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/feedforward_layernorm/reduce_sum" stack_frame_id=92} +} + +%region_202.210 (reduce_sum.1046: f32[], reduce_sum.1047: f32[]) -> f32[] { + %reduce_sum.1046 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1047 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1048 = f32[] add(%reduce_sum.1046, %reduce_sum.1047), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention_layernorm/reduce_sum" stack_frame_id=73} +} + +%region_203.211 (reduce_max.346: f32[], reduce_max.347: f32[]) -> f32[] { + %reduce_max.346 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.347 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.348 = f32[] maximum(%reduce_max.346, %reduce_max.347), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/vmap()/reduce_max" stack_frame_id=83} +} + +%region_204.212 (reduce_sum.1053: f32[], reduce_sum.1054: f32[]) -> f32[] { + %reduce_sum.1053 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1054 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1055 = f32[] add(%reduce_sum.1053, %reduce_sum.1054), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/vmap()/reduce_sum" stack_frame_id=83} +} + +%region_205.213 (reduce_sum.1060: f32[], reduce_sum.1061: f32[]) -> f32[] { + %reduce_sum.1060 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1061 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1062 = f32[] add(%reduce_sum.1060, %reduce_sum.1061), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/feedforward_layernorm/reduce_sum" stack_frame_id=92} +} + +%region_206.214 (reduce_sum.1067: f32[], reduce_sum.1068: f32[]) -> f32[] { + %reduce_sum.1067 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1068 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1069 = f32[] add(%reduce_sum.1067, %reduce_sum.1068), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention_layernorm/reduce_sum" stack_frame_id=73} +} + +%region_207.215 (reduce_max.353: f32[], reduce_max.354: f32[]) -> f32[] { + %reduce_max.353 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.354 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.355 = f32[] maximum(%reduce_max.353, %reduce_max.354), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/vmap()/reduce_max" stack_frame_id=83} +} + +%region_208.216 (reduce_sum.1074: f32[], reduce_sum.1075: f32[]) -> f32[] { + %reduce_sum.1074 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1075 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1076 = f32[] add(%reduce_sum.1074, %reduce_sum.1075), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/vmap()/reduce_sum" stack_frame_id=83} +} + +%region_209.217 (reduce_sum.1081: f32[], reduce_sum.1082: f32[]) -> f32[] { + %reduce_sum.1081 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1082 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1083 = f32[] add(%reduce_sum.1081, %reduce_sum.1082), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/feedforward_layernorm/reduce_sum" stack_frame_id=92} +} + +%region_210.218 (reduce_sum.1088: f32[], reduce_sum.1089: f32[]) -> f32[] { + %reduce_sum.1088 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1089 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1090 = f32[] add(%reduce_sum.1088, %reduce_sum.1089), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention_layernorm/reduce_sum" stack_frame_id=73} +} + +%region_211.219 (reduce_max.360: f32[], reduce_max.361: f32[]) -> f32[] { + %reduce_max.360 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.361 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.362 = f32[] maximum(%reduce_max.360, %reduce_max.361), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/vmap()/reduce_max" stack_frame_id=83} +} + +%region_212.220 (reduce_sum.1095: f32[], reduce_sum.1096: f32[]) -> f32[] { + %reduce_sum.1095 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1096 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1097 = f32[] add(%reduce_sum.1095, %reduce_sum.1096), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/vmap()/reduce_sum" stack_frame_id=83} +} + +%region_213.221 (reduce_sum.1102: f32[], reduce_sum.1103: f32[]) -> f32[] { + %reduce_sum.1102 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1103 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1104 = f32[] add(%reduce_sum.1102, %reduce_sum.1103), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/feedforward_layernorm/reduce_sum" stack_frame_id=92} +} + +%region_214.222 (reduce_sum.1109: f32[], reduce_sum.1110: f32[]) -> f32[] { + %reduce_sum.1109 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1110 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1111 = f32[] add(%reduce_sum.1109, %reduce_sum.1110), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention_layernorm/reduce_sum" stack_frame_id=73} +} + +%region_215.223 (reduce_max.367: f32[], reduce_max.368: f32[]) -> f32[] { + %reduce_max.367 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.368 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.369 = f32[] maximum(%reduce_max.367, %reduce_max.368), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/vmap()/reduce_max" stack_frame_id=83} +} + +%region_216.224 (reduce_sum.1116: f32[], reduce_sum.1117: f32[]) -> f32[] { + %reduce_sum.1116 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1117 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1118 = f32[] add(%reduce_sum.1116, %reduce_sum.1117), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/vmap()/reduce_sum" stack_frame_id=83} +} + +%region_217.225 (reduce_sum.1123: f32[], reduce_sum.1124: f32[]) -> f32[] { + %reduce_sum.1123 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1124 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1125 = f32[] add(%reduce_sum.1123, %reduce_sum.1124), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/feedforward_layernorm/reduce_sum" stack_frame_id=92} +} + +%region_218.226 (reduce_sum.1130: f32[], reduce_sum.1131: f32[]) -> f32[] { + %reduce_sum.1130 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1131 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1132 = f32[] add(%reduce_sum.1130, %reduce_sum.1131), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention_layernorm/reduce_sum" stack_frame_id=73} +} + +%region_219.227 (reduce_max.374: f32[], reduce_max.375: f32[]) -> f32[] { + %reduce_max.374 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.375 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.376 = f32[] maximum(%reduce_max.374, %reduce_max.375), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/vmap()/reduce_max" stack_frame_id=83} +} + +%region_220.228 (reduce_sum.1137: f32[], reduce_sum.1138: f32[]) -> f32[] { + %reduce_sum.1137 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1138 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1139 = f32[] add(%reduce_sum.1137, %reduce_sum.1138), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/vmap()/reduce_sum" stack_frame_id=83} +} + +%region_221.229 (reduce_sum.1144: f32[], reduce_sum.1145: f32[]) -> f32[] { + %reduce_sum.1144 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1145 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1146 = f32[] add(%reduce_sum.1144, %reduce_sum.1145), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/feedforward_layernorm/reduce_sum" stack_frame_id=92} +} + +%region_222.230 (reduce_sum.1151: f32[], reduce_sum.1152: f32[]) -> f32[] { + %reduce_sum.1151 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1152 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1153 = f32[] add(%reduce_sum.1151, %reduce_sum.1152), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention_layernorm/reduce_sum" stack_frame_id=73} +} + +%region_223.231 (reduce_max.381: f32[], reduce_max.382: f32[]) -> f32[] { + %reduce_max.381 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.382 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.383 = f32[] maximum(%reduce_max.381, %reduce_max.382), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/vmap()/reduce_max" stack_frame_id=83} +} + +%region_224.232 (reduce_sum.1158: f32[], reduce_sum.1159: f32[]) -> f32[] { + %reduce_sum.1158 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1159 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1160 = f32[] add(%reduce_sum.1158, %reduce_sum.1159), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/vmap()/reduce_sum" stack_frame_id=83} +} + +%region_225.233 (reduce_sum.1165: f32[], reduce_sum.1166: f32[]) -> f32[] { + %reduce_sum.1165 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1166 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1167 = f32[] add(%reduce_sum.1165, %reduce_sum.1166), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/feedforward_layernorm/reduce_sum" stack_frame_id=92} +} + +%region_226.234 (reduce_sum.1172: f32[], reduce_sum.1173: f32[]) -> f32[] { + %reduce_sum.1172 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1173 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1174 = f32[] add(%reduce_sum.1172, %reduce_sum.1173), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention_layernorm/reduce_sum" stack_frame_id=73} +} + +%region_227.235 (reduce_max.388: f32[], reduce_max.389: f32[]) -> f32[] { + %reduce_max.388 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.389 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.390 = f32[] maximum(%reduce_max.388, %reduce_max.389), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/vmap()/reduce_max" stack_frame_id=83} +} + +%region_228.236 (reduce_sum.1179: f32[], reduce_sum.1180: f32[]) -> f32[] { + %reduce_sum.1179 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1180 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1181 = f32[] add(%reduce_sum.1179, %reduce_sum.1180), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/vmap()/reduce_sum" stack_frame_id=83} +} + +%region_229.237 (reduce_sum.1186: f32[], reduce_sum.1187: f32[]) -> f32[] { + %reduce_sum.1186 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1187 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1188 = f32[] add(%reduce_sum.1186, %reduce_sum.1187), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/feedforward_layernorm/reduce_sum" stack_frame_id=92} +} + +%region_230.238 (reduce_sum.1193: f32[], reduce_sum.1194: f32[]) -> f32[] { + %reduce_sum.1193 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1194 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1195 = f32[] add(%reduce_sum.1193, %reduce_sum.1194), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention_layernorm/reduce_sum" stack_frame_id=73} +} + +%region_231.239 (reduce_max.395: f32[], reduce_max.396: f32[]) -> f32[] { + %reduce_max.395 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.396 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.397 = f32[] maximum(%reduce_max.395, %reduce_max.396), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/vmap()/reduce_max" stack_frame_id=83} +} + +%region_232.240 (reduce_sum.1200: f32[], reduce_sum.1201: f32[]) -> f32[] { + %reduce_sum.1200 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1201 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1202 = f32[] add(%reduce_sum.1200, %reduce_sum.1201), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/vmap()/reduce_sum" stack_frame_id=83} +} + +%region_233.241 (reduce_sum.1207: f32[], reduce_sum.1208: f32[]) -> f32[] { + %reduce_sum.1207 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1208 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1209 = f32[] add(%reduce_sum.1207, %reduce_sum.1208), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/feedforward_layernorm/reduce_sum" stack_frame_id=92} +} + +%region_234.242 (reduce_sum.1214: f32[], reduce_sum.1215: f32[]) -> f32[] { + %reduce_sum.1214 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1215 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1216 = f32[] add(%reduce_sum.1214, %reduce_sum.1215), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention_layernorm/reduce_sum" stack_frame_id=73} +} + +%region_235.243 (reduce_max.402: f32[], reduce_max.403: f32[]) -> f32[] { + %reduce_max.402 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.403 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.404 = f32[] maximum(%reduce_max.402, %reduce_max.403), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/vmap()/reduce_max" stack_frame_id=83} +} + +%region_236.244 (reduce_sum.1221: f32[], reduce_sum.1222: f32[]) -> f32[] { + %reduce_sum.1221 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1222 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1223 = f32[] add(%reduce_sum.1221, %reduce_sum.1222), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/vmap()/reduce_sum" stack_frame_id=83} +} + +%region_237.245 (reduce_sum.1228: f32[], reduce_sum.1229: f32[]) -> f32[] { + %reduce_sum.1228 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1229 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1230 = f32[] add(%reduce_sum.1228, %reduce_sum.1229), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/feedforward_layernorm/reduce_sum" stack_frame_id=92} +} + +%region_238.246 (reduce_sum.1235: f32[], reduce_sum.1236: f32[]) -> f32[] { + %reduce_sum.1235 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1236 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1237 = f32[] add(%reduce_sum.1235, %reduce_sum.1236), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention_layernorm/reduce_sum" stack_frame_id=73} +} + +%region_239.247 (reduce_max.409: f32[], reduce_max.410: f32[]) -> f32[] { + %reduce_max.409 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.410 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.411 = f32[] maximum(%reduce_max.409, %reduce_max.410), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/vmap()/reduce_max" stack_frame_id=83} +} + +%region_240.248 (reduce_sum.1242: f32[], reduce_sum.1243: f32[]) -> f32[] { + %reduce_sum.1242 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1243 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1244 = f32[] add(%reduce_sum.1242, %reduce_sum.1243), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/vmap()/reduce_sum" stack_frame_id=83} +} + +%region_241.249 (reduce_sum.1249: f32[], reduce_sum.1250: f32[]) -> f32[] { + %reduce_sum.1249 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1250 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1251 = f32[] add(%reduce_sum.1249, %reduce_sum.1250), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/feedforward_layernorm/reduce_sum" stack_frame_id=92} +} + +%region_242.250 (reduce_sum.1256: f32[], reduce_sum.1257: f32[]) -> f32[] { + %reduce_sum.1256 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1257 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1258 = f32[] add(%reduce_sum.1256, %reduce_sum.1257), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention_layernorm/reduce_sum" stack_frame_id=73} +} + +%region_243.251 (reduce_max.416: f32[], reduce_max.417: f32[]) -> f32[] { + %reduce_max.416 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.417 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.418 = f32[] maximum(%reduce_max.416, %reduce_max.417), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/vmap()/reduce_max" stack_frame_id=83} +} + +%region_244.252 (reduce_sum.1263: f32[], reduce_sum.1264: f32[]) -> f32[] { + %reduce_sum.1263 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1264 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1265 = f32[] add(%reduce_sum.1263, %reduce_sum.1264), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/vmap()/reduce_sum" stack_frame_id=83} +} + +%region_245.253 (reduce_sum.1270: f32[], reduce_sum.1271: f32[]) -> f32[] { + %reduce_sum.1270 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1271 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1272 = f32[] add(%reduce_sum.1270, %reduce_sum.1271), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/feedforward_layernorm/reduce_sum" stack_frame_id=92} +} + +%region_246.254 (reduce_sum.1277: f32[], reduce_sum.1278: f32[]) -> f32[] { + %reduce_sum.1277 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1278 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1279 = f32[] add(%reduce_sum.1277, %reduce_sum.1278), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention_layernorm/reduce_sum" stack_frame_id=73} +} + +%region_247.255 (reduce_max.423: f32[], reduce_max.424: f32[]) -> f32[] { + %reduce_max.423 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.424 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.425 = f32[] maximum(%reduce_max.423, %reduce_max.424), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/vmap()/reduce_max" stack_frame_id=83} +} + +%region_248.256 (reduce_sum.1284: f32[], reduce_sum.1285: f32[]) -> f32[] { + %reduce_sum.1284 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1285 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1286 = f32[] add(%reduce_sum.1284, %reduce_sum.1285), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/vmap()/reduce_sum" stack_frame_id=83} +} + +%region_249.257 (reduce_sum.1291: f32[], reduce_sum.1292: f32[]) -> f32[] { + %reduce_sum.1291 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1292 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1293 = f32[] add(%reduce_sum.1291, %reduce_sum.1292), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/feedforward_layernorm/reduce_sum" stack_frame_id=92} +} + +%region_250.258 (reduce_sum.1298: f32[], reduce_sum.1299: f32[]) -> f32[] { + %reduce_sum.1298 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1299 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1300 = f32[] add(%reduce_sum.1298, %reduce_sum.1299), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention_layernorm/reduce_sum" stack_frame_id=73} +} + +%region_251.259 (reduce_max.430: f32[], reduce_max.431: f32[]) -> f32[] { + %reduce_max.430 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.431 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.432 = f32[] maximum(%reduce_max.430, %reduce_max.431), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/vmap()/reduce_max" stack_frame_id=83} +} + +%region_252.260 (reduce_sum.1305: f32[], reduce_sum.1306: f32[]) -> f32[] { + %reduce_sum.1305 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1306 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1307 = f32[] add(%reduce_sum.1305, %reduce_sum.1306), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/vmap()/reduce_sum" stack_frame_id=83} +} + +%region_253.261 (reduce_sum.1312: f32[], reduce_sum.1313: f32[]) -> f32[] { + %reduce_sum.1312 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1313 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1314 = f32[] add(%reduce_sum.1312, %reduce_sum.1313), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/feedforward_layernorm/reduce_sum" stack_frame_id=92} +} + +%region_254.262 (reduce_sum.1319: f32[], reduce_sum.1320: f32[]) -> f32[] { + %reduce_sum.1319 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1320 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1321 = f32[] add(%reduce_sum.1319, %reduce_sum.1320), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention_layernorm/reduce_sum" stack_frame_id=73} +} + +%region_255.263 (reduce_max.437: f32[], reduce_max.438: f32[]) -> f32[] { + %reduce_max.437 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.438 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.439 = f32[] maximum(%reduce_max.437, %reduce_max.438), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/vmap()/reduce_max" stack_frame_id=83} +} + +%region_256.264 (reduce_sum.1326: f32[], reduce_sum.1327: f32[]) -> f32[] { + %reduce_sum.1326 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1327 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1328 = f32[] add(%reduce_sum.1326, %reduce_sum.1327), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/vmap()/reduce_sum" stack_frame_id=83} +} + +%region_257.265 (reduce_sum.1333: f32[], reduce_sum.1334: f32[]) -> f32[] { + %reduce_sum.1333 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1334 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1335 = f32[] add(%reduce_sum.1333, %reduce_sum.1334), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/feedforward_layernorm/reduce_sum" stack_frame_id=92} +} + +%region_258.266 (reduce_sum.1340: f32[], reduce_sum.1341: f32[]) -> f32[] { + %reduce_sum.1340 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1341 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1342 = f32[] add(%reduce_sum.1340, %reduce_sum.1341), metadata={op_name="jit(compiled_generate_function)/while/body/sequence_output_layernorm/reduce_sum" stack_frame_id=101} +} + +%region_259.267 (reduce_max.444: f32[], reduce_max.445: f32[]) -> f32[] { + %reduce_max.444 = f32[] parameter(0), metadata={op_name="reduce_max"} + %reduce_max.445 = f32[] parameter(1), metadata={op_name="reduce_max"} + ROOT %reduce_max.446 = f32[] maximum(%reduce_max.444, %reduce_max.445), metadata={op_name="jit(compiled_generate_function)/while/body/reduce_max" stack_frame_id=106} +} + +%region_260.268 (reduce_sum.1347: f32[], reduce_sum.1348: f32[]) -> f32[] { + %reduce_sum.1347 = f32[] parameter(0), metadata={op_name="reduce_sum"} + %reduce_sum.1348 = f32[] parameter(1), metadata={op_name="reduce_sum"} + ROOT %reduce_sum.1349 = f32[] add(%reduce_sum.1347, %reduce_sum.1348), metadata={op_name="jit(compiled_generate_function)/while/body/reduce_sum" stack_frame_id=106} +} + +%closed_call.269 (Arg_0.20: s32[], Arg_1.15: u32[1,5,1], Arg_2.11: u32[1,5,1], Arg_3.2: u32[], Arg_4.1: u32[], Arg_5.1: u32[], Arg_6.1: u32[4], Arg_7.1: u32[4]) -> (s32[], u32[1,5,1], u32[1,5,1], u32[], u32[], /*index=5*/u32[], u32[4], u32[4]) { + %Arg_0.20 = s32[] parameter(0) + %constant.81 = s32[] constant(1) + %add.664 = s32[] add(%Arg_0.20, %constant.81), metadata={op_name="add"} + %Arg_1.15 = u32[1,5,1]{2,1,0} parameter(1) + %Arg_2.11 = u32[1,5,1]{2,1,0} parameter(2) + %add.665 = u32[1,5,1]{2,1,0} add(%Arg_1.15, %Arg_2.11), metadata={op_name="add"} + %Arg_6.1 = u32[4]{0} parameter(6) + %unstack.12 = u32[1]{0} slice(%Arg_6.1), slice={[0:1]}, metadata={op_name="unstack"} + %unstack.13 = u32[] reshape(%unstack.12), metadata={op_name="unstack"} + %shift_left.8 = u32[1,5,1]{2,1,0} broadcast(%unstack.13), dimensions={}, metadata={op_name="shift_left"} + %shift_left.9 = u32[1,5,1]{2,1,0} shift-left(%Arg_2.11, %shift_left.8), metadata={op_name="shift_left"} + %constant.80 = u32[] constant(32) + %sub.262 = u32[] subtract(%constant.80, %unstack.13), metadata={op_name="sub"} + %shift_right_logical.8 = u32[1,5,1]{2,1,0} broadcast(%sub.262), dimensions={}, metadata={op_name="shift_right_logical"} + %shift_right_logical.9 = u32[1,5,1]{2,1,0} shift-right-logical(%Arg_2.11, %shift_right_logical.8), metadata={op_name="shift_right_logical"} + %or.4 = u32[1,5,1]{2,1,0} or(%shift_left.9, %shift_right_logical.9), metadata={op_name="or"} + %xor.6 = u32[1,5,1]{2,1,0} xor(%add.665, %or.4), metadata={op_name="xor"} + %add.666 = u32[1,5,1]{2,1,0} add(%add.665, %xor.6), metadata={op_name="add"} + %unstack.14 = u32[1]{0} slice(%Arg_6.1), slice={[1:2]}, metadata={op_name="unstack"} + %unstack.15 = u32[] reshape(%unstack.14), metadata={op_name="unstack"} + %shift_left.10 = u32[1,5,1]{2,1,0} broadcast(%unstack.15), dimensions={}, metadata={op_name="shift_left"} + %shift_left.11 = u32[1,5,1]{2,1,0} shift-left(%xor.6, %shift_left.10), metadata={op_name="shift_left"} + %sub.263 = u32[] subtract(%constant.80, %unstack.15), metadata={op_name="sub"} + %shift_right_logical.10 = u32[1,5,1]{2,1,0} broadcast(%sub.263), dimensions={}, metadata={op_name="shift_right_logical"} + %shift_right_logical.11 = u32[1,5,1]{2,1,0} shift-right-logical(%xor.6, %shift_right_logical.10), metadata={op_name="shift_right_logical"} + %or.5 = u32[1,5,1]{2,1,0} or(%shift_left.11, %shift_right_logical.11), metadata={op_name="or"} + %xor.7 = u32[1,5,1]{2,1,0} xor(%add.666, %or.5), metadata={op_name="xor"} + %add.667 = u32[1,5,1]{2,1,0} add(%add.666, %xor.7), metadata={op_name="add"} + %unstack.16 = u32[1]{0} slice(%Arg_6.1), slice={[2:3]}, metadata={op_name="unstack"} + %unstack.17 = u32[] reshape(%unstack.16), metadata={op_name="unstack"} + %shift_left.12 = u32[1,5,1]{2,1,0} broadcast(%unstack.17), dimensions={}, metadata={op_name="shift_left"} + %shift_left.13 = u32[1,5,1]{2,1,0} shift-left(%xor.7, %shift_left.12), metadata={op_name="shift_left"} + %sub.264 = u32[] subtract(%constant.80, %unstack.17), metadata={op_name="sub"} + %shift_right_logical.12 = u32[1,5,1]{2,1,0} broadcast(%sub.264), dimensions={}, metadata={op_name="shift_right_logical"} + %shift_right_logical.13 = u32[1,5,1]{2,1,0} shift-right-logical(%xor.7, %shift_right_logical.12), metadata={op_name="shift_right_logical"} + %or.6 = u32[1,5,1]{2,1,0} or(%shift_left.13, %shift_right_logical.13), metadata={op_name="or"} + %xor.8 = u32[1,5,1]{2,1,0} xor(%add.667, %or.6), metadata={op_name="xor"} + %add.668 = u32[1,5,1]{2,1,0} add(%add.667, %xor.8), metadata={op_name="add"} + %Arg_3.2 = u32[] parameter(3) + %add.669 = u32[1,5,1]{2,1,0} broadcast(%Arg_3.2), dimensions={}, metadata={op_name="add"} + %add.670 = u32[1,5,1]{2,1,0} add(%add.668, %add.669), metadata={op_name="add"} + %unstack.18 = u32[1]{0} slice(%Arg_6.1), slice={[3:4]}, metadata={op_name="unstack"} + %unstack.19 = u32[] reshape(%unstack.18), metadata={op_name="unstack"} + %shift_left.14 = u32[1,5,1]{2,1,0} broadcast(%unstack.19), dimensions={}, metadata={op_name="shift_left"} + %shift_left.15 = u32[1,5,1]{2,1,0} shift-left(%xor.8, %shift_left.14), metadata={op_name="shift_left"} + %sub.265 = u32[] subtract(%constant.80, %unstack.19), metadata={op_name="sub"} + %shift_right_logical.14 = u32[1,5,1]{2,1,0} broadcast(%sub.265), dimensions={}, metadata={op_name="shift_right_logical"} + %shift_right_logical.15 = u32[1,5,1]{2,1,0} shift-right-logical(%xor.8, %shift_right_logical.14), metadata={op_name="shift_right_logical"} + %or.7 = u32[1,5,1]{2,1,0} or(%shift_left.15, %shift_right_logical.15), metadata={op_name="or"} + %xor.9 = u32[1,5,1]{2,1,0} xor(%add.668, %or.7), metadata={op_name="xor"} + %Arg_4.1 = u32[] parameter(4) + %add.671 = u32[1,5,1]{2,1,0} broadcast(%Arg_4.1), dimensions={}, metadata={op_name="add"} + %add.672 = u32[1,5,1]{2,1,0} add(%xor.9, %add.671), metadata={op_name="add"} + %add.673 = s32[] add(%Arg_0.20, %constant.81), metadata={op_name="add"} + %convert_element_type.1116 = u32[] convert(%add.673), metadata={op_name="convert_element_type"} + %add.674 = u32[1,5,1]{2,1,0} broadcast(%convert_element_type.1116), dimensions={}, metadata={op_name="add"} + %add.675 = u32[1,5,1]{2,1,0} add(%add.672, %add.674), metadata={op_name="add"} + %Arg_5.1 = u32[] parameter(5) + %Arg_7.1 = u32[4]{0} parameter(7) + ROOT %tuple.1 = (s32[], u32[1,5,1]{2,1,0}, u32[1,5,1]{2,1,0}, u32[], u32[], /*index=5*/u32[], u32[4]{0}, u32[4]{0}) tuple(%add.664, %add.670, %add.675, %Arg_4.1, %Arg_5.1, /*index=5*/%Arg_3.2, %Arg_7.1, %Arg_6.1) +} + +%region_261.270 (arg_tuple.2: (s32[], s32[], u32[1,5,1], u32[1,5,1], u32[], /*index=5*/u32[], u32[], u32[4], u32[4])) -> (s32[], s32[], u32[1,5,1], u32[1,5,1], u32[], /*index=5*/u32[], u32[], u32[4], u32[4]) { + %arg_tuple.2 = (s32[], s32[], u32[1,5,1]{2,1,0}, u32[1,5,1]{2,1,0}, u32[], /*index=5*/u32[], u32[], u32[4]{0}, u32[4]{0}) parameter(0) + %constant.82 = s32[] constant(1) + %add.677 = s32[] add(%arg_tuple.2#0, %constant.82), metadata={op_name="while/body/add"} + %closed_call.9 = (s32[], u32[1,5,1]{2,1,0}, u32[1,5,1]{2,1,0}, u32[], u32[], /*index=5*/u32[], u32[4]{0}, u32[4]{0}) call(%arg_tuple.2#1, %arg_tuple.2#2, %arg_tuple.2#3, %arg_tuple.2#4, %arg_tuple.2#5, /*index=5*/%arg_tuple.2#6, %arg_tuple.2#7, %arg_tuple.2#8), to_apply=%closed_call.269, metadata={op_name="while/body/closed_call"} + ROOT %tuple.3 = (s32[], s32[], u32[1,5,1]{2,1,0}, u32[1,5,1]{2,1,0}, u32[], /*index=5*/u32[], u32[], u32[4]{0}, u32[4]{0}) tuple(%add.677, %closed_call.9#0, %closed_call.9#1, %closed_call.9#2, %closed_call.9#3, /*index=5*/%closed_call.9#4, %closed_call.9#5, %closed_call.9#6, %closed_call.9#7) +} + +%region_262.271 (arg_tuple.4: (s32[], s32[], u32[1,5,1], u32[1,5,1], u32[], /*index=5*/u32[], u32[], u32[4], u32[4])) -> pred[] { + %arg_tuple.4 = (s32[], s32[], u32[1,5,1]{2,1,0}, u32[1,5,1]{2,1,0}, u32[], /*index=5*/u32[], u32[], u32[4]{0}, u32[4]{0}) parameter(0) + %constant.84 = s32[] constant(5) + ROOT %lt.417 = pred[] compare(%arg_tuple.4#0, %constant.84), direction=LT, metadata={op_name="while/cond/lt"} +} + +%threefry2x32.272 (Arg_0.21: u32[], Arg_1.16: u32[], Arg_2.12: u32[1,5,1], Arg_3.3: u32[1,5,1]) -> (u32[1,5,1], u32[1,5,1]) { + %constant.85 = s32[] constant(0) + %Arg_2.12 = u32[1,5,1]{2,1,0} parameter(2) + %Arg_0.21 = u32[] parameter(0) + %add.678 = u32[1,5,1]{2,1,0} broadcast(%Arg_0.21), dimensions={}, metadata={op_name="add"} + %add.679 = u32[1,5,1]{2,1,0} add(%Arg_2.12, %add.678), metadata={op_name="add"} + %Arg_3.3 = u32[1,5,1]{2,1,0} parameter(3) + %Arg_1.16 = u32[] parameter(1) + %add.680 = u32[1,5,1]{2,1,0} broadcast(%Arg_1.16), dimensions={}, metadata={op_name="add"} + %add.681 = u32[1,5,1]{2,1,0} add(%Arg_3.3, %add.680), metadata={op_name="add"} + %xor.10 = u32[] xor(%Arg_0.21, %Arg_1.16), metadata={op_name="xor"} + %constant.86 = u32[] constant(466688986) + %xor.11 = u32[] xor(%xor.10, %constant.86), metadata={op_name="xor"} + %constant.87 = u32[4]{0} constant({13, 15, 26, 6}) + %constant.88 = u32[4]{0} constant({17, 29, 16, 24}) + %while.11 = (s32[], s32[], u32[1,5,1]{2,1,0}, u32[1,5,1]{2,1,0}, u32[], /*index=5*/u32[], u32[], u32[4]{0}, u32[4]{0}) tuple(%constant.85, %constant.85, %add.679, %add.681, %Arg_1.16, /*index=5*/%xor.11, %Arg_0.21, %constant.87, %constant.88), metadata={op_name="while"} + %while.12 = (s32[], s32[], u32[1,5,1]{2,1,0}, u32[1,5,1]{2,1,0}, u32[], /*index=5*/u32[], u32[], u32[4]{0}, u32[4]{0}) while(%while.11), condition=%region_262.271, body=%region_261.270, metadata={op_name="while"} + ROOT %tuple.5 = (u32[1,5,1]{2,1,0}, u32[1,5,1]{2,1,0}) tuple(%while.12#2, %while.12#3) +} + +%_uniform.273 (Arg_0.22: u32[2], Arg_1.17: f32[], Arg_2.13: f32[]) -> f32[1,5,1] { + %Arg_2.13 = f32[] parameter(2) + %broadcast_in_dim.1032 = f32[1,1,1]{2,1,0} reshape(%Arg_2.13), metadata={op_name="broadcast_in_dim"} + %max.68 = f32[1,1,1]{2,1,0} broadcast(%broadcast_in_dim.1032), dimensions={0,1,2}, metadata={op_name="max"} + %max.69 = f32[1,1]{1,0} reshape(%max.68), metadata={op_name="max"} + %max.70 = f32[1,5,1]{2,1,0} broadcast(%max.69), dimensions={0,2}, metadata={op_name="max"} + %Arg_0.22 = u32[2]{0} parameter(0) + %unstack.20 = u32[1]{0} slice(%Arg_0.22), slice={[0:1]}, metadata={op_name="unstack"} + %unstack.21 = u32[] reshape(%unstack.20), metadata={op_name="unstack"} + %unstack.22 = u32[1]{0} slice(%Arg_0.22), slice={[1:2]}, metadata={op_name="unstack"} + %unstack.23 = u32[] reshape(%unstack.22), metadata={op_name="unstack"} + %iota_2x32_shape.5 = u64[5]{0} iota(), iota_dimension=0, metadata={op_name="iota_2x32_shape"} + %iota_2x32_shape.6 = u64[1,5,1]{2,1,0} reshape(%iota_2x32_shape.5), metadata={op_name="iota_2x32_shape"} + %constant.92 = u64[] constant(32) + %broadcast.33 = u64[1,5,1]{2,1,0} broadcast(%constant.92), dimensions={} + %iota_2x32_shape.7 = u64[1,5,1]{2,1,0} shift-right-logical(%iota_2x32_shape.6, %broadcast.33), metadata={op_name="iota_2x32_shape"} + %iota_2x32_shape.9 = u32[1,5,1]{2,1,0} convert(%iota_2x32_shape.7), metadata={op_name="iota_2x32_shape"} + %iota_2x32_shape.8 = u32[1,5,1]{2,1,0} convert(%iota_2x32_shape.6), metadata={op_name="iota_2x32_shape"} + %call.1 = (u32[1,5,1]{2,1,0}, u32[1,5,1]{2,1,0}) call(%unstack.21, %unstack.23, %iota_2x32_shape.9, %iota_2x32_shape.8), to_apply=%threefry2x32.272 + %xor.13 = u32[1,5,1]{2,1,0} xor(%call.1#0, %call.1#1), metadata={op_name="xor"} + %constant.91 = u32[] constant(9) + %broadcast.32 = u32[1,5,1]{2,1,0} broadcast(%constant.91), dimensions={} + %shift_right_logical.17 = u32[1,5,1]{2,1,0} shift-right-logical(%xor.13, %broadcast.32), metadata={op_name="shift_right_logical"} + %constant.90 = u32[] constant(1065353216) + %broadcast.31 = u32[1,5,1]{2,1,0} broadcast(%constant.90), dimensions={} + %or.9 = u32[1,5,1]{2,1,0} or(%shift_right_logical.17, %broadcast.31), metadata={op_name="or"} + %bitcast_convert_type.1 = f32[1,5,1]{2,1,0} bitcast-convert(%or.9), metadata={op_name="bitcast_convert_type"} + %constant.89 = f32[] constant(1) + %broadcast.30 = f32[1,5,1]{2,1,0} broadcast(%constant.89), dimensions={} + %sub.268 = f32[1,5,1]{2,1,0} subtract(%bitcast_convert_type.1, %broadcast.30), metadata={op_name="sub"} + %Arg_1.17 = f32[] parameter(1) + %broadcast_in_dim.1033 = f32[1,1,1]{2,1,0} reshape(%Arg_1.17), metadata={op_name="broadcast_in_dim"} + %sub.269 = f32[1,1,1]{2,1,0} subtract(%broadcast_in_dim.1033, %broadcast_in_dim.1032), metadata={op_name="sub"} + %mul.2106 = f32[1,1,1]{2,1,0} broadcast(%sub.269), dimensions={0,1,2}, metadata={op_name="mul"} + %mul.2107 = f32[1,1]{1,0} reshape(%mul.2106), metadata={op_name="mul"} + %mul.2108 = f32[1,5,1]{2,1,0} broadcast(%mul.2107), dimensions={0,2}, metadata={op_name="mul"} + %mul.2109 = f32[1,5,1]{2,1,0} multiply(%sub.268, %mul.2108), metadata={op_name="mul"} + %add.686 = f32[1,1,1]{2,1,0} broadcast(%broadcast_in_dim.1032), dimensions={0,1,2}, metadata={op_name="add"} + %add.687 = f32[1,1]{1,0} reshape(%add.686), metadata={op_name="add"} + %add.688 = f32[1,5,1]{2,1,0} broadcast(%add.687), dimensions={0,2}, metadata={op_name="add"} + %add.689 = f32[1,5,1]{2,1,0} add(%mul.2109, %add.688), metadata={op_name="add"} + ROOT %max.71 = f32[1,5,1]{2,1,0} maximum(%max.70, %add.689), metadata={op_name="max"} +} + +%_gumbel.274 (Arg_0.23: u32[2]) -> f32[1,5,1] { + %Arg_0.23 = u32[2]{0} parameter(0) + %constant.93 = f32[] constant(1) + %constant.94 = f32[] constant(1.17549435e-38) + %jit__uniform_.1 = f32[1,5,1]{2,1,0} call(%Arg_0.23, %constant.93, %constant.94), to_apply=%_uniform.273, metadata={op_name="jit(_uniform)"} + %log.3 = f32[1,5,1]{2,1,0} log(%jit__uniform_.1), metadata={op_name="log"} + %neg.260 = f32[1,5,1]{2,1,0} negate(%log.3), metadata={op_name="neg"} + %log.4 = f32[1,5,1]{2,1,0} log(%neg.260), metadata={op_name="log"} + ROOT %neg.261 = f32[1,5,1]{2,1,0} negate(%log.4), metadata={op_name="neg"} +} + +%region_263.275 (reduce.4: f32[], reduce.5: s32[], reduce.6: f32[], reduce.7: s32[]) -> (f32[], s32[]) { + %reduce.4 = f32[] parameter(0), metadata={op_name="reduce"} + %reduce.6 = f32[] parameter(2), metadata={op_name="reduce"} + %gt.1 = pred[] compare(%reduce.4, %reduce.6), direction=GT, metadata={op_name="gt" stack_frame_id=118} + %ne.1 = pred[] compare(%reduce.4, %reduce.4), direction=NE, metadata={op_name="ne" stack_frame_id=118} + %or.12 = pred[] or(%gt.1, %ne.1), metadata={op_name="or" stack_frame_id=118} + %select_n.79 = f32[] select(%or.12, %reduce.4, %reduce.6), metadata={op_name="select_n" stack_frame_id=118} + %eq.1 = pred[] compare(%reduce.4, %reduce.6), direction=EQ, metadata={op_name="eq" stack_frame_id=118} + %reduce.5 = s32[] parameter(1), metadata={op_name="reduce"} + %reduce.7 = s32[] parameter(3), metadata={op_name="reduce"} + %lt.419 = pred[] compare(%reduce.5, %reduce.7), direction=LT, metadata={op_name="lt" stack_frame_id=118} + %and.194 = pred[] and(%eq.1, %lt.419), metadata={op_name="and" stack_frame_id=118} + %or.13 = pred[] or(%or.12, %and.194), metadata={op_name="or" stack_frame_id=118} + %select_n.80 = s32[] select(%or.13, %reduce.5, %reduce.7), metadata={op_name="select_n" stack_frame_id=118} + ROOT %tuple.7 = (f32[], s32[]) tuple(%select_n.79, %select_n.80) +} + +%argmax.276 (body.1: f32[1,5,1]) -> s32[1,1] { + %body.1 = f32[1,5,1]{2,1,0} parameter(0), metadata={op_name="while/body"} + %iota.380 = s32[5]{0} iota(), iota_dimension=0, metadata={op_name="iota" stack_frame_id=112} + %iota.381 = s32[1,5,1]{2,1,0} reshape(%iota.380), metadata={op_name="iota" stack_frame_id=112} + %constant.98 = f32[] constant(-inf) + %constant.97 = s32[] constant(0) + %reduce.11 = (f32[1,1]{1,0}, s32[1,1]{1,0}) reduce(%body.1, %iota.381, %constant.98, %constant.97), dimensions={1}, to_apply=%region_263.275, metadata={op_name="reduce" stack_frame_id=112} + ROOT %reduce.13 = s32[1,1]{1,0} get-tuple-element(%reduce.11), index=1, metadata={op_name="reduce" stack_frame_id=112} +} + +%region_264.277 (reduce_and.19: pred[], reduce_and.20: pred[]) -> pred[] { + %reduce_and.19 = pred[] parameter(0), metadata={op_name="reduce_and"} + %reduce_and.20 = pred[] parameter(1), metadata={op_name="reduce_and"} + ROOT %reduce_and.21 = pred[] and(%reduce_and.19, %reduce_and.20), metadata={op_name="reduce_and"} +} + +%take_along_axis.278 (Arg_0.25: s32[1,5], Arg_1.19: s32[1,1]) -> s32[1,1] { + %Arg_1.19 = s32[1,1]{1,0} parameter(1) + %constant.108 = s32[1,1]{1,0} constant({ {0} }) + %lt.421 = pred[1,1]{1,0} compare(%Arg_1.19, %constant.108), direction=LT, metadata={op_name="lt"} + %constant.106 = s32[1,1]{1,0} constant({ {5} }) + %add.692 = s32[1,1]{1,0} add(%Arg_1.19, %constant.106), metadata={op_name="add"} + %select_n.83 = s32[1,1]{1,0} select(%lt.421, %add.692, %Arg_1.19), metadata={op_name="select_n"} + %ge.352 = pred[1,1]{1,0} compare(%select_n.83, %constant.108), direction=GE, metadata={op_name="ge"} + %constant.105 = s32[1,1]{1,0} constant({ {4} }) + %le.7 = pred[1,1]{1,0} compare(%select_n.83, %constant.105), direction=LE, metadata={op_name="le"} + %and.196 = pred[1,1]{1,0} and(%ge.352, %le.7), metadata={op_name="and"} + %constant.107 = pred[] constant(true) + %reduce_and.23 = pred[1]{0} reduce(%and.196, %constant.107), dimensions={1}, to_apply=%region_264.277, metadata={op_name="reduce_and"} + %broadcast_in_dim.1035 = pred[1,1]{1,0} reshape(%reduce_and.23), metadata={op_name="broadcast_in_dim"} + %Arg_0.25 = s32[1,5]{1,0} parameter(0) + %gather.5 = s32[1,1]{1,0} gather(%Arg_0.25, %select_n.83), offset_dims={0}, collapsed_slice_dims={1}, start_index_map={1}, index_vector_dim=1, slice_sizes={1,1}, metadata={op_name="gather"} + %constant.104 = s32[1,1]{1,0} constant({ {-2147483648} }) + ROOT %select_n.84 = s32[1,1]{1,0} select(%broadcast_in_dim.1035, %gather.5, %constant.104), metadata={op_name="select_n"} +} + +%_where_5.279 (Arg_0.27: pred[1], Arg_1.21: s32[1], Arg_2.15: s32[1]) -> s32[1] { + %Arg_0.27 = pred[1]{0} parameter(0) + %Arg_1.21 = s32[1]{0} parameter(1) + %Arg_2.15 = s32[1]{0} parameter(2) + ROOT %select_n.88 = s32[1]{0} select(%Arg_0.27, %Arg_1.21, %Arg_2.15), metadata={op_name="select_n"} +} + +%region_128.280 (arg_tuple.5: (u32[2], s32[1,41], bf16[1,32,2,41,8,128], s32[], s32[], /*index=5*/pred[1,41], s32[], bf16[32000,4096], bf16[4096], bf16[4096,32,128], /*index=10*/bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], /*index=15*/bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], /*index=20*/bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], /*index=25*/bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], /*index=30*/bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], /*index=35*/bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], /*index=40*/bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], /*index=45*/bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], /*index=50*/bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], /*index=55*/bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], /*index=60*/bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], /*index=65*/bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], /*index=70*/bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], /*index=75*/bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], /*index=80*/bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], /*index=85*/bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], /*index=90*/bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], /*index=95*/bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], /*index=100*/bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], /*index=105*/bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], /*index=110*/bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], /*index=115*/bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], /*index=120*/bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], /*index=125*/bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], /*index=130*/bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], /*index=135*/bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], /*index=140*/bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], /*index=145*/bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], /*index=150*/bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], /*index=155*/bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], /*index=160*/bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], /*index=165*/bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], /*index=170*/bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], /*index=175*/bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], /*index=180*/bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], /*index=185*/bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], /*index=190*/bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], /*index=195*/bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], /*index=200*/bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], /*index=205*/bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], /*index=210*/bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], /*index=215*/bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], /*index=220*/bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], /*index=225*/bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], /*index=230*/bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], /*index=235*/bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], /*index=240*/bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], /*index=245*/bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], /*index=250*/bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], /*index=255*/bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], /*index=260*/bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], /*index=265*/bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], /*index=270*/bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], /*index=275*/bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], /*index=280*/bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], /*index=285*/bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], /*index=290*/bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], /*index=295*/bf16[14336,4096], bf16[4096], bf16[4096,32000])) -> (u32[2], s32[1,41], bf16[1,32,2,41,8,128], s32[], s32[], /*index=5*/pred[1,41], s32[], bf16[32000,4096], bf16[4096], bf16[4096,32,128], /*index=10*/bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], /*index=15*/bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], /*index=20*/bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], /*index=25*/bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], /*index=30*/bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], /*index=35*/bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], /*index=40*/bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], /*index=45*/bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], /*index=50*/bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], /*index=55*/bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], /*index=60*/bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], /*index=65*/bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], /*index=70*/bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], /*index=75*/bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], /*index=80*/bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], /*index=85*/bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], /*index=90*/bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], /*index=95*/bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], /*index=100*/bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], /*index=105*/bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], /*index=110*/bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], /*index=115*/bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], /*index=120*/bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], /*index=125*/bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], /*index=130*/bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], /*index=135*/bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], /*index=140*/bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], /*index=145*/bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], /*index=150*/bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], /*index=155*/bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], /*index=160*/bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], /*index=165*/bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], /*index=170*/bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], /*index=175*/bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], /*index=180*/bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], /*index=185*/bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], /*index=190*/bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], /*index=195*/bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], /*index=200*/bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], /*index=205*/bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], /*index=210*/bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], /*index=215*/bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], /*index=220*/bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], /*index=225*/bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], /*index=230*/bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], /*index=235*/bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], /*index=240*/bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], /*index=245*/bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], /*index=250*/bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], /*index=255*/bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], /*index=260*/bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], /*index=265*/bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], /*index=270*/bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], /*index=275*/bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], /*index=280*/bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], /*index=285*/bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], /*index=290*/bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], /*index=295*/bf16[14336,4096], bf16[4096], bf16[4096,32000]) { + %arg_tuple.5 = (u32[2]{0}, s32[1,41]{1,0}, bf16[1,32,2,41,8,128]{5,4,3,2,1,0}, s32[], s32[], /*index=5*/pred[1,41]{1,0}, s32[], bf16[32000,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=10*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=15*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=20*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=25*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=30*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=35*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=40*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=45*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=50*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=55*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=60*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=65*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=70*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=75*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=80*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=85*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=90*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=95*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=100*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=105*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=110*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=115*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=120*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=125*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=130*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=135*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=140*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=145*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=150*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=155*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=160*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=165*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=170*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=175*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=180*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=185*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=190*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=195*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=200*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=205*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=210*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=215*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=220*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=225*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=230*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=235*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=240*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=245*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=250*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=255*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=260*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=265*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=270*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=275*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=280*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=285*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=290*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=295*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32000]{1,0}) parameter(0) + %constant.109 = u32[2]{0} constant({0, 1}) + %add.1116 = u32[2]{0} add(%arg_tuple.5#0, %constant.109), metadata={op_name="jit(compiled_generate_function)/while/body/add" stack_frame_id=467} + %constant.127 = s32[] constant(0) + %lt.618 = pred[] compare(%arg_tuple.5#3, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/lt" stack_frame_id=476} + %constant.126 = s32[] constant(41) + %add.1118 = s32[] add(%arg_tuple.5#3, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/add" stack_frame_id=476} + %select_n.155 = s32[] select(%lt.618, %add.1118, %arg_tuple.5#3), metadata={op_name="jit(compiled_generate_function)/while/body/select_n" stack_frame_id=476} + %dynamic_slice.4 = pred[1,1]{1,0} dynamic-slice(%arg_tuple.5#5, %constant.127, %select_n.155), dynamic_slice_sizes={1,1}, metadata={op_name="jit(compiled_generate_function)/while/body/dynamic_slice" stack_frame_id=476} + %squeeze.198 = pred[1]{0} reshape(%dynamic_slice.4), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=476} + %lt.619 = pred[] compare(%arg_tuple.5#3, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/lt" stack_frame_id=477} + %add.1119 = s32[] add(%arg_tuple.5#3, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/add" stack_frame_id=477} + %select_n.156 = s32[] select(%lt.619, %add.1119, %arg_tuple.5#3), metadata={op_name="jit(compiled_generate_function)/while/body/select_n" stack_frame_id=477} + %dynamic_slice.5 = s32[1,1]{1,0} dynamic-slice(%arg_tuple.5#1, %constant.127, %select_n.156), dynamic_slice_sizes={1,1}, metadata={op_name="jit(compiled_generate_function)/while/body/dynamic_slice" stack_frame_id=477} + %squeeze.199 = s32[1]{0} reshape(%dynamic_slice.5), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=477} + %slice.189 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} slice(%arg_tuple.5#2), slice={[0:1], [31:32], [0:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/slice" stack_frame_id=133} + %squeeze.193 = bf16[1,2,41,8,128]{4,3,2,1,0} reshape(%slice.189), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=133} + %slice.191 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.193), slice={[0:1], [1:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/slice" stack_frame_id=246} + %squeeze.195 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.191), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/squeeze" stack_frame_id=246} + %slice.186 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} slice(%arg_tuple.5#2), slice={[0:1], [30:31], [0:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/slice" stack_frame_id=133} + %squeeze.190 = bf16[1,2,41,8,128]{4,3,2,1,0} reshape(%slice.186), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=133} + %slice.188 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.190), slice={[0:1], [1:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/slice" stack_frame_id=246} + %squeeze.192 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.188), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/squeeze" stack_frame_id=246} + %slice.183 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} slice(%arg_tuple.5#2), slice={[0:1], [29:30], [0:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/slice" stack_frame_id=133} + %squeeze.187 = bf16[1,2,41,8,128]{4,3,2,1,0} reshape(%slice.183), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=133} + %slice.185 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.187), slice={[0:1], [1:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/slice" stack_frame_id=246} + %squeeze.189 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.185), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/squeeze" stack_frame_id=246} + %slice.180 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} slice(%arg_tuple.5#2), slice={[0:1], [28:29], [0:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/slice" stack_frame_id=133} + %squeeze.184 = bf16[1,2,41,8,128]{4,3,2,1,0} reshape(%slice.180), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=133} + %slice.182 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.184), slice={[0:1], [1:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/slice" stack_frame_id=246} + %squeeze.186 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.182), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/squeeze" stack_frame_id=246} + %slice.177 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} slice(%arg_tuple.5#2), slice={[0:1], [27:28], [0:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/slice" stack_frame_id=133} + %squeeze.181 = bf16[1,2,41,8,128]{4,3,2,1,0} reshape(%slice.177), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=133} + %slice.179 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.181), slice={[0:1], [1:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/slice" stack_frame_id=246} + %squeeze.183 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.179), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/squeeze" stack_frame_id=246} + %slice.174 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} slice(%arg_tuple.5#2), slice={[0:1], [26:27], [0:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/slice" stack_frame_id=133} + %squeeze.178 = bf16[1,2,41,8,128]{4,3,2,1,0} reshape(%slice.174), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=133} + %slice.176 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.178), slice={[0:1], [1:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/slice" stack_frame_id=246} + %squeeze.180 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.176), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/squeeze" stack_frame_id=246} + %slice.171 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} slice(%arg_tuple.5#2), slice={[0:1], [25:26], [0:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/slice" stack_frame_id=133} + %squeeze.175 = bf16[1,2,41,8,128]{4,3,2,1,0} reshape(%slice.171), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=133} + %slice.173 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.175), slice={[0:1], [1:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/slice" stack_frame_id=246} + %squeeze.177 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.173), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/squeeze" stack_frame_id=246} + %slice.168 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} slice(%arg_tuple.5#2), slice={[0:1], [24:25], [0:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/slice" stack_frame_id=133} + %squeeze.172 = bf16[1,2,41,8,128]{4,3,2,1,0} reshape(%slice.168), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=133} + %slice.170 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.172), slice={[0:1], [1:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/slice" stack_frame_id=246} + %squeeze.174 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.170), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/squeeze" stack_frame_id=246} + %slice.165 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} slice(%arg_tuple.5#2), slice={[0:1], [23:24], [0:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/slice" stack_frame_id=133} + %squeeze.169 = bf16[1,2,41,8,128]{4,3,2,1,0} reshape(%slice.165), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=133} + %slice.167 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.169), slice={[0:1], [1:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/slice" stack_frame_id=246} + %squeeze.171 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.167), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/squeeze" stack_frame_id=246} + %slice.162 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} slice(%arg_tuple.5#2), slice={[0:1], [22:23], [0:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/slice" stack_frame_id=133} + %squeeze.166 = bf16[1,2,41,8,128]{4,3,2,1,0} reshape(%slice.162), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=133} + %slice.164 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.166), slice={[0:1], [1:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/slice" stack_frame_id=246} + %squeeze.168 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.164), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/squeeze" stack_frame_id=246} + %slice.159 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} slice(%arg_tuple.5#2), slice={[0:1], [21:22], [0:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/slice" stack_frame_id=133} + %squeeze.163 = bf16[1,2,41,8,128]{4,3,2,1,0} reshape(%slice.159), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=133} + %slice.161 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.163), slice={[0:1], [1:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/slice" stack_frame_id=246} + %squeeze.165 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.161), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/squeeze" stack_frame_id=246} + %slice.156 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} slice(%arg_tuple.5#2), slice={[0:1], [20:21], [0:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/slice" stack_frame_id=133} + %squeeze.160 = bf16[1,2,41,8,128]{4,3,2,1,0} reshape(%slice.156), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=133} + %slice.158 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.160), slice={[0:1], [1:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/slice" stack_frame_id=246} + %squeeze.162 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.158), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/squeeze" stack_frame_id=246} + %slice.153 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} slice(%arg_tuple.5#2), slice={[0:1], [19:20], [0:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/slice" stack_frame_id=133} + %squeeze.157 = bf16[1,2,41,8,128]{4,3,2,1,0} reshape(%slice.153), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=133} + %slice.155 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.157), slice={[0:1], [1:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/slice" stack_frame_id=246} + %squeeze.159 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.155), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/squeeze" stack_frame_id=246} + %slice.150 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} slice(%arg_tuple.5#2), slice={[0:1], [18:19], [0:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/slice" stack_frame_id=133} + %squeeze.154 = bf16[1,2,41,8,128]{4,3,2,1,0} reshape(%slice.150), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=133} + %slice.152 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.154), slice={[0:1], [1:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/slice" stack_frame_id=246} + %squeeze.156 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.152), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/squeeze" stack_frame_id=246} + %slice.147 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} slice(%arg_tuple.5#2), slice={[0:1], [17:18], [0:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/slice" stack_frame_id=133} + %squeeze.151 = bf16[1,2,41,8,128]{4,3,2,1,0} reshape(%slice.147), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=133} + %slice.149 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.151), slice={[0:1], [1:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/slice" stack_frame_id=246} + %squeeze.153 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.149), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/squeeze" stack_frame_id=246} + %slice.144 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} slice(%arg_tuple.5#2), slice={[0:1], [16:17], [0:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/slice" stack_frame_id=133} + %squeeze.148 = bf16[1,2,41,8,128]{4,3,2,1,0} reshape(%slice.144), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=133} + %slice.146 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.148), slice={[0:1], [1:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/slice" stack_frame_id=246} + %squeeze.150 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.146), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/squeeze" stack_frame_id=246} + %slice.141 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} slice(%arg_tuple.5#2), slice={[0:1], [15:16], [0:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/slice" stack_frame_id=133} + %squeeze.145 = bf16[1,2,41,8,128]{4,3,2,1,0} reshape(%slice.141), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=133} + %slice.143 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.145), slice={[0:1], [1:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/slice" stack_frame_id=246} + %squeeze.147 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.143), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/squeeze" stack_frame_id=246} + %slice.138 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} slice(%arg_tuple.5#2), slice={[0:1], [14:15], [0:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/slice" stack_frame_id=133} + %squeeze.142 = bf16[1,2,41,8,128]{4,3,2,1,0} reshape(%slice.138), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=133} + %slice.140 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.142), slice={[0:1], [1:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/slice" stack_frame_id=246} + %squeeze.144 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.140), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/squeeze" stack_frame_id=246} + %slice.135 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} slice(%arg_tuple.5#2), slice={[0:1], [13:14], [0:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/slice" stack_frame_id=133} + %squeeze.139 = bf16[1,2,41,8,128]{4,3,2,1,0} reshape(%slice.135), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=133} + %slice.137 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.139), slice={[0:1], [1:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/slice" stack_frame_id=246} + %squeeze.141 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.137), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/squeeze" stack_frame_id=246} + %slice.132 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} slice(%arg_tuple.5#2), slice={[0:1], [12:13], [0:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/slice" stack_frame_id=133} + %squeeze.136 = bf16[1,2,41,8,128]{4,3,2,1,0} reshape(%slice.132), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=133} + %slice.134 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.136), slice={[0:1], [1:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/slice" stack_frame_id=246} + %squeeze.138 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.134), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/squeeze" stack_frame_id=246} + %slice.129 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} slice(%arg_tuple.5#2), slice={[0:1], [11:12], [0:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/slice" stack_frame_id=133} + %squeeze.133 = bf16[1,2,41,8,128]{4,3,2,1,0} reshape(%slice.129), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=133} + %slice.131 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.133), slice={[0:1], [1:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/slice" stack_frame_id=246} + %squeeze.135 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.131), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/squeeze" stack_frame_id=246} + %slice.126 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} slice(%arg_tuple.5#2), slice={[0:1], [10:11], [0:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/slice" stack_frame_id=133} + %squeeze.130 = bf16[1,2,41,8,128]{4,3,2,1,0} reshape(%slice.126), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=133} + %slice.128 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.130), slice={[0:1], [1:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/slice" stack_frame_id=246} + %squeeze.132 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.128), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/squeeze" stack_frame_id=246} + %slice.123 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} slice(%arg_tuple.5#2), slice={[0:1], [9:10], [0:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/slice" stack_frame_id=133} + %squeeze.127 = bf16[1,2,41,8,128]{4,3,2,1,0} reshape(%slice.123), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=133} + %slice.125 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.127), slice={[0:1], [1:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/slice" stack_frame_id=246} + %squeeze.129 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.125), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/squeeze" stack_frame_id=246} + %slice.120 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} slice(%arg_tuple.5#2), slice={[0:1], [8:9], [0:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/slice" stack_frame_id=133} + %squeeze.124 = bf16[1,2,41,8,128]{4,3,2,1,0} reshape(%slice.120), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=133} + %slice.122 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.124), slice={[0:1], [1:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/slice" stack_frame_id=246} + %squeeze.126 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/squeeze" stack_frame_id=246} + %slice.117 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} slice(%arg_tuple.5#2), slice={[0:1], [7:8], [0:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/slice" stack_frame_id=133} + %squeeze.121 = bf16[1,2,41,8,128]{4,3,2,1,0} reshape(%slice.117), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=133} + %slice.119 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.121), slice={[0:1], [1:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/slice" stack_frame_id=246} + %squeeze.123 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.119), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/squeeze" stack_frame_id=246} + %slice.114 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} slice(%arg_tuple.5#2), slice={[0:1], [6:7], [0:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/slice" stack_frame_id=133} + %squeeze.118 = bf16[1,2,41,8,128]{4,3,2,1,0} reshape(%slice.114), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=133} + %slice.116 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.118), slice={[0:1], [1:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/slice" stack_frame_id=246} + %squeeze.120 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.116), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/squeeze" stack_frame_id=246} + %slice.111 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} slice(%arg_tuple.5#2), slice={[0:1], [5:6], [0:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/slice" stack_frame_id=133} + %squeeze.115 = bf16[1,2,41,8,128]{4,3,2,1,0} reshape(%slice.111), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=133} + %slice.113 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.115), slice={[0:1], [1:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/slice" stack_frame_id=246} + %squeeze.117 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.113), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/squeeze" stack_frame_id=246} + %slice.108 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} slice(%arg_tuple.5#2), slice={[0:1], [4:5], [0:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/slice" stack_frame_id=133} + %squeeze.112 = bf16[1,2,41,8,128]{4,3,2,1,0} reshape(%slice.108), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=133} + %slice.110 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.112), slice={[0:1], [1:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/slice" stack_frame_id=246} + %squeeze.114 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.110), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/squeeze" stack_frame_id=246} + %slice.105 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} slice(%arg_tuple.5#2), slice={[0:1], [3:4], [0:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/slice" stack_frame_id=133} + %squeeze.109 = bf16[1,2,41,8,128]{4,3,2,1,0} reshape(%slice.105), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=133} + %slice.107 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.109), slice={[0:1], [1:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/slice" stack_frame_id=246} + %squeeze.111 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.107), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/squeeze" stack_frame_id=246} + %slice.102 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} slice(%arg_tuple.5#2), slice={[0:1], [2:3], [0:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/slice" stack_frame_id=133} + %squeeze.106 = bf16[1,2,41,8,128]{4,3,2,1,0} reshape(%slice.102), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=133} + %slice.104 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.106), slice={[0:1], [1:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/slice" stack_frame_id=246} + %squeeze.108 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.104), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/squeeze" stack_frame_id=246} + %slice.99 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} slice(%arg_tuple.5#2), slice={[0:1], [1:2], [0:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/slice" stack_frame_id=133} + %squeeze.103 = bf16[1,2,41,8,128]{4,3,2,1,0} reshape(%slice.99), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=133} + %slice.101 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.103), slice={[0:1], [1:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/slice" stack_frame_id=246} + %squeeze.105 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.101), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/squeeze" stack_frame_id=246} + %slice.96 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} slice(%arg_tuple.5#2), slice={[0:1], [0:1], [0:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/slice" stack_frame_id=133} + %squeeze.100 = bf16[1,2,41,8,128]{4,3,2,1,0} reshape(%slice.96), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=133} + %slice.98 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.100), slice={[0:1], [1:2], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/slice" stack_frame_id=246} + %squeeze.102 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.98), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/squeeze" stack_frame_id=246} + %constant.128 = s32[] constant(1) + %sub.270 = s32[] subtract(%arg_tuple.5#3, %constant.128), metadata={op_name="jit(compiled_generate_function)/while/body/sub" stack_frame_id=119} + %lt.425 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/lt" stack_frame_id=122} + %add.698 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/add" stack_frame_id=122} + %select_n.90 = s32[] select(%lt.425, %add.698, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/select_n" stack_frame_id=122} + %dynamic_slice.3 = s32[1,1]{1,0} dynamic-slice(%arg_tuple.5#1, %constant.127, %select_n.90), dynamic_slice_sizes={1,1}, metadata={op_name="jit(compiled_generate_function)/while/body/dynamic_slice" stack_frame_id=122} + %jit__take_.2 = bf16[1,1,4096]{2,1,0} call(%arg_tuple.5#7, %dynamic_slice.3), to_apply=%_take_1.135, metadata={op_name="jit(compiled_generate_function)/while/body/token_embedding/jit(_take)" stack_frame_id=132} + %convert_element_type.1120 = f32[1,1,4096]{2,1,0} convert(%jit__take_.2), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention_layernorm/convert_element_type" stack_frame_id=164} + %constant.124 = f32[] constant(2) + %broadcast.42 = f32[1,1,4096]{2,1,0} broadcast(%constant.124), dimensions={} + %pow.255 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1120, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention_layernorm/pow" stack_frame_id=167} + %constant.123 = f32[] constant(0) + %reduce_sum.1351 = f32[1,1]{1,0} reduce(%pow.255, %constant.123), dimensions={2}, to_apply=%region_130.136, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention_layernorm/reduce_sum" stack_frame_id=73} + %broadcast_in_dim.1041 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1351), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention_layernorm/broadcast_in_dim" stack_frame_id=73} + %constant.122 = f32[1,1,1]{2,1,0} constant({ { {4096} } }) + %div.643 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1041, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention_layernorm/div" stack_frame_id=73} + %constant.121 = f32[1,1,1]{2,1,0} constant({ { {1e-05} } }) + %add.702 = f32[1,1,1]{2,1,0} add(%div.643, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention_layernorm/add" stack_frame_id=168} + %rsqrt.128 = f32[1,1,1]{2,1,0} rsqrt(%add.702), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention_layernorm/rsqrt" stack_frame_id=171} + %mul.2110 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.128), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2111 = f32[1,1]{1,0} reshape(%mul.2110), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2112 = f32[1,1,4096]{2,1,0} broadcast(%mul.2111), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2113 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1120, %mul.2112), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention_layernorm/mul" stack_frame_id=172} + %convert_element_type.1121 = f32[4096]{0} convert(%arg_tuple.5#8), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention_layernorm/convert_element_type" stack_frame_id=173} + %broadcast_in_dim.1042 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention_layernorm/broadcast_in_dim" stack_frame_id=173} + %mul.2114 = f32[1,1,4096]{2,1,0} multiply(%mul.2113, %broadcast_in_dim.1042), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention_layernorm/mul" stack_frame_id=173} + %convert_element_type.1122 = bf16[1,1,4096]{2,1,0} convert(%mul.2114), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention_layernorm/convert_element_type" stack_frame_id=177} + %dot_general.700 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1122, %arg_tuple.5#11), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=265} + %lt.431 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/lt" stack_frame_id=329} + %add.708 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/add" stack_frame_id=329} + %select_n.92 = s32[] select(%lt.431, %add.708, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/select_n" stack_frame_id=329} + %dynamic_update_slice.66 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.102, %dot_general.700, %constant.127, %select_n.92, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/dynamic_update_slice" stack_frame_id=329} + %broadcast_in_dim.1044 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.66), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/broadcast_in_dim" stack_frame_id=338} + %reshape.479 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1044), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/reshape" stack_frame_id=338} + %convert_element_type.1117 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/convert_element_type" stack_frame_id=139} + %add.699 = f32[1,1]{1,0} reshape(%convert_element_type.1117), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_0/add"} + %ge.353 = f32[1,1]{1,0} broadcast(%add.699), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/ge" stack_frame_id=143} + %ge.354 = f32[1]{0} reshape(%ge.353), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/ge" stack_frame_id=143} + %ge.355 = f32[1,41]{1,0} broadcast(%ge.354), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/ge" stack_frame_id=143} + %iota.382 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/iota" stack_frame_id=142} + %broadcast_in_dim.1037 = f32[1,41]{1,0} reshape(%iota.382), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/broadcast_in_dim" stack_frame_id=143} + %ge.356 = pred[1,41]{1,0} compare(%ge.355, %broadcast_in_dim.1037), direction=GE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/ge" stack_frame_id=143} + %broadcast_in_dim.1038 = pred[1,1,41]{2,1,0} reshape(%ge.356), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/broadcast_in_dim" stack_frame_id=146} + %convert_element_type.1119 = s32[1,1,41]{2,1,0} convert(%broadcast_in_dim.1038), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/convert_element_type" stack_frame_id=160} + %add.700 = s32[1,1]{1,0} reshape(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/add" stack_frame_id=147} + %lt.426 = s32[1,1]{1,0} broadcast(%add.700), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/lt" stack_frame_id=152} + %lt.427 = s32[1]{0} reshape(%lt.426), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/lt" stack_frame_id=152} + %lt.428 = s32[1,41]{1,0} broadcast(%lt.427), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/lt" stack_frame_id=152} + %iota.383 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/iota" stack_frame_id=150} + %broadcast_in_dim.1039 = s32[1,41]{1,0} reshape(%iota.383), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/broadcast_in_dim" stack_frame_id=148} + %constant.125 = s32[] constant(4096) + %broadcast.43 = s32[1,41]{1,0} broadcast(%constant.125), dimensions={} + %add.701 = s32[1,41]{1,0} add(%broadcast_in_dim.1039, %broadcast.43), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/add" stack_frame_id=151} + %lt.429 = pred[1,41]{1,0} compare(%lt.428, %add.701), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/lt" stack_frame_id=152} + %convert_element_type.1118 = s32[1,41]{1,0} convert(%lt.429), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/convert_element_type" stack_frame_id=156} + %broadcast_in_dim.1040 = s32[1,1,41]{2,1,0} reshape(%convert_element_type.1118), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/broadcast_in_dim" stack_frame_id=160} + %min.63 = s32[1,1,41]{2,1,0} minimum(%convert_element_type.1119, %broadcast_in_dim.1040), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/min" stack_frame_id=160} + %broadcast_in_dim.1045 = s32[1,1,1,41]{3,2,1,0} reshape(%min.63), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/broadcast_in_dim" stack_frame_id=341} + %constant.116 = s32[] constant(0) + %broadcast.37 = s32[1,1,1,41]{3,2,1,0} broadcast(%constant.116), dimensions={} + %convert_element_type.1129 = pred[1,1,1,41]{3,2,1,0} compare(%broadcast_in_dim.1045, %broadcast.37), direction=NE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/convert_element_type" stack_frame_id=345} + %broadcast_in_dim.1046 = pred[1,1,1,1,41]{4,3,2,1,0} reshape(%convert_element_type.1129), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/broadcast_in_dim" stack_frame_id=83} + %and.197 = pred[1,1,1,1,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1046), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/vmap()/and" stack_frame_id=83} + %and.198 = pred[1,1,1,41]{3,2,1,0} reshape(%and.197), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/vmap()/and" stack_frame_id=83} + %and.199 = pred[1,1,32,1,41]{4,3,2,1,0} broadcast(%and.198), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/vmap()/and" stack_frame_id=83} + %dot_general.697 = bf16[1,1,32,128]{3,2,1,0} dot(%convert_element_type.1122, %arg_tuple.5#9), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=186} + %convert_element_type.1123 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/convert_element_type" stack_frame_id=208} + %add.703 = f32[1,1]{1,0} reshape(%convert_element_type.1123), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/add"} + %constant.118 = f32[] constant(10000) + %broadcast.39 = f32[64]{0} broadcast(%constant.118), dimensions={} + %iota.384 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/iota" stack_frame_id=197} + %constant.120 = f32[] constant(2) + %broadcast.41 = f32[64]{0} broadcast(%constant.120), dimensions={} + %mul.2115 = f32[64]{0} multiply(%iota.384, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=197} + %constant.119 = f32[] constant(128) + %broadcast.40 = f32[64]{0} broadcast(%constant.119), dimensions={} + %div.644 = f32[64]{0} divide(%mul.2115, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/div" stack_frame_id=198} + %neg.262 = f32[64]{0} negate(%div.644), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/neg" stack_frame_id=199} + %pow.256 = f32[64]{0} power(%broadcast.39, %neg.262), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/pow" stack_frame_id=202} + %constant.117 = f32[] constant(1) + %broadcast.38 = f32[64]{0} broadcast(%constant.117), dimensions={} + %div.645 = f32[64]{0} divide(%pow.256, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/div" stack_frame_id=203} + %dot_general.698 = f32[1,1,64]{2,1,0} dot(%add.703, %div.645), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/bi,j->bij/dot_general" stack_frame_id=211} + %stack.1020 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.698), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/stack" stack_frame_id=214} + %stack.1021 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.698), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/stack" stack_frame_id=214} + %stack.1022 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1020, %stack.1021), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/stack" stack_frame_id=214} + %reshape.474 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1022), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/reshape"} + %cos.127 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.474), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/cos" stack_frame_id=218} + %convert_element_type.1124 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/convert_element_type" stack_frame_id=222} + %mul.2116 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1124), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=242} + %mul.2117 = bf16[1,1,128]{2,1,0} reshape(%mul.2116), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=242} + %mul.2118 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2117), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=242} + %mul.2119 = bf16[1,1,32,128]{3,2,1,0} multiply(%dot_general.697, %mul.2118), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=242} + %split.255 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.697), slice={[0:1], [0:1], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/split" stack_frame_id=234} + %neg.263 = bf16[1,1,32,64]{3,2,1,0} negate(%split.255), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/neg" stack_frame_id=235} + %stack.1023 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%neg.263), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/stack" stack_frame_id=238} + %split.254 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.697), slice={[0:1], [0:1], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/split" stack_frame_id=234} + %stack.1024 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%split.254), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/stack" stack_frame_id=238} + %stack.1025 = bf16[1,1,32,2,64]{4,3,2,1,0} concatenate(%stack.1023, %stack.1024), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/stack" stack_frame_id=238} + %reshape.475 = bf16[1,1,32,128]{3,2,1,0} reshape(%stack.1025), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/reshape" stack_frame_id=241} + %sin.127 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.474), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/sin" stack_frame_id=226} + %convert_element_type.1125 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/convert_element_type" stack_frame_id=230} + %mul.2120 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1125), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=243} + %mul.2121 = bf16[1,1,128]{2,1,0} reshape(%mul.2120), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=243} + %mul.2122 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2121), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=243} + %mul.2123 = bf16[1,1,32,128]{3,2,1,0} multiply(%reshape.475, %mul.2122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=243} + %add.704 = bf16[1,1,32,128]{3,2,1,0} add(%mul.2119, %mul.2123), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/add" stack_frame_id=244} + %reshape.480 = bf16[1,1,32,1,128]{4,3,2,1,0} reshape(%add.704), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/reshape" stack_frame_id=83} + %slice.97 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.100), slice={[0:1], [0:1], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/slice" stack_frame_id=245} + %squeeze.101 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.97), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/squeeze" stack_frame_id=245} + %dot_general.699 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1122, %arg_tuple.5#10), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=256} + %convert_element_type.1126 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/convert_element_type" stack_frame_id=287} + %add.705 = f32[1,1]{1,0} reshape(%convert_element_type.1126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/add"} + %iota.385 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/iota" stack_frame_id=276} + %mul.2124 = f32[64]{0} multiply(%iota.385, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=276} + %div.646 = f32[64]{0} divide(%mul.2124, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/div" stack_frame_id=277} + %neg.264 = f32[64]{0} negate(%div.646), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/neg" stack_frame_id=278} + %pow.257 = f32[64]{0} power(%broadcast.39, %neg.264), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/pow" stack_frame_id=281} + %div.647 = f32[64]{0} divide(%pow.257, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/div" stack_frame_id=282} + %dot_general.701 = f32[1,1,64]{2,1,0} dot(%add.705, %div.647), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/bi,j->bij/dot_general" stack_frame_id=290} + %stack.1026 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.701), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/stack" stack_frame_id=293} + %stack.1027 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.701), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/stack" stack_frame_id=293} + %stack.1028 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1026, %stack.1027), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/stack" stack_frame_id=293} + %reshape.476 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1028), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/reshape"} + %cos.128 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.476), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/cos" stack_frame_id=297} + %convert_element_type.1127 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.128), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/convert_element_type" stack_frame_id=301} + %mul.2125 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1127), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=321} + %mul.2126 = bf16[1,1,128]{2,1,0} reshape(%mul.2125), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=321} + %mul.2127 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2126), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=321} + %mul.2128 = bf16[1,1,8,128]{3,2,1,0} multiply(%dot_general.699, %mul.2127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=321} + %split.257 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.699), slice={[0:1], [0:1], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/split" stack_frame_id=313} + %neg.265 = bf16[1,1,8,64]{3,2,1,0} negate(%split.257), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/neg" stack_frame_id=314} + %stack.1029 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%neg.265), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/stack" stack_frame_id=317} + %split.256 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.699), slice={[0:1], [0:1], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/split" stack_frame_id=313} + %stack.1030 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%split.256), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/stack" stack_frame_id=317} + %stack.1031 = bf16[1,1,8,2,64]{4,3,2,1,0} concatenate(%stack.1029, %stack.1030), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/stack" stack_frame_id=317} + %reshape.477 = bf16[1,1,8,128]{3,2,1,0} reshape(%stack.1031), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/reshape" stack_frame_id=320} + %sin.128 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.476), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/sin" stack_frame_id=305} + %convert_element_type.1128 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.128), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/convert_element_type" stack_frame_id=309} + %mul.2129 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1128), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=322} + %mul.2130 = bf16[1,1,128]{2,1,0} reshape(%mul.2129), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=322} + %mul.2131 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2130), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=322} + %mul.2132 = bf16[1,1,8,128]{3,2,1,0} multiply(%reshape.477, %mul.2131), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=322} + %add.706 = bf16[1,1,8,128]{3,2,1,0} add(%mul.2128, %mul.2132), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/rotary_embedding/add" stack_frame_id=323} + %lt.430 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/lt" stack_frame_id=326} + %add.707 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/add" stack_frame_id=326} + %select_n.91 = s32[] select(%lt.430, %add.707, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/select_n" stack_frame_id=326} + %dynamic_update_slice.65 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.101, %add.706, %constant.127, %select_n.91, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/dynamic_update_slice" stack_frame_id=326} + %broadcast_in_dim.1043 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.65), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/broadcast_in_dim" stack_frame_id=335} + %reshape.478 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1043), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/reshape" stack_frame_id=335} + %dot_general.702 = f32[1,32,1,1,41]{4,3,2,1,0} dot(%reshape.480, %reshape.478), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=83} + %constant.115 = f32[] constant(0.0883883461) + %broadcast.36 = f32[1,32,1,1,41]{4,3,2,1,0} broadcast(%constant.115), dimensions={} + %mul.2133 = f32[1,32,1,1,41]{4,3,2,1,0} multiply(%dot_general.702, %broadcast.36), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/vmap()/mul" stack_frame_id=83} + %constant.114 = f32[] constant(-2.38197633e+38) + %vmap_jit__where__.63 = f32[1,1,32,1,41]{4,3,2,1,0} call(%and.199, %mul.2133, %constant.114), to_apply=%_where_3.137, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/vmap(jit(_where))" stack_frame_id=83} + %constant.113 = f32[] constant(-inf) + %reduce_max.448 = f32[1,1,32,1]{3,2,1,0} reduce(%vmap_jit__where__.63, %constant.113), dimensions={4}, to_apply=%region_131.138, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/vmap()/reduce_max" stack_frame_id=83} + %constant.112 = f32[] constant(-inf) + %broadcast.35 = f32[1,1,32,1]{3,2,1,0} broadcast(%constant.112), dimensions={} + %max.72 = f32[1,1,32,1]{3,2,1,0} maximum(%reduce_max.448, %broadcast.35), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/vmap()/max" stack_frame_id=83} + %broadcast_in_dim.1047 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%max.72), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %sub.271 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1047), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/vmap()/sub" stack_frame_id=83} + %sub.272 = f32[1,1,32,1]{3,2,1,0} reshape(%sub.271), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/vmap()/sub" stack_frame_id=83} + %sub.273 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%sub.272), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/vmap()/sub" stack_frame_id=83} + %sub.274 = f32[1,1,32,1,41]{4,3,2,1,0} subtract(%vmap_jit__where__.63, %sub.273), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/vmap()/sub" stack_frame_id=83} + %exp.68 = f32[1,1,32,1,41]{4,3,2,1,0} exponential(%sub.274), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/vmap()/exp" stack_frame_id=83} + %reduce_sum.1352 = f32[1,1,32,1]{3,2,1,0} reduce(%exp.68, %constant.123), dimensions={4}, to_apply=%region_132.139, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/vmap()/reduce_sum" stack_frame_id=83} + %broadcast_in_dim.1048 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%reduce_sum.1352), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %div.648 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1048), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/vmap()/div" stack_frame_id=83} + %div.649 = f32[1,1,32,1]{3,2,1,0} reshape(%div.648), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/vmap()/div" stack_frame_id=83} + %div.650 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%div.649), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/vmap()/div" stack_frame_id=83} + %div.651 = f32[1,1,32,1,41]{4,3,2,1,0} divide(%exp.68, %div.650), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/vmap()/div" stack_frame_id=83} + %convert_element_type.1130 = bf16[1,1,32,1,41]{4,3,2,1,0} convert(%div.651), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/vmap()/convert_element_type" stack_frame_id=83} + %dot_general.703 = bf16[1,32,128,1,1]{4,3,2,1,0} dot(%reshape.479, %convert_element_type.1130), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=83} + %transpose.67 = bf16[1,1,32,128]{3,2,1,0} reshape(%dot_general.703), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/reshape;jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/vmap()/transpose;jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %dot_general.704 = bf16[1,1,4096]{2,1,0} dot(%transpose.67, %arg_tuple.5#12), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=354} + %add.709 = bf16[1,1,4096]{2,1,0} add(%dot_general.704, %jit__take_.2), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/add" stack_frame_id=355} + %convert_element_type.1131 = f32[1,1,4096]{2,1,0} convert(%add.709), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/feedforward_layernorm/convert_element_type" stack_frame_id=359} + %pow.258 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1131, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/feedforward_layernorm/pow" stack_frame_id=362} + %reduce_sum.1353 = f32[1,1]{1,0} reduce(%pow.258, %constant.123), dimensions={2}, to_apply=%region_133.140, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/feedforward_layernorm/reduce_sum" stack_frame_id=92} + %broadcast_in_dim.1049 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1353), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/feedforward_layernorm/broadcast_in_dim" stack_frame_id=92} + %div.652 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1049, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/feedforward_layernorm/div" stack_frame_id=92} + %add.710 = f32[1,1,1]{2,1,0} add(%div.652, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/feedforward_layernorm/add" stack_frame_id=363} + %rsqrt.129 = f32[1,1,1]{2,1,0} rsqrt(%add.710), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/feedforward_layernorm/rsqrt" stack_frame_id=366} + %mul.2134 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.129), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2135 = f32[1,1]{1,0} reshape(%mul.2134), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2136 = f32[1,1,4096]{2,1,0} broadcast(%mul.2135), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2137 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1131, %mul.2136), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/feedforward_layernorm/mul" stack_frame_id=367} + %convert_element_type.1132 = f32[4096]{0} convert(%arg_tuple.5#13), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/feedforward_layernorm/convert_element_type" stack_frame_id=368} + %broadcast_in_dim.1050 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1132), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/feedforward_layernorm/broadcast_in_dim" stack_frame_id=368} + %mul.2138 = f32[1,1,4096]{2,1,0} multiply(%mul.2137, %broadcast_in_dim.1050), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/feedforward_layernorm/mul" stack_frame_id=368} + %convert_element_type.1133 = bf16[1,1,4096]{2,1,0} convert(%mul.2138), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/feedforward_layernorm/convert_element_type" stack_frame_id=372} + %dot_general.706 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1133, %arg_tuple.5#15), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/feedforward_intermediate_dense/dot_general" stack_frame_id=401} + %dot_general.705 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1133, %arg_tuple.5#14), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/feedforward_gate_dense/dot_general" stack_frame_id=381} + %convert_element_type.1134 = f32[1,1,14336]{2,1,0} convert(%dot_general.705), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/convert_element_type" stack_frame_id=385} + %jit_silu_.63 = f32[1,1,14336]{2,1,0} call(%convert_element_type.1134), to_apply=%silu_4.141, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/jit(silu)" stack_frame_id=388} + %convert_element_type.1135 = bf16[1,1,14336]{2,1,0} convert(%jit_silu_.63), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/convert_element_type" stack_frame_id=392} + %mul.2139 = bf16[1,1,14336]{2,1,0} multiply(%dot_general.706, %convert_element_type.1135), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/mul" stack_frame_id=404} + %dot_general.707 = bf16[1,1,4096]{2,1,0} dot(%mul.2139, %arg_tuple.5#16), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/feedforward_output_dense/dot_general" stack_frame_id=413} + %add.711 = bf16[1,1,4096]{2,1,0} add(%dot_general.707, %add.709), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/add" stack_frame_id=414} + %convert_element_type.1139 = f32[1,1,4096]{2,1,0} convert(%add.711), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention_layernorm/convert_element_type" stack_frame_id=164} + %pow.259 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1139, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention_layernorm/pow" stack_frame_id=167} + %reduce_sum.1354 = f32[1,1]{1,0} reduce(%pow.259, %constant.123), dimensions={2}, to_apply=%region_134.142, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention_layernorm/reduce_sum" stack_frame_id=73} + %broadcast_in_dim.1055 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1354), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention_layernorm/broadcast_in_dim" stack_frame_id=73} + %div.653 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1055, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention_layernorm/div" stack_frame_id=73} + %add.715 = f32[1,1,1]{2,1,0} add(%div.653, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention_layernorm/add" stack_frame_id=168} + %rsqrt.130 = f32[1,1,1]{2,1,0} rsqrt(%add.715), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention_layernorm/rsqrt" stack_frame_id=171} + %mul.2140 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.130), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2141 = f32[1,1]{1,0} reshape(%mul.2140), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2142 = f32[1,1,4096]{2,1,0} broadcast(%mul.2141), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2143 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1139, %mul.2142), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention_layernorm/mul" stack_frame_id=172} + %convert_element_type.1140 = f32[4096]{0} convert(%arg_tuple.5#17), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention_layernorm/convert_element_type" stack_frame_id=173} + %broadcast_in_dim.1056 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1140), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention_layernorm/broadcast_in_dim" stack_frame_id=173} + %mul.2144 = f32[1,1,4096]{2,1,0} multiply(%mul.2143, %broadcast_in_dim.1056), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention_layernorm/mul" stack_frame_id=173} + %convert_element_type.1141 = bf16[1,1,4096]{2,1,0} convert(%mul.2144), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention_layernorm/convert_element_type" stack_frame_id=177} + %dot_general.711 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1141, %arg_tuple.5#20), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=265} + %lt.437 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/lt" stack_frame_id=329} + %add.721 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/add" stack_frame_id=329} + %select_n.94 = s32[] select(%lt.437, %add.721, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/select_n" stack_frame_id=329} + %dynamic_update_slice.68 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.105, %dot_general.711, %constant.127, %select_n.94, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/dynamic_update_slice" stack_frame_id=329} + %broadcast_in_dim.1058 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.68), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/broadcast_in_dim" stack_frame_id=338} + %reshape.486 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1058), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/reshape" stack_frame_id=338} + %convert_element_type.1136 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/convert_element_type" stack_frame_id=139} + %add.712 = f32[1,1]{1,0} reshape(%convert_element_type.1136), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_1/add"} + %ge.357 = f32[1,1]{1,0} broadcast(%add.712), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/ge" stack_frame_id=143} + %ge.358 = f32[1]{0} reshape(%ge.357), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/ge" stack_frame_id=143} + %ge.359 = f32[1,41]{1,0} broadcast(%ge.358), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/ge" stack_frame_id=143} + %iota.386 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/iota" stack_frame_id=142} + %broadcast_in_dim.1051 = f32[1,41]{1,0} reshape(%iota.386), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/broadcast_in_dim" stack_frame_id=143} + %ge.360 = pred[1,41]{1,0} compare(%ge.359, %broadcast_in_dim.1051), direction=GE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/ge" stack_frame_id=143} + %broadcast_in_dim.1052 = pred[1,1,41]{2,1,0} reshape(%ge.360), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/broadcast_in_dim" stack_frame_id=146} + %convert_element_type.1138 = s32[1,1,41]{2,1,0} convert(%broadcast_in_dim.1052), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/convert_element_type" stack_frame_id=160} + %add.713 = s32[1,1]{1,0} reshape(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/add" stack_frame_id=147} + %lt.432 = s32[1,1]{1,0} broadcast(%add.713), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/lt" stack_frame_id=152} + %lt.433 = s32[1]{0} reshape(%lt.432), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/lt" stack_frame_id=152} + %lt.434 = s32[1,41]{1,0} broadcast(%lt.433), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/lt" stack_frame_id=152} + %iota.387 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/iota" stack_frame_id=150} + %broadcast_in_dim.1053 = s32[1,41]{1,0} reshape(%iota.387), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/broadcast_in_dim" stack_frame_id=148} + %add.714 = s32[1,41]{1,0} add(%broadcast_in_dim.1053, %broadcast.43), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/add" stack_frame_id=151} + %lt.435 = pred[1,41]{1,0} compare(%lt.434, %add.714), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/lt" stack_frame_id=152} + %convert_element_type.1137 = s32[1,41]{1,0} convert(%lt.435), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/convert_element_type" stack_frame_id=156} + %broadcast_in_dim.1054 = s32[1,1,41]{2,1,0} reshape(%convert_element_type.1137), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/broadcast_in_dim" stack_frame_id=160} + %min.64 = s32[1,1,41]{2,1,0} minimum(%convert_element_type.1138, %broadcast_in_dim.1054), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/min" stack_frame_id=160} + %broadcast_in_dim.1059 = s32[1,1,1,41]{3,2,1,0} reshape(%min.64), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/broadcast_in_dim" stack_frame_id=341} + %convert_element_type.1148 = pred[1,1,1,41]{3,2,1,0} compare(%broadcast_in_dim.1059, %broadcast.37), direction=NE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/convert_element_type" stack_frame_id=345} + %broadcast_in_dim.1060 = pred[1,1,1,1,41]{4,3,2,1,0} reshape(%convert_element_type.1148), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/broadcast_in_dim" stack_frame_id=83} + %and.200 = pred[1,1,1,1,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1060), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/vmap()/and" stack_frame_id=83} + %and.201 = pred[1,1,1,41]{3,2,1,0} reshape(%and.200), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/vmap()/and" stack_frame_id=83} + %and.202 = pred[1,1,32,1,41]{4,3,2,1,0} broadcast(%and.201), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/vmap()/and" stack_frame_id=83} + %dot_general.708 = bf16[1,1,32,128]{3,2,1,0} dot(%convert_element_type.1141, %arg_tuple.5#18), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=186} + %convert_element_type.1142 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/convert_element_type" stack_frame_id=208} + %add.716 = f32[1,1]{1,0} reshape(%convert_element_type.1142), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/add"} + %iota.388 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/iota" stack_frame_id=197} + %mul.2145 = f32[64]{0} multiply(%iota.388, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=197} + %div.654 = f32[64]{0} divide(%mul.2145, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/div" stack_frame_id=198} + %neg.266 = f32[64]{0} negate(%div.654), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/neg" stack_frame_id=199} + %pow.260 = f32[64]{0} power(%broadcast.39, %neg.266), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/pow" stack_frame_id=202} + %div.655 = f32[64]{0} divide(%pow.260, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/div" stack_frame_id=203} + %dot_general.709 = f32[1,1,64]{2,1,0} dot(%add.716, %div.655), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/bi,j->bij/dot_general" stack_frame_id=211} + %stack.1035 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.709), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/stack" stack_frame_id=214} + %stack.1036 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.709), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/stack" stack_frame_id=214} + %stack.1037 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1035, %stack.1036), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/stack" stack_frame_id=214} + %reshape.481 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1037), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/reshape"} + %cos.129 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.481), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/cos" stack_frame_id=218} + %convert_element_type.1143 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.129), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/convert_element_type" stack_frame_id=222} + %mul.2146 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1143), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=242} + %mul.2147 = bf16[1,1,128]{2,1,0} reshape(%mul.2146), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=242} + %mul.2148 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2147), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=242} + %mul.2149 = bf16[1,1,32,128]{3,2,1,0} multiply(%dot_general.708, %mul.2148), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=242} + %split.259 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.708), slice={[0:1], [0:1], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/split" stack_frame_id=234} + %neg.267 = bf16[1,1,32,64]{3,2,1,0} negate(%split.259), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/neg" stack_frame_id=235} + %stack.1038 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%neg.267), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/stack" stack_frame_id=238} + %split.258 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.708), slice={[0:1], [0:1], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/split" stack_frame_id=234} + %stack.1039 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%split.258), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/stack" stack_frame_id=238} + %stack.1040 = bf16[1,1,32,2,64]{4,3,2,1,0} concatenate(%stack.1038, %stack.1039), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/stack" stack_frame_id=238} + %reshape.482 = bf16[1,1,32,128]{3,2,1,0} reshape(%stack.1040), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/reshape" stack_frame_id=241} + %sin.129 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.481), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/sin" stack_frame_id=226} + %convert_element_type.1144 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.129), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/convert_element_type" stack_frame_id=230} + %mul.2150 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1144), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=243} + %mul.2151 = bf16[1,1,128]{2,1,0} reshape(%mul.2150), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=243} + %mul.2152 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2151), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=243} + %mul.2153 = bf16[1,1,32,128]{3,2,1,0} multiply(%reshape.482, %mul.2152), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=243} + %add.717 = bf16[1,1,32,128]{3,2,1,0} add(%mul.2149, %mul.2153), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/add" stack_frame_id=244} + %reshape.487 = bf16[1,1,32,1,128]{4,3,2,1,0} reshape(%add.717), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/reshape" stack_frame_id=83} + %slice.100 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.103), slice={[0:1], [0:1], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/slice" stack_frame_id=245} + %squeeze.104 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.100), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/squeeze" stack_frame_id=245} + %dot_general.710 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1141, %arg_tuple.5#19), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=256} + %convert_element_type.1145 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/convert_element_type" stack_frame_id=287} + %add.718 = f32[1,1]{1,0} reshape(%convert_element_type.1145), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/add"} + %iota.389 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/iota" stack_frame_id=276} + %mul.2154 = f32[64]{0} multiply(%iota.389, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=276} + %div.656 = f32[64]{0} divide(%mul.2154, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/div" stack_frame_id=277} + %neg.268 = f32[64]{0} negate(%div.656), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/neg" stack_frame_id=278} + %pow.261 = f32[64]{0} power(%broadcast.39, %neg.268), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/pow" stack_frame_id=281} + %div.657 = f32[64]{0} divide(%pow.261, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/div" stack_frame_id=282} + %dot_general.712 = f32[1,1,64]{2,1,0} dot(%add.718, %div.657), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/bi,j->bij/dot_general" stack_frame_id=290} + %stack.1041 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.712), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/stack" stack_frame_id=293} + %stack.1042 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.712), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/stack" stack_frame_id=293} + %stack.1043 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1041, %stack.1042), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/stack" stack_frame_id=293} + %reshape.483 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1043), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/reshape"} + %cos.130 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.483), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/cos" stack_frame_id=297} + %convert_element_type.1146 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.130), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/convert_element_type" stack_frame_id=301} + %mul.2155 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1146), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=321} + %mul.2156 = bf16[1,1,128]{2,1,0} reshape(%mul.2155), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=321} + %mul.2157 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2156), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=321} + %mul.2158 = bf16[1,1,8,128]{3,2,1,0} multiply(%dot_general.710, %mul.2157), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=321} + %split.261 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.710), slice={[0:1], [0:1], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/split" stack_frame_id=313} + %neg.269 = bf16[1,1,8,64]{3,2,1,0} negate(%split.261), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/neg" stack_frame_id=314} + %stack.1044 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%neg.269), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/stack" stack_frame_id=317} + %split.260 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.710), slice={[0:1], [0:1], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/split" stack_frame_id=313} + %stack.1045 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%split.260), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/stack" stack_frame_id=317} + %stack.1046 = bf16[1,1,8,2,64]{4,3,2,1,0} concatenate(%stack.1044, %stack.1045), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/stack" stack_frame_id=317} + %reshape.484 = bf16[1,1,8,128]{3,2,1,0} reshape(%stack.1046), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/reshape" stack_frame_id=320} + %sin.130 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.483), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/sin" stack_frame_id=305} + %convert_element_type.1147 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.130), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/convert_element_type" stack_frame_id=309} + %mul.2159 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1147), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=322} + %mul.2160 = bf16[1,1,128]{2,1,0} reshape(%mul.2159), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=322} + %mul.2161 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2160), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=322} + %mul.2162 = bf16[1,1,8,128]{3,2,1,0} multiply(%reshape.484, %mul.2161), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=322} + %add.719 = bf16[1,1,8,128]{3,2,1,0} add(%mul.2158, %mul.2162), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/rotary_embedding_1/add" stack_frame_id=323} + %lt.436 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/lt" stack_frame_id=326} + %add.720 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/add" stack_frame_id=326} + %select_n.93 = s32[] select(%lt.436, %add.720, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/select_n" stack_frame_id=326} + %dynamic_update_slice.67 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.104, %add.719, %constant.127, %select_n.93, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/dynamic_update_slice" stack_frame_id=326} + %broadcast_in_dim.1057 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.67), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/broadcast_in_dim" stack_frame_id=335} + %reshape.485 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1057), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/reshape" stack_frame_id=335} + %dot_general.713 = f32[1,32,1,1,41]{4,3,2,1,0} dot(%reshape.487, %reshape.485), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=83} + %mul.2163 = f32[1,32,1,1,41]{4,3,2,1,0} multiply(%dot_general.713, %broadcast.36), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/vmap()/mul" stack_frame_id=83} + %vmap_jit__where__.64 = f32[1,1,32,1,41]{4,3,2,1,0} call(%and.202, %mul.2163, %constant.114), to_apply=%_where_3.137, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/vmap(jit(_where))" stack_frame_id=83} + %reduce_max.449 = f32[1,1,32,1]{3,2,1,0} reduce(%vmap_jit__where__.64, %constant.113), dimensions={4}, to_apply=%region_135.143, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/vmap()/reduce_max" stack_frame_id=83} + %max.73 = f32[1,1,32,1]{3,2,1,0} maximum(%reduce_max.449, %broadcast.35), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/vmap()/max" stack_frame_id=83} + %broadcast_in_dim.1061 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%max.73), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %sub.275 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1061), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/vmap()/sub" stack_frame_id=83} + %sub.276 = f32[1,1,32,1]{3,2,1,0} reshape(%sub.275), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/vmap()/sub" stack_frame_id=83} + %sub.277 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%sub.276), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/vmap()/sub" stack_frame_id=83} + %sub.278 = f32[1,1,32,1,41]{4,3,2,1,0} subtract(%vmap_jit__where__.64, %sub.277), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/vmap()/sub" stack_frame_id=83} + %exp.69 = f32[1,1,32,1,41]{4,3,2,1,0} exponential(%sub.278), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/vmap()/exp" stack_frame_id=83} + %reduce_sum.1355 = f32[1,1,32,1]{3,2,1,0} reduce(%exp.69, %constant.123), dimensions={4}, to_apply=%region_136.144, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/vmap()/reduce_sum" stack_frame_id=83} + %broadcast_in_dim.1062 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%reduce_sum.1355), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %div.658 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1062), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/vmap()/div" stack_frame_id=83} + %div.659 = f32[1,1,32,1]{3,2,1,0} reshape(%div.658), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/vmap()/div" stack_frame_id=83} + %div.660 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%div.659), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/vmap()/div" stack_frame_id=83} + %div.661 = f32[1,1,32,1,41]{4,3,2,1,0} divide(%exp.69, %div.660), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/vmap()/div" stack_frame_id=83} + %convert_element_type.1149 = bf16[1,1,32,1,41]{4,3,2,1,0} convert(%div.661), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/vmap()/convert_element_type" stack_frame_id=83} + %dot_general.714 = bf16[1,32,128,1,1]{4,3,2,1,0} dot(%reshape.486, %convert_element_type.1149), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=83} + %transpose.68 = bf16[1,1,32,128]{3,2,1,0} reshape(%dot_general.714), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/reshape;jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/vmap()/transpose;jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %dot_general.715 = bf16[1,1,4096]{2,1,0} dot(%transpose.68, %arg_tuple.5#21), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=354} + %add.722 = bf16[1,1,4096]{2,1,0} add(%dot_general.715, %add.711), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/add" stack_frame_id=355} + %convert_element_type.1150 = f32[1,1,4096]{2,1,0} convert(%add.722), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/feedforward_layernorm/convert_element_type" stack_frame_id=359} + %pow.262 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1150, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/feedforward_layernorm/pow" stack_frame_id=362} + %reduce_sum.1356 = f32[1,1]{1,0} reduce(%pow.262, %constant.123), dimensions={2}, to_apply=%region_137.145, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/feedforward_layernorm/reduce_sum" stack_frame_id=92} + %broadcast_in_dim.1063 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1356), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/feedforward_layernorm/broadcast_in_dim" stack_frame_id=92} + %div.662 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1063, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/feedforward_layernorm/div" stack_frame_id=92} + %add.723 = f32[1,1,1]{2,1,0} add(%div.662, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/feedforward_layernorm/add" stack_frame_id=363} + %rsqrt.131 = f32[1,1,1]{2,1,0} rsqrt(%add.723), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/feedforward_layernorm/rsqrt" stack_frame_id=366} + %mul.2164 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.131), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2165 = f32[1,1]{1,0} reshape(%mul.2164), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2166 = f32[1,1,4096]{2,1,0} broadcast(%mul.2165), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2167 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1150, %mul.2166), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/feedforward_layernorm/mul" stack_frame_id=367} + %convert_element_type.1151 = f32[4096]{0} convert(%arg_tuple.5#22), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/feedforward_layernorm/convert_element_type" stack_frame_id=368} + %broadcast_in_dim.1064 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1151), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/feedforward_layernorm/broadcast_in_dim" stack_frame_id=368} + %mul.2168 = f32[1,1,4096]{2,1,0} multiply(%mul.2167, %broadcast_in_dim.1064), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/feedforward_layernorm/mul" stack_frame_id=368} + %convert_element_type.1152 = bf16[1,1,4096]{2,1,0} convert(%mul.2168), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/feedforward_layernorm/convert_element_type" stack_frame_id=372} + %dot_general.717 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1152, %arg_tuple.5#24), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/feedforward_intermediate_dense/dot_general" stack_frame_id=401} + %dot_general.716 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1152, %arg_tuple.5#23), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/feedforward_gate_dense/dot_general" stack_frame_id=381} + %convert_element_type.1153 = f32[1,1,14336]{2,1,0} convert(%dot_general.716), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/convert_element_type" stack_frame_id=385} + %jit_silu_.64 = f32[1,1,14336]{2,1,0} call(%convert_element_type.1153), to_apply=%silu_4.141, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/jit(silu)" stack_frame_id=388} + %convert_element_type.1154 = bf16[1,1,14336]{2,1,0} convert(%jit_silu_.64), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/convert_element_type" stack_frame_id=392} + %mul.2169 = bf16[1,1,14336]{2,1,0} multiply(%dot_general.717, %convert_element_type.1154), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/mul" stack_frame_id=404} + %dot_general.718 = bf16[1,1,4096]{2,1,0} dot(%mul.2169, %arg_tuple.5#25), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/feedforward_output_dense/dot_general" stack_frame_id=413} + %add.724 = bf16[1,1,4096]{2,1,0} add(%dot_general.718, %add.722), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/add" stack_frame_id=414} + %convert_element_type.1158 = f32[1,1,4096]{2,1,0} convert(%add.724), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention_layernorm/convert_element_type" stack_frame_id=164} + %pow.263 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1158, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention_layernorm/pow" stack_frame_id=167} + %reduce_sum.1357 = f32[1,1]{1,0} reduce(%pow.263, %constant.123), dimensions={2}, to_apply=%region_138.146, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention_layernorm/reduce_sum" stack_frame_id=73} + %broadcast_in_dim.1069 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1357), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention_layernorm/broadcast_in_dim" stack_frame_id=73} + %div.663 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1069, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention_layernorm/div" stack_frame_id=73} + %add.728 = f32[1,1,1]{2,1,0} add(%div.663, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention_layernorm/add" stack_frame_id=168} + %rsqrt.132 = f32[1,1,1]{2,1,0} rsqrt(%add.728), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention_layernorm/rsqrt" stack_frame_id=171} + %mul.2170 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.132), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2171 = f32[1,1]{1,0} reshape(%mul.2170), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2172 = f32[1,1,4096]{2,1,0} broadcast(%mul.2171), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2173 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1158, %mul.2172), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention_layernorm/mul" stack_frame_id=172} + %convert_element_type.1159 = f32[4096]{0} convert(%arg_tuple.5#26), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention_layernorm/convert_element_type" stack_frame_id=173} + %broadcast_in_dim.1070 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1159), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention_layernorm/broadcast_in_dim" stack_frame_id=173} + %mul.2174 = f32[1,1,4096]{2,1,0} multiply(%mul.2173, %broadcast_in_dim.1070), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention_layernorm/mul" stack_frame_id=173} + %convert_element_type.1160 = bf16[1,1,4096]{2,1,0} convert(%mul.2174), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention_layernorm/convert_element_type" stack_frame_id=177} + %dot_general.722 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1160, %arg_tuple.5#29), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=265} + %lt.443 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/lt" stack_frame_id=329} + %add.734 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/add" stack_frame_id=329} + %select_n.96 = s32[] select(%lt.443, %add.734, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/select_n" stack_frame_id=329} + %dynamic_update_slice.70 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.108, %dot_general.722, %constant.127, %select_n.96, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/dynamic_update_slice" stack_frame_id=329} + %broadcast_in_dim.1072 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.70), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/broadcast_in_dim" stack_frame_id=338} + %reshape.493 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1072), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/reshape" stack_frame_id=338} + %convert_element_type.1155 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/convert_element_type" stack_frame_id=139} + %add.725 = f32[1,1]{1,0} reshape(%convert_element_type.1155), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_2/add"} + %ge.361 = f32[1,1]{1,0} broadcast(%add.725), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/ge" stack_frame_id=143} + %ge.362 = f32[1]{0} reshape(%ge.361), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/ge" stack_frame_id=143} + %ge.363 = f32[1,41]{1,0} broadcast(%ge.362), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/ge" stack_frame_id=143} + %iota.390 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/iota" stack_frame_id=142} + %broadcast_in_dim.1065 = f32[1,41]{1,0} reshape(%iota.390), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/broadcast_in_dim" stack_frame_id=143} + %ge.364 = pred[1,41]{1,0} compare(%ge.363, %broadcast_in_dim.1065), direction=GE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/ge" stack_frame_id=143} + %broadcast_in_dim.1066 = pred[1,1,41]{2,1,0} reshape(%ge.364), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/broadcast_in_dim" stack_frame_id=146} + %convert_element_type.1157 = s32[1,1,41]{2,1,0} convert(%broadcast_in_dim.1066), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/convert_element_type" stack_frame_id=160} + %add.726 = s32[1,1]{1,0} reshape(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/add" stack_frame_id=147} + %lt.438 = s32[1,1]{1,0} broadcast(%add.726), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/lt" stack_frame_id=152} + %lt.439 = s32[1]{0} reshape(%lt.438), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/lt" stack_frame_id=152} + %lt.440 = s32[1,41]{1,0} broadcast(%lt.439), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/lt" stack_frame_id=152} + %iota.391 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/iota" stack_frame_id=150} + %broadcast_in_dim.1067 = s32[1,41]{1,0} reshape(%iota.391), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/broadcast_in_dim" stack_frame_id=148} + %add.727 = s32[1,41]{1,0} add(%broadcast_in_dim.1067, %broadcast.43), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/add" stack_frame_id=151} + %lt.441 = pred[1,41]{1,0} compare(%lt.440, %add.727), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/lt" stack_frame_id=152} + %convert_element_type.1156 = s32[1,41]{1,0} convert(%lt.441), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/convert_element_type" stack_frame_id=156} + %broadcast_in_dim.1068 = s32[1,1,41]{2,1,0} reshape(%convert_element_type.1156), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/broadcast_in_dim" stack_frame_id=160} + %min.65 = s32[1,1,41]{2,1,0} minimum(%convert_element_type.1157, %broadcast_in_dim.1068), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/min" stack_frame_id=160} + %broadcast_in_dim.1073 = s32[1,1,1,41]{3,2,1,0} reshape(%min.65), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/broadcast_in_dim" stack_frame_id=341} + %convert_element_type.1167 = pred[1,1,1,41]{3,2,1,0} compare(%broadcast_in_dim.1073, %broadcast.37), direction=NE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/convert_element_type" stack_frame_id=345} + %broadcast_in_dim.1074 = pred[1,1,1,1,41]{4,3,2,1,0} reshape(%convert_element_type.1167), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/broadcast_in_dim" stack_frame_id=83} + %and.203 = pred[1,1,1,1,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1074), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/vmap()/and" stack_frame_id=83} + %and.204 = pred[1,1,1,41]{3,2,1,0} reshape(%and.203), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/vmap()/and" stack_frame_id=83} + %and.205 = pred[1,1,32,1,41]{4,3,2,1,0} broadcast(%and.204), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/vmap()/and" stack_frame_id=83} + %dot_general.719 = bf16[1,1,32,128]{3,2,1,0} dot(%convert_element_type.1160, %arg_tuple.5#27), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=186} + %convert_element_type.1161 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/convert_element_type" stack_frame_id=208} + %add.729 = f32[1,1]{1,0} reshape(%convert_element_type.1161), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/add"} + %iota.392 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/iota" stack_frame_id=197} + %mul.2175 = f32[64]{0} multiply(%iota.392, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=197} + %div.664 = f32[64]{0} divide(%mul.2175, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/div" stack_frame_id=198} + %neg.270 = f32[64]{0} negate(%div.664), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/neg" stack_frame_id=199} + %pow.264 = f32[64]{0} power(%broadcast.39, %neg.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/pow" stack_frame_id=202} + %div.665 = f32[64]{0} divide(%pow.264, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/div" stack_frame_id=203} + %dot_general.720 = f32[1,1,64]{2,1,0} dot(%add.729, %div.665), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/bi,j->bij/dot_general" stack_frame_id=211} + %stack.1050 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.720), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/stack" stack_frame_id=214} + %stack.1051 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.720), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/stack" stack_frame_id=214} + %stack.1052 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1050, %stack.1051), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/stack" stack_frame_id=214} + %reshape.488 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1052), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/reshape"} + %cos.131 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.488), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/cos" stack_frame_id=218} + %convert_element_type.1162 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.131), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/convert_element_type" stack_frame_id=222} + %mul.2176 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1162), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=242} + %mul.2177 = bf16[1,1,128]{2,1,0} reshape(%mul.2176), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=242} + %mul.2178 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2177), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=242} + %mul.2179 = bf16[1,1,32,128]{3,2,1,0} multiply(%dot_general.719, %mul.2178), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=242} + %split.263 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.719), slice={[0:1], [0:1], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/split" stack_frame_id=234} + %neg.271 = bf16[1,1,32,64]{3,2,1,0} negate(%split.263), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/neg" stack_frame_id=235} + %stack.1053 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%neg.271), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/stack" stack_frame_id=238} + %split.262 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.719), slice={[0:1], [0:1], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/split" stack_frame_id=234} + %stack.1054 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%split.262), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/stack" stack_frame_id=238} + %stack.1055 = bf16[1,1,32,2,64]{4,3,2,1,0} concatenate(%stack.1053, %stack.1054), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/stack" stack_frame_id=238} + %reshape.489 = bf16[1,1,32,128]{3,2,1,0} reshape(%stack.1055), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/reshape" stack_frame_id=241} + %sin.131 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.488), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/sin" stack_frame_id=226} + %convert_element_type.1163 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.131), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/convert_element_type" stack_frame_id=230} + %mul.2180 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1163), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=243} + %mul.2181 = bf16[1,1,128]{2,1,0} reshape(%mul.2180), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=243} + %mul.2182 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2181), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=243} + %mul.2183 = bf16[1,1,32,128]{3,2,1,0} multiply(%reshape.489, %mul.2182), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=243} + %add.730 = bf16[1,1,32,128]{3,2,1,0} add(%mul.2179, %mul.2183), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/add" stack_frame_id=244} + %reshape.494 = bf16[1,1,32,1,128]{4,3,2,1,0} reshape(%add.730), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/reshape" stack_frame_id=83} + %slice.103 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.106), slice={[0:1], [0:1], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/slice" stack_frame_id=245} + %squeeze.107 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.103), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/squeeze" stack_frame_id=245} + %dot_general.721 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1160, %arg_tuple.5#28), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=256} + %convert_element_type.1164 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/convert_element_type" stack_frame_id=287} + %add.731 = f32[1,1]{1,0} reshape(%convert_element_type.1164), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/add"} + %iota.393 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/iota" stack_frame_id=276} + %mul.2184 = f32[64]{0} multiply(%iota.393, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=276} + %div.666 = f32[64]{0} divide(%mul.2184, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/div" stack_frame_id=277} + %neg.272 = f32[64]{0} negate(%div.666), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/neg" stack_frame_id=278} + %pow.265 = f32[64]{0} power(%broadcast.39, %neg.272), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/pow" stack_frame_id=281} + %div.667 = f32[64]{0} divide(%pow.265, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/div" stack_frame_id=282} + %dot_general.723 = f32[1,1,64]{2,1,0} dot(%add.731, %div.667), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/bi,j->bij/dot_general" stack_frame_id=290} + %stack.1056 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.723), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/stack" stack_frame_id=293} + %stack.1057 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.723), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/stack" stack_frame_id=293} + %stack.1058 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1056, %stack.1057), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/stack" stack_frame_id=293} + %reshape.490 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1058), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/reshape"} + %cos.132 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.490), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/cos" stack_frame_id=297} + %convert_element_type.1165 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.132), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/convert_element_type" stack_frame_id=301} + %mul.2185 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1165), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=321} + %mul.2186 = bf16[1,1,128]{2,1,0} reshape(%mul.2185), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=321} + %mul.2187 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2186), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=321} + %mul.2188 = bf16[1,1,8,128]{3,2,1,0} multiply(%dot_general.721, %mul.2187), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=321} + %split.265 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.721), slice={[0:1], [0:1], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/split" stack_frame_id=313} + %neg.273 = bf16[1,1,8,64]{3,2,1,0} negate(%split.265), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/neg" stack_frame_id=314} + %stack.1059 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%neg.273), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/stack" stack_frame_id=317} + %split.264 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.721), slice={[0:1], [0:1], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/split" stack_frame_id=313} + %stack.1060 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%split.264), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/stack" stack_frame_id=317} + %stack.1061 = bf16[1,1,8,2,64]{4,3,2,1,0} concatenate(%stack.1059, %stack.1060), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/stack" stack_frame_id=317} + %reshape.491 = bf16[1,1,8,128]{3,2,1,0} reshape(%stack.1061), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/reshape" stack_frame_id=320} + %sin.132 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.490), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/sin" stack_frame_id=305} + %convert_element_type.1166 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.132), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/convert_element_type" stack_frame_id=309} + %mul.2189 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1166), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=322} + %mul.2190 = bf16[1,1,128]{2,1,0} reshape(%mul.2189), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=322} + %mul.2191 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2190), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=322} + %mul.2192 = bf16[1,1,8,128]{3,2,1,0} multiply(%reshape.491, %mul.2191), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=322} + %add.732 = bf16[1,1,8,128]{3,2,1,0} add(%mul.2188, %mul.2192), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/rotary_embedding_2/add" stack_frame_id=323} + %lt.442 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/lt" stack_frame_id=326} + %add.733 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/add" stack_frame_id=326} + %select_n.95 = s32[] select(%lt.442, %add.733, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/select_n" stack_frame_id=326} + %dynamic_update_slice.69 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.107, %add.732, %constant.127, %select_n.95, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/dynamic_update_slice" stack_frame_id=326} + %broadcast_in_dim.1071 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.69), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/broadcast_in_dim" stack_frame_id=335} + %reshape.492 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1071), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/reshape" stack_frame_id=335} + %dot_general.724 = f32[1,32,1,1,41]{4,3,2,1,0} dot(%reshape.494, %reshape.492), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=83} + %mul.2193 = f32[1,32,1,1,41]{4,3,2,1,0} multiply(%dot_general.724, %broadcast.36), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/vmap()/mul" stack_frame_id=83} + %vmap_jit__where__.65 = f32[1,1,32,1,41]{4,3,2,1,0} call(%and.205, %mul.2193, %constant.114), to_apply=%_where_3.137, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/vmap(jit(_where))" stack_frame_id=83} + %reduce_max.450 = f32[1,1,32,1]{3,2,1,0} reduce(%vmap_jit__where__.65, %constant.113), dimensions={4}, to_apply=%region_139.147, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/vmap()/reduce_max" stack_frame_id=83} + %max.74 = f32[1,1,32,1]{3,2,1,0} maximum(%reduce_max.450, %broadcast.35), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/vmap()/max" stack_frame_id=83} + %broadcast_in_dim.1075 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%max.74), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %sub.279 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1075), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/vmap()/sub" stack_frame_id=83} + %sub.280 = f32[1,1,32,1]{3,2,1,0} reshape(%sub.279), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/vmap()/sub" stack_frame_id=83} + %sub.281 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%sub.280), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/vmap()/sub" stack_frame_id=83} + %sub.282 = f32[1,1,32,1,41]{4,3,2,1,0} subtract(%vmap_jit__where__.65, %sub.281), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/vmap()/sub" stack_frame_id=83} + %exp.70 = f32[1,1,32,1,41]{4,3,2,1,0} exponential(%sub.282), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/vmap()/exp" stack_frame_id=83} + %reduce_sum.1358 = f32[1,1,32,1]{3,2,1,0} reduce(%exp.70, %constant.123), dimensions={4}, to_apply=%region_140.148, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/vmap()/reduce_sum" stack_frame_id=83} + %broadcast_in_dim.1076 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%reduce_sum.1358), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %div.668 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1076), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/vmap()/div" stack_frame_id=83} + %div.669 = f32[1,1,32,1]{3,2,1,0} reshape(%div.668), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/vmap()/div" stack_frame_id=83} + %div.670 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%div.669), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/vmap()/div" stack_frame_id=83} + %div.671 = f32[1,1,32,1,41]{4,3,2,1,0} divide(%exp.70, %div.670), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/vmap()/div" stack_frame_id=83} + %convert_element_type.1168 = bf16[1,1,32,1,41]{4,3,2,1,0} convert(%div.671), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/vmap()/convert_element_type" stack_frame_id=83} + %dot_general.725 = bf16[1,32,128,1,1]{4,3,2,1,0} dot(%reshape.493, %convert_element_type.1168), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=83} + %transpose.69 = bf16[1,1,32,128]{3,2,1,0} reshape(%dot_general.725), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/reshape;jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/vmap()/transpose;jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %dot_general.726 = bf16[1,1,4096]{2,1,0} dot(%transpose.69, %arg_tuple.5#30), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=354} + %add.735 = bf16[1,1,4096]{2,1,0} add(%dot_general.726, %add.724), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/add" stack_frame_id=355} + %convert_element_type.1169 = f32[1,1,4096]{2,1,0} convert(%add.735), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/feedforward_layernorm/convert_element_type" stack_frame_id=359} + %pow.266 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1169, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/feedforward_layernorm/pow" stack_frame_id=362} + %reduce_sum.1359 = f32[1,1]{1,0} reduce(%pow.266, %constant.123), dimensions={2}, to_apply=%region_141.149, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/feedforward_layernorm/reduce_sum" stack_frame_id=92} + %broadcast_in_dim.1077 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1359), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/feedforward_layernorm/broadcast_in_dim" stack_frame_id=92} + %div.672 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1077, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/feedforward_layernorm/div" stack_frame_id=92} + %add.736 = f32[1,1,1]{2,1,0} add(%div.672, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/feedforward_layernorm/add" stack_frame_id=363} + %rsqrt.133 = f32[1,1,1]{2,1,0} rsqrt(%add.736), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/feedforward_layernorm/rsqrt" stack_frame_id=366} + %mul.2194 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.133), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2195 = f32[1,1]{1,0} reshape(%mul.2194), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2196 = f32[1,1,4096]{2,1,0} broadcast(%mul.2195), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2197 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1169, %mul.2196), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/feedforward_layernorm/mul" stack_frame_id=367} + %convert_element_type.1170 = f32[4096]{0} convert(%arg_tuple.5#31), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/feedforward_layernorm/convert_element_type" stack_frame_id=368} + %broadcast_in_dim.1078 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1170), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/feedforward_layernorm/broadcast_in_dim" stack_frame_id=368} + %mul.2198 = f32[1,1,4096]{2,1,0} multiply(%mul.2197, %broadcast_in_dim.1078), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/feedforward_layernorm/mul" stack_frame_id=368} + %convert_element_type.1171 = bf16[1,1,4096]{2,1,0} convert(%mul.2198), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/feedforward_layernorm/convert_element_type" stack_frame_id=372} + %dot_general.728 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1171, %arg_tuple.5#33), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/feedforward_intermediate_dense/dot_general" stack_frame_id=401} + %dot_general.727 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1171, %arg_tuple.5#32), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/feedforward_gate_dense/dot_general" stack_frame_id=381} + %convert_element_type.1172 = f32[1,1,14336]{2,1,0} convert(%dot_general.727), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/convert_element_type" stack_frame_id=385} + %jit_silu_.65 = f32[1,1,14336]{2,1,0} call(%convert_element_type.1172), to_apply=%silu_4.141, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/jit(silu)" stack_frame_id=388} + %convert_element_type.1173 = bf16[1,1,14336]{2,1,0} convert(%jit_silu_.65), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/convert_element_type" stack_frame_id=392} + %mul.2199 = bf16[1,1,14336]{2,1,0} multiply(%dot_general.728, %convert_element_type.1173), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/mul" stack_frame_id=404} + %dot_general.729 = bf16[1,1,4096]{2,1,0} dot(%mul.2199, %arg_tuple.5#34), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/feedforward_output_dense/dot_general" stack_frame_id=413} + %add.737 = bf16[1,1,4096]{2,1,0} add(%dot_general.729, %add.735), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/add" stack_frame_id=414} + %convert_element_type.1177 = f32[1,1,4096]{2,1,0} convert(%add.737), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention_layernorm/convert_element_type" stack_frame_id=164} + %pow.267 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1177, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention_layernorm/pow" stack_frame_id=167} + %reduce_sum.1360 = f32[1,1]{1,0} reduce(%pow.267, %constant.123), dimensions={2}, to_apply=%region_142.150, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention_layernorm/reduce_sum" stack_frame_id=73} + %broadcast_in_dim.1083 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1360), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention_layernorm/broadcast_in_dim" stack_frame_id=73} + %div.673 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1083, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention_layernorm/div" stack_frame_id=73} + %add.741 = f32[1,1,1]{2,1,0} add(%div.673, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention_layernorm/add" stack_frame_id=168} + %rsqrt.134 = f32[1,1,1]{2,1,0} rsqrt(%add.741), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention_layernorm/rsqrt" stack_frame_id=171} + %mul.2200 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.134), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2201 = f32[1,1]{1,0} reshape(%mul.2200), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2202 = f32[1,1,4096]{2,1,0} broadcast(%mul.2201), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2203 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1177, %mul.2202), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention_layernorm/mul" stack_frame_id=172} + %convert_element_type.1178 = f32[4096]{0} convert(%arg_tuple.5#35), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention_layernorm/convert_element_type" stack_frame_id=173} + %broadcast_in_dim.1084 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1178), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention_layernorm/broadcast_in_dim" stack_frame_id=173} + %mul.2204 = f32[1,1,4096]{2,1,0} multiply(%mul.2203, %broadcast_in_dim.1084), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention_layernorm/mul" stack_frame_id=173} + %convert_element_type.1179 = bf16[1,1,4096]{2,1,0} convert(%mul.2204), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention_layernorm/convert_element_type" stack_frame_id=177} + %dot_general.733 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1179, %arg_tuple.5#38), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=265} + %lt.449 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/lt" stack_frame_id=329} + %add.747 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/add" stack_frame_id=329} + %select_n.98 = s32[] select(%lt.449, %add.747, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/select_n" stack_frame_id=329} + %dynamic_update_slice.72 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.111, %dot_general.733, %constant.127, %select_n.98, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/dynamic_update_slice" stack_frame_id=329} + %broadcast_in_dim.1086 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.72), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/broadcast_in_dim" stack_frame_id=338} + %reshape.500 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1086), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/reshape" stack_frame_id=338} + %convert_element_type.1174 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/convert_element_type" stack_frame_id=139} + %add.738 = f32[1,1]{1,0} reshape(%convert_element_type.1174), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_3/add"} + %ge.365 = f32[1,1]{1,0} broadcast(%add.738), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/ge" stack_frame_id=143} + %ge.366 = f32[1]{0} reshape(%ge.365), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/ge" stack_frame_id=143} + %ge.367 = f32[1,41]{1,0} broadcast(%ge.366), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/ge" stack_frame_id=143} + %iota.394 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/iota" stack_frame_id=142} + %broadcast_in_dim.1079 = f32[1,41]{1,0} reshape(%iota.394), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/broadcast_in_dim" stack_frame_id=143} + %ge.368 = pred[1,41]{1,0} compare(%ge.367, %broadcast_in_dim.1079), direction=GE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/ge" stack_frame_id=143} + %broadcast_in_dim.1080 = pred[1,1,41]{2,1,0} reshape(%ge.368), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/broadcast_in_dim" stack_frame_id=146} + %convert_element_type.1176 = s32[1,1,41]{2,1,0} convert(%broadcast_in_dim.1080), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/convert_element_type" stack_frame_id=160} + %add.739 = s32[1,1]{1,0} reshape(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/add" stack_frame_id=147} + %lt.444 = s32[1,1]{1,0} broadcast(%add.739), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/lt" stack_frame_id=152} + %lt.445 = s32[1]{0} reshape(%lt.444), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/lt" stack_frame_id=152} + %lt.446 = s32[1,41]{1,0} broadcast(%lt.445), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/lt" stack_frame_id=152} + %iota.395 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/iota" stack_frame_id=150} + %broadcast_in_dim.1081 = s32[1,41]{1,0} reshape(%iota.395), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/broadcast_in_dim" stack_frame_id=148} + %add.740 = s32[1,41]{1,0} add(%broadcast_in_dim.1081, %broadcast.43), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/add" stack_frame_id=151} + %lt.447 = pred[1,41]{1,0} compare(%lt.446, %add.740), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/lt" stack_frame_id=152} + %convert_element_type.1175 = s32[1,41]{1,0} convert(%lt.447), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/convert_element_type" stack_frame_id=156} + %broadcast_in_dim.1082 = s32[1,1,41]{2,1,0} reshape(%convert_element_type.1175), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/broadcast_in_dim" stack_frame_id=160} + %min.66 = s32[1,1,41]{2,1,0} minimum(%convert_element_type.1176, %broadcast_in_dim.1082), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/min" stack_frame_id=160} + %broadcast_in_dim.1087 = s32[1,1,1,41]{3,2,1,0} reshape(%min.66), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/broadcast_in_dim" stack_frame_id=341} + %convert_element_type.1186 = pred[1,1,1,41]{3,2,1,0} compare(%broadcast_in_dim.1087, %broadcast.37), direction=NE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/convert_element_type" stack_frame_id=345} + %broadcast_in_dim.1088 = pred[1,1,1,1,41]{4,3,2,1,0} reshape(%convert_element_type.1186), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/broadcast_in_dim" stack_frame_id=83} + %and.206 = pred[1,1,1,1,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1088), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/vmap()/and" stack_frame_id=83} + %and.207 = pred[1,1,1,41]{3,2,1,0} reshape(%and.206), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/vmap()/and" stack_frame_id=83} + %and.208 = pred[1,1,32,1,41]{4,3,2,1,0} broadcast(%and.207), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/vmap()/and" stack_frame_id=83} + %dot_general.730 = bf16[1,1,32,128]{3,2,1,0} dot(%convert_element_type.1179, %arg_tuple.5#36), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=186} + %convert_element_type.1180 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/convert_element_type" stack_frame_id=208} + %add.742 = f32[1,1]{1,0} reshape(%convert_element_type.1180), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/add"} + %iota.396 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/iota" stack_frame_id=197} + %mul.2205 = f32[64]{0} multiply(%iota.396, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=197} + %div.674 = f32[64]{0} divide(%mul.2205, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/div" stack_frame_id=198} + %neg.274 = f32[64]{0} negate(%div.674), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/neg" stack_frame_id=199} + %pow.268 = f32[64]{0} power(%broadcast.39, %neg.274), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/pow" stack_frame_id=202} + %div.675 = f32[64]{0} divide(%pow.268, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/div" stack_frame_id=203} + %dot_general.731 = f32[1,1,64]{2,1,0} dot(%add.742, %div.675), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/bi,j->bij/dot_general" stack_frame_id=211} + %stack.1065 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.731), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/stack" stack_frame_id=214} + %stack.1066 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.731), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/stack" stack_frame_id=214} + %stack.1067 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1065, %stack.1066), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/stack" stack_frame_id=214} + %reshape.495 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1067), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/reshape"} + %cos.133 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.495), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/cos" stack_frame_id=218} + %convert_element_type.1181 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.133), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/convert_element_type" stack_frame_id=222} + %mul.2206 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1181), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=242} + %mul.2207 = bf16[1,1,128]{2,1,0} reshape(%mul.2206), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=242} + %mul.2208 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2207), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=242} + %mul.2209 = bf16[1,1,32,128]{3,2,1,0} multiply(%dot_general.730, %mul.2208), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=242} + %split.267 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.730), slice={[0:1], [0:1], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/split" stack_frame_id=234} + %neg.275 = bf16[1,1,32,64]{3,2,1,0} negate(%split.267), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/neg" stack_frame_id=235} + %stack.1068 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%neg.275), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/stack" stack_frame_id=238} + %split.266 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.730), slice={[0:1], [0:1], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/split" stack_frame_id=234} + %stack.1069 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%split.266), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/stack" stack_frame_id=238} + %stack.1070 = bf16[1,1,32,2,64]{4,3,2,1,0} concatenate(%stack.1068, %stack.1069), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/stack" stack_frame_id=238} + %reshape.496 = bf16[1,1,32,128]{3,2,1,0} reshape(%stack.1070), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/reshape" stack_frame_id=241} + %sin.133 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.495), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/sin" stack_frame_id=226} + %convert_element_type.1182 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.133), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/convert_element_type" stack_frame_id=230} + %mul.2210 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1182), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=243} + %mul.2211 = bf16[1,1,128]{2,1,0} reshape(%mul.2210), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=243} + %mul.2212 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2211), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=243} + %mul.2213 = bf16[1,1,32,128]{3,2,1,0} multiply(%reshape.496, %mul.2212), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=243} + %add.743 = bf16[1,1,32,128]{3,2,1,0} add(%mul.2209, %mul.2213), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/add" stack_frame_id=244} + %reshape.501 = bf16[1,1,32,1,128]{4,3,2,1,0} reshape(%add.743), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/reshape" stack_frame_id=83} + %slice.106 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.109), slice={[0:1], [0:1], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/slice" stack_frame_id=245} + %squeeze.110 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.106), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/squeeze" stack_frame_id=245} + %dot_general.732 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1179, %arg_tuple.5#37), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=256} + %convert_element_type.1183 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/convert_element_type" stack_frame_id=287} + %add.744 = f32[1,1]{1,0} reshape(%convert_element_type.1183), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/add"} + %iota.397 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/iota" stack_frame_id=276} + %mul.2214 = f32[64]{0} multiply(%iota.397, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=276} + %div.676 = f32[64]{0} divide(%mul.2214, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/div" stack_frame_id=277} + %neg.276 = f32[64]{0} negate(%div.676), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/neg" stack_frame_id=278} + %pow.269 = f32[64]{0} power(%broadcast.39, %neg.276), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/pow" stack_frame_id=281} + %div.677 = f32[64]{0} divide(%pow.269, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/div" stack_frame_id=282} + %dot_general.734 = f32[1,1,64]{2,1,0} dot(%add.744, %div.677), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/bi,j->bij/dot_general" stack_frame_id=290} + %stack.1071 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.734), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/stack" stack_frame_id=293} + %stack.1072 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.734), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/stack" stack_frame_id=293} + %stack.1073 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1071, %stack.1072), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/stack" stack_frame_id=293} + %reshape.497 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1073), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/reshape"} + %cos.134 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.497), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/cos" stack_frame_id=297} + %convert_element_type.1184 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.134), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/convert_element_type" stack_frame_id=301} + %mul.2215 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1184), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=321} + %mul.2216 = bf16[1,1,128]{2,1,0} reshape(%mul.2215), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=321} + %mul.2217 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2216), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=321} + %mul.2218 = bf16[1,1,8,128]{3,2,1,0} multiply(%dot_general.732, %mul.2217), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=321} + %split.269 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.732), slice={[0:1], [0:1], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/split" stack_frame_id=313} + %neg.277 = bf16[1,1,8,64]{3,2,1,0} negate(%split.269), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/neg" stack_frame_id=314} + %stack.1074 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%neg.277), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/stack" stack_frame_id=317} + %split.268 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.732), slice={[0:1], [0:1], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/split" stack_frame_id=313} + %stack.1075 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%split.268), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/stack" stack_frame_id=317} + %stack.1076 = bf16[1,1,8,2,64]{4,3,2,1,0} concatenate(%stack.1074, %stack.1075), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/stack" stack_frame_id=317} + %reshape.498 = bf16[1,1,8,128]{3,2,1,0} reshape(%stack.1076), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/reshape" stack_frame_id=320} + %sin.134 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.497), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/sin" stack_frame_id=305} + %convert_element_type.1185 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.134), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/convert_element_type" stack_frame_id=309} + %mul.2219 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1185), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=322} + %mul.2220 = bf16[1,1,128]{2,1,0} reshape(%mul.2219), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=322} + %mul.2221 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2220), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=322} + %mul.2222 = bf16[1,1,8,128]{3,2,1,0} multiply(%reshape.498, %mul.2221), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=322} + %add.745 = bf16[1,1,8,128]{3,2,1,0} add(%mul.2218, %mul.2222), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/rotary_embedding_3/add" stack_frame_id=323} + %lt.448 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/lt" stack_frame_id=326} + %add.746 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/add" stack_frame_id=326} + %select_n.97 = s32[] select(%lt.448, %add.746, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/select_n" stack_frame_id=326} + %dynamic_update_slice.71 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.110, %add.745, %constant.127, %select_n.97, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/dynamic_update_slice" stack_frame_id=326} + %broadcast_in_dim.1085 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.71), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/broadcast_in_dim" stack_frame_id=335} + %reshape.499 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1085), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/reshape" stack_frame_id=335} + %dot_general.735 = f32[1,32,1,1,41]{4,3,2,1,0} dot(%reshape.501, %reshape.499), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=83} + %mul.2223 = f32[1,32,1,1,41]{4,3,2,1,0} multiply(%dot_general.735, %broadcast.36), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/vmap()/mul" stack_frame_id=83} + %vmap_jit__where__.66 = f32[1,1,32,1,41]{4,3,2,1,0} call(%and.208, %mul.2223, %constant.114), to_apply=%_where_3.137, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/vmap(jit(_where))" stack_frame_id=83} + %reduce_max.451 = f32[1,1,32,1]{3,2,1,0} reduce(%vmap_jit__where__.66, %constant.113), dimensions={4}, to_apply=%region_143.151, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/vmap()/reduce_max" stack_frame_id=83} + %max.75 = f32[1,1,32,1]{3,2,1,0} maximum(%reduce_max.451, %broadcast.35), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/vmap()/max" stack_frame_id=83} + %broadcast_in_dim.1089 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%max.75), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %sub.283 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1089), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/vmap()/sub" stack_frame_id=83} + %sub.284 = f32[1,1,32,1]{3,2,1,0} reshape(%sub.283), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/vmap()/sub" stack_frame_id=83} + %sub.285 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%sub.284), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/vmap()/sub" stack_frame_id=83} + %sub.286 = f32[1,1,32,1,41]{4,3,2,1,0} subtract(%vmap_jit__where__.66, %sub.285), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/vmap()/sub" stack_frame_id=83} + %exp.71 = f32[1,1,32,1,41]{4,3,2,1,0} exponential(%sub.286), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/vmap()/exp" stack_frame_id=83} + %reduce_sum.1361 = f32[1,1,32,1]{3,2,1,0} reduce(%exp.71, %constant.123), dimensions={4}, to_apply=%region_144.152, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/vmap()/reduce_sum" stack_frame_id=83} + %broadcast_in_dim.1090 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%reduce_sum.1361), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %div.678 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1090), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/vmap()/div" stack_frame_id=83} + %div.679 = f32[1,1,32,1]{3,2,1,0} reshape(%div.678), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/vmap()/div" stack_frame_id=83} + %div.680 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%div.679), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/vmap()/div" stack_frame_id=83} + %div.681 = f32[1,1,32,1,41]{4,3,2,1,0} divide(%exp.71, %div.680), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/vmap()/div" stack_frame_id=83} + %convert_element_type.1187 = bf16[1,1,32,1,41]{4,3,2,1,0} convert(%div.681), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/vmap()/convert_element_type" stack_frame_id=83} + %dot_general.736 = bf16[1,32,128,1,1]{4,3,2,1,0} dot(%reshape.500, %convert_element_type.1187), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=83} + %transpose.70 = bf16[1,1,32,128]{3,2,1,0} reshape(%dot_general.736), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/reshape;jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/vmap()/transpose;jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %dot_general.737 = bf16[1,1,4096]{2,1,0} dot(%transpose.70, %arg_tuple.5#39), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=354} + %add.748 = bf16[1,1,4096]{2,1,0} add(%dot_general.737, %add.737), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/add" stack_frame_id=355} + %convert_element_type.1188 = f32[1,1,4096]{2,1,0} convert(%add.748), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/feedforward_layernorm/convert_element_type" stack_frame_id=359} + %pow.270 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1188, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/feedforward_layernorm/pow" stack_frame_id=362} + %reduce_sum.1362 = f32[1,1]{1,0} reduce(%pow.270, %constant.123), dimensions={2}, to_apply=%region_145.153, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/feedforward_layernorm/reduce_sum" stack_frame_id=92} + %broadcast_in_dim.1091 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1362), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/feedforward_layernorm/broadcast_in_dim" stack_frame_id=92} + %div.682 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1091, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/feedforward_layernorm/div" stack_frame_id=92} + %add.749 = f32[1,1,1]{2,1,0} add(%div.682, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/feedforward_layernorm/add" stack_frame_id=363} + %rsqrt.135 = f32[1,1,1]{2,1,0} rsqrt(%add.749), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/feedforward_layernorm/rsqrt" stack_frame_id=366} + %mul.2224 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.135), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2225 = f32[1,1]{1,0} reshape(%mul.2224), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2226 = f32[1,1,4096]{2,1,0} broadcast(%mul.2225), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2227 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1188, %mul.2226), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/feedforward_layernorm/mul" stack_frame_id=367} + %convert_element_type.1189 = f32[4096]{0} convert(%arg_tuple.5#40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/feedforward_layernorm/convert_element_type" stack_frame_id=368} + %broadcast_in_dim.1092 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1189), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/feedforward_layernorm/broadcast_in_dim" stack_frame_id=368} + %mul.2228 = f32[1,1,4096]{2,1,0} multiply(%mul.2227, %broadcast_in_dim.1092), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/feedforward_layernorm/mul" stack_frame_id=368} + %convert_element_type.1190 = bf16[1,1,4096]{2,1,0} convert(%mul.2228), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/feedforward_layernorm/convert_element_type" stack_frame_id=372} + %dot_general.739 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1190, %arg_tuple.5#42), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/feedforward_intermediate_dense/dot_general" stack_frame_id=401} + %dot_general.738 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1190, %arg_tuple.5#41), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/feedforward_gate_dense/dot_general" stack_frame_id=381} + %convert_element_type.1191 = f32[1,1,14336]{2,1,0} convert(%dot_general.738), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/convert_element_type" stack_frame_id=385} + %jit_silu_.66 = f32[1,1,14336]{2,1,0} call(%convert_element_type.1191), to_apply=%silu_4.141, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/jit(silu)" stack_frame_id=388} + %convert_element_type.1192 = bf16[1,1,14336]{2,1,0} convert(%jit_silu_.66), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/convert_element_type" stack_frame_id=392} + %mul.2229 = bf16[1,1,14336]{2,1,0} multiply(%dot_general.739, %convert_element_type.1192), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/mul" stack_frame_id=404} + %dot_general.740 = bf16[1,1,4096]{2,1,0} dot(%mul.2229, %arg_tuple.5#43), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/feedforward_output_dense/dot_general" stack_frame_id=413} + %add.750 = bf16[1,1,4096]{2,1,0} add(%dot_general.740, %add.748), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/add" stack_frame_id=414} + %convert_element_type.1196 = f32[1,1,4096]{2,1,0} convert(%add.750), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention_layernorm/convert_element_type" stack_frame_id=164} + %pow.271 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1196, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention_layernorm/pow" stack_frame_id=167} + %reduce_sum.1363 = f32[1,1]{1,0} reduce(%pow.271, %constant.123), dimensions={2}, to_apply=%region_146.154, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention_layernorm/reduce_sum" stack_frame_id=73} + %broadcast_in_dim.1097 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1363), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention_layernorm/broadcast_in_dim" stack_frame_id=73} + %div.683 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1097, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention_layernorm/div" stack_frame_id=73} + %add.754 = f32[1,1,1]{2,1,0} add(%div.683, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention_layernorm/add" stack_frame_id=168} + %rsqrt.136 = f32[1,1,1]{2,1,0} rsqrt(%add.754), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention_layernorm/rsqrt" stack_frame_id=171} + %mul.2230 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.136), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2231 = f32[1,1]{1,0} reshape(%mul.2230), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2232 = f32[1,1,4096]{2,1,0} broadcast(%mul.2231), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2233 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1196, %mul.2232), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention_layernorm/mul" stack_frame_id=172} + %convert_element_type.1197 = f32[4096]{0} convert(%arg_tuple.5#44), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention_layernorm/convert_element_type" stack_frame_id=173} + %broadcast_in_dim.1098 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1197), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention_layernorm/broadcast_in_dim" stack_frame_id=173} + %mul.2234 = f32[1,1,4096]{2,1,0} multiply(%mul.2233, %broadcast_in_dim.1098), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention_layernorm/mul" stack_frame_id=173} + %convert_element_type.1198 = bf16[1,1,4096]{2,1,0} convert(%mul.2234), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention_layernorm/convert_element_type" stack_frame_id=177} + %dot_general.744 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1198, %arg_tuple.5#47), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=265} + %lt.455 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/lt" stack_frame_id=329} + %add.760 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/add" stack_frame_id=329} + %select_n.100 = s32[] select(%lt.455, %add.760, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/select_n" stack_frame_id=329} + %dynamic_update_slice.74 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.114, %dot_general.744, %constant.127, %select_n.100, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/dynamic_update_slice" stack_frame_id=329} + %broadcast_in_dim.1100 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.74), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/broadcast_in_dim" stack_frame_id=338} + %reshape.507 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1100), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/reshape" stack_frame_id=338} + %convert_element_type.1193 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/convert_element_type" stack_frame_id=139} + %add.751 = f32[1,1]{1,0} reshape(%convert_element_type.1193), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_4/add"} + %ge.369 = f32[1,1]{1,0} broadcast(%add.751), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/ge" stack_frame_id=143} + %ge.370 = f32[1]{0} reshape(%ge.369), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/ge" stack_frame_id=143} + %ge.371 = f32[1,41]{1,0} broadcast(%ge.370), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/ge" stack_frame_id=143} + %iota.398 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/iota" stack_frame_id=142} + %broadcast_in_dim.1093 = f32[1,41]{1,0} reshape(%iota.398), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/broadcast_in_dim" stack_frame_id=143} + %ge.372 = pred[1,41]{1,0} compare(%ge.371, %broadcast_in_dim.1093), direction=GE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/ge" stack_frame_id=143} + %broadcast_in_dim.1094 = pred[1,1,41]{2,1,0} reshape(%ge.372), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/broadcast_in_dim" stack_frame_id=146} + %convert_element_type.1195 = s32[1,1,41]{2,1,0} convert(%broadcast_in_dim.1094), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/convert_element_type" stack_frame_id=160} + %add.752 = s32[1,1]{1,0} reshape(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/add" stack_frame_id=147} + %lt.450 = s32[1,1]{1,0} broadcast(%add.752), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/lt" stack_frame_id=152} + %lt.451 = s32[1]{0} reshape(%lt.450), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/lt" stack_frame_id=152} + %lt.452 = s32[1,41]{1,0} broadcast(%lt.451), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/lt" stack_frame_id=152} + %iota.399 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/iota" stack_frame_id=150} + %broadcast_in_dim.1095 = s32[1,41]{1,0} reshape(%iota.399), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/broadcast_in_dim" stack_frame_id=148} + %add.753 = s32[1,41]{1,0} add(%broadcast_in_dim.1095, %broadcast.43), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/add" stack_frame_id=151} + %lt.453 = pred[1,41]{1,0} compare(%lt.452, %add.753), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/lt" stack_frame_id=152} + %convert_element_type.1194 = s32[1,41]{1,0} convert(%lt.453), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/convert_element_type" stack_frame_id=156} + %broadcast_in_dim.1096 = s32[1,1,41]{2,1,0} reshape(%convert_element_type.1194), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/broadcast_in_dim" stack_frame_id=160} + %min.67 = s32[1,1,41]{2,1,0} minimum(%convert_element_type.1195, %broadcast_in_dim.1096), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/min" stack_frame_id=160} + %broadcast_in_dim.1101 = s32[1,1,1,41]{3,2,1,0} reshape(%min.67), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/broadcast_in_dim" stack_frame_id=341} + %convert_element_type.1205 = pred[1,1,1,41]{3,2,1,0} compare(%broadcast_in_dim.1101, %broadcast.37), direction=NE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/convert_element_type" stack_frame_id=345} + %broadcast_in_dim.1102 = pred[1,1,1,1,41]{4,3,2,1,0} reshape(%convert_element_type.1205), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/broadcast_in_dim" stack_frame_id=83} + %and.209 = pred[1,1,1,1,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1102), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/vmap()/and" stack_frame_id=83} + %and.210 = pred[1,1,1,41]{3,2,1,0} reshape(%and.209), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/vmap()/and" stack_frame_id=83} + %and.211 = pred[1,1,32,1,41]{4,3,2,1,0} broadcast(%and.210), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/vmap()/and" stack_frame_id=83} + %dot_general.741 = bf16[1,1,32,128]{3,2,1,0} dot(%convert_element_type.1198, %arg_tuple.5#45), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=186} + %convert_element_type.1199 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/convert_element_type" stack_frame_id=208} + %add.755 = f32[1,1]{1,0} reshape(%convert_element_type.1199), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/add"} + %iota.400 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/iota" stack_frame_id=197} + %mul.2235 = f32[64]{0} multiply(%iota.400, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=197} + %div.684 = f32[64]{0} divide(%mul.2235, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/div" stack_frame_id=198} + %neg.278 = f32[64]{0} negate(%div.684), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/neg" stack_frame_id=199} + %pow.272 = f32[64]{0} power(%broadcast.39, %neg.278), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/pow" stack_frame_id=202} + %div.685 = f32[64]{0} divide(%pow.272, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/div" stack_frame_id=203} + %dot_general.742 = f32[1,1,64]{2,1,0} dot(%add.755, %div.685), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/bi,j->bij/dot_general" stack_frame_id=211} + %stack.1080 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.742), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/stack" stack_frame_id=214} + %stack.1081 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.742), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/stack" stack_frame_id=214} + %stack.1082 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1080, %stack.1081), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/stack" stack_frame_id=214} + %reshape.502 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1082), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/reshape"} + %cos.135 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.502), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/cos" stack_frame_id=218} + %convert_element_type.1200 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.135), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/convert_element_type" stack_frame_id=222} + %mul.2236 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1200), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=242} + %mul.2237 = bf16[1,1,128]{2,1,0} reshape(%mul.2236), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=242} + %mul.2238 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2237), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=242} + %mul.2239 = bf16[1,1,32,128]{3,2,1,0} multiply(%dot_general.741, %mul.2238), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=242} + %split.271 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.741), slice={[0:1], [0:1], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/split" stack_frame_id=234} + %neg.279 = bf16[1,1,32,64]{3,2,1,0} negate(%split.271), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/neg" stack_frame_id=235} + %stack.1083 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%neg.279), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/stack" stack_frame_id=238} + %split.270 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.741), slice={[0:1], [0:1], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/split" stack_frame_id=234} + %stack.1084 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%split.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/stack" stack_frame_id=238} + %stack.1085 = bf16[1,1,32,2,64]{4,3,2,1,0} concatenate(%stack.1083, %stack.1084), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/stack" stack_frame_id=238} + %reshape.503 = bf16[1,1,32,128]{3,2,1,0} reshape(%stack.1085), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/reshape" stack_frame_id=241} + %sin.135 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.502), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/sin" stack_frame_id=226} + %convert_element_type.1201 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.135), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/convert_element_type" stack_frame_id=230} + %mul.2240 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1201), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=243} + %mul.2241 = bf16[1,1,128]{2,1,0} reshape(%mul.2240), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=243} + %mul.2242 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2241), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=243} + %mul.2243 = bf16[1,1,32,128]{3,2,1,0} multiply(%reshape.503, %mul.2242), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=243} + %add.756 = bf16[1,1,32,128]{3,2,1,0} add(%mul.2239, %mul.2243), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/add" stack_frame_id=244} + %reshape.508 = bf16[1,1,32,1,128]{4,3,2,1,0} reshape(%add.756), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/reshape" stack_frame_id=83} + %slice.109 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.112), slice={[0:1], [0:1], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/slice" stack_frame_id=245} + %squeeze.113 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.109), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/squeeze" stack_frame_id=245} + %dot_general.743 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1198, %arg_tuple.5#46), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=256} + %convert_element_type.1202 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/convert_element_type" stack_frame_id=287} + %add.757 = f32[1,1]{1,0} reshape(%convert_element_type.1202), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/add"} + %iota.401 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/iota" stack_frame_id=276} + %mul.2244 = f32[64]{0} multiply(%iota.401, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=276} + %div.686 = f32[64]{0} divide(%mul.2244, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/div" stack_frame_id=277} + %neg.280 = f32[64]{0} negate(%div.686), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/neg" stack_frame_id=278} + %pow.273 = f32[64]{0} power(%broadcast.39, %neg.280), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/pow" stack_frame_id=281} + %div.687 = f32[64]{0} divide(%pow.273, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/div" stack_frame_id=282} + %dot_general.745 = f32[1,1,64]{2,1,0} dot(%add.757, %div.687), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/bi,j->bij/dot_general" stack_frame_id=290} + %stack.1086 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.745), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/stack" stack_frame_id=293} + %stack.1087 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.745), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/stack" stack_frame_id=293} + %stack.1088 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1086, %stack.1087), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/stack" stack_frame_id=293} + %reshape.504 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1088), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/reshape"} + %cos.136 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.504), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/cos" stack_frame_id=297} + %convert_element_type.1203 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.136), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/convert_element_type" stack_frame_id=301} + %mul.2245 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1203), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=321} + %mul.2246 = bf16[1,1,128]{2,1,0} reshape(%mul.2245), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=321} + %mul.2247 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2246), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=321} + %mul.2248 = bf16[1,1,8,128]{3,2,1,0} multiply(%dot_general.743, %mul.2247), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=321} + %split.273 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.743), slice={[0:1], [0:1], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/split" stack_frame_id=313} + %neg.281 = bf16[1,1,8,64]{3,2,1,0} negate(%split.273), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/neg" stack_frame_id=314} + %stack.1089 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%neg.281), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/stack" stack_frame_id=317} + %split.272 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.743), slice={[0:1], [0:1], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/split" stack_frame_id=313} + %stack.1090 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%split.272), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/stack" stack_frame_id=317} + %stack.1091 = bf16[1,1,8,2,64]{4,3,2,1,0} concatenate(%stack.1089, %stack.1090), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/stack" stack_frame_id=317} + %reshape.505 = bf16[1,1,8,128]{3,2,1,0} reshape(%stack.1091), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/reshape" stack_frame_id=320} + %sin.136 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.504), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/sin" stack_frame_id=305} + %convert_element_type.1204 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.136), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/convert_element_type" stack_frame_id=309} + %mul.2249 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1204), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=322} + %mul.2250 = bf16[1,1,128]{2,1,0} reshape(%mul.2249), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=322} + %mul.2251 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2250), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=322} + %mul.2252 = bf16[1,1,8,128]{3,2,1,0} multiply(%reshape.505, %mul.2251), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=322} + %add.758 = bf16[1,1,8,128]{3,2,1,0} add(%mul.2248, %mul.2252), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/rotary_embedding_4/add" stack_frame_id=323} + %lt.454 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/lt" stack_frame_id=326} + %add.759 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/add" stack_frame_id=326} + %select_n.99 = s32[] select(%lt.454, %add.759, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/select_n" stack_frame_id=326} + %dynamic_update_slice.73 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.113, %add.758, %constant.127, %select_n.99, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/dynamic_update_slice" stack_frame_id=326} + %broadcast_in_dim.1099 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.73), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/broadcast_in_dim" stack_frame_id=335} + %reshape.506 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1099), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/reshape" stack_frame_id=335} + %dot_general.746 = f32[1,32,1,1,41]{4,3,2,1,0} dot(%reshape.508, %reshape.506), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=83} + %mul.2253 = f32[1,32,1,1,41]{4,3,2,1,0} multiply(%dot_general.746, %broadcast.36), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/vmap()/mul" stack_frame_id=83} + %vmap_jit__where__.67 = f32[1,1,32,1,41]{4,3,2,1,0} call(%and.211, %mul.2253, %constant.114), to_apply=%_where_3.137, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/vmap(jit(_where))" stack_frame_id=83} + %reduce_max.452 = f32[1,1,32,1]{3,2,1,0} reduce(%vmap_jit__where__.67, %constant.113), dimensions={4}, to_apply=%region_147.155, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/vmap()/reduce_max" stack_frame_id=83} + %max.76 = f32[1,1,32,1]{3,2,1,0} maximum(%reduce_max.452, %broadcast.35), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/vmap()/max" stack_frame_id=83} + %broadcast_in_dim.1103 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%max.76), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %sub.287 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1103), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/vmap()/sub" stack_frame_id=83} + %sub.288 = f32[1,1,32,1]{3,2,1,0} reshape(%sub.287), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/vmap()/sub" stack_frame_id=83} + %sub.289 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%sub.288), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/vmap()/sub" stack_frame_id=83} + %sub.290 = f32[1,1,32,1,41]{4,3,2,1,0} subtract(%vmap_jit__where__.67, %sub.289), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/vmap()/sub" stack_frame_id=83} + %exp.72 = f32[1,1,32,1,41]{4,3,2,1,0} exponential(%sub.290), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/vmap()/exp" stack_frame_id=83} + %reduce_sum.1364 = f32[1,1,32,1]{3,2,1,0} reduce(%exp.72, %constant.123), dimensions={4}, to_apply=%region_148.156, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/vmap()/reduce_sum" stack_frame_id=83} + %broadcast_in_dim.1104 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%reduce_sum.1364), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %div.688 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1104), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/vmap()/div" stack_frame_id=83} + %div.689 = f32[1,1,32,1]{3,2,1,0} reshape(%div.688), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/vmap()/div" stack_frame_id=83} + %div.690 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%div.689), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/vmap()/div" stack_frame_id=83} + %div.691 = f32[1,1,32,1,41]{4,3,2,1,0} divide(%exp.72, %div.690), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/vmap()/div" stack_frame_id=83} + %convert_element_type.1206 = bf16[1,1,32,1,41]{4,3,2,1,0} convert(%div.691), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/vmap()/convert_element_type" stack_frame_id=83} + %dot_general.747 = bf16[1,32,128,1,1]{4,3,2,1,0} dot(%reshape.507, %convert_element_type.1206), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=83} + %transpose.71 = bf16[1,1,32,128]{3,2,1,0} reshape(%dot_general.747), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/reshape;jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/vmap()/transpose;jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %dot_general.748 = bf16[1,1,4096]{2,1,0} dot(%transpose.71, %arg_tuple.5#48), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=354} + %add.761 = bf16[1,1,4096]{2,1,0} add(%dot_general.748, %add.750), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/add" stack_frame_id=355} + %convert_element_type.1207 = f32[1,1,4096]{2,1,0} convert(%add.761), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/feedforward_layernorm/convert_element_type" stack_frame_id=359} + %pow.274 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1207, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/feedforward_layernorm/pow" stack_frame_id=362} + %reduce_sum.1365 = f32[1,1]{1,0} reduce(%pow.274, %constant.123), dimensions={2}, to_apply=%region_149.157, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/feedforward_layernorm/reduce_sum" stack_frame_id=92} + %broadcast_in_dim.1105 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1365), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/feedforward_layernorm/broadcast_in_dim" stack_frame_id=92} + %div.692 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1105, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/feedforward_layernorm/div" stack_frame_id=92} + %add.762 = f32[1,1,1]{2,1,0} add(%div.692, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/feedforward_layernorm/add" stack_frame_id=363} + %rsqrt.137 = f32[1,1,1]{2,1,0} rsqrt(%add.762), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/feedforward_layernorm/rsqrt" stack_frame_id=366} + %mul.2254 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.137), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2255 = f32[1,1]{1,0} reshape(%mul.2254), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2256 = f32[1,1,4096]{2,1,0} broadcast(%mul.2255), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2257 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1207, %mul.2256), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/feedforward_layernorm/mul" stack_frame_id=367} + %convert_element_type.1208 = f32[4096]{0} convert(%arg_tuple.5#49), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/feedforward_layernorm/convert_element_type" stack_frame_id=368} + %broadcast_in_dim.1106 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1208), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/feedforward_layernorm/broadcast_in_dim" stack_frame_id=368} + %mul.2258 = f32[1,1,4096]{2,1,0} multiply(%mul.2257, %broadcast_in_dim.1106), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/feedforward_layernorm/mul" stack_frame_id=368} + %convert_element_type.1209 = bf16[1,1,4096]{2,1,0} convert(%mul.2258), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/feedforward_layernorm/convert_element_type" stack_frame_id=372} + %dot_general.750 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1209, %arg_tuple.5#51), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/feedforward_intermediate_dense/dot_general" stack_frame_id=401} + %dot_general.749 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1209, %arg_tuple.5#50), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/feedforward_gate_dense/dot_general" stack_frame_id=381} + %convert_element_type.1210 = f32[1,1,14336]{2,1,0} convert(%dot_general.749), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/convert_element_type" stack_frame_id=385} + %jit_silu_.67 = f32[1,1,14336]{2,1,0} call(%convert_element_type.1210), to_apply=%silu_4.141, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/jit(silu)" stack_frame_id=388} + %convert_element_type.1211 = bf16[1,1,14336]{2,1,0} convert(%jit_silu_.67), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/convert_element_type" stack_frame_id=392} + %mul.2259 = bf16[1,1,14336]{2,1,0} multiply(%dot_general.750, %convert_element_type.1211), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/mul" stack_frame_id=404} + %dot_general.751 = bf16[1,1,4096]{2,1,0} dot(%mul.2259, %arg_tuple.5#52), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/feedforward_output_dense/dot_general" stack_frame_id=413} + %add.763 = bf16[1,1,4096]{2,1,0} add(%dot_general.751, %add.761), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/add" stack_frame_id=414} + %convert_element_type.1215 = f32[1,1,4096]{2,1,0} convert(%add.763), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention_layernorm/convert_element_type" stack_frame_id=164} + %pow.275 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1215, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention_layernorm/pow" stack_frame_id=167} + %reduce_sum.1366 = f32[1,1]{1,0} reduce(%pow.275, %constant.123), dimensions={2}, to_apply=%region_150.158, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention_layernorm/reduce_sum" stack_frame_id=73} + %broadcast_in_dim.1111 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1366), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention_layernorm/broadcast_in_dim" stack_frame_id=73} + %div.693 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1111, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention_layernorm/div" stack_frame_id=73} + %add.767 = f32[1,1,1]{2,1,0} add(%div.693, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention_layernorm/add" stack_frame_id=168} + %rsqrt.138 = f32[1,1,1]{2,1,0} rsqrt(%add.767), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention_layernorm/rsqrt" stack_frame_id=171} + %mul.2260 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.138), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2261 = f32[1,1]{1,0} reshape(%mul.2260), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2262 = f32[1,1,4096]{2,1,0} broadcast(%mul.2261), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2263 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1215, %mul.2262), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention_layernorm/mul" stack_frame_id=172} + %convert_element_type.1216 = f32[4096]{0} convert(%arg_tuple.5#53), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention_layernorm/convert_element_type" stack_frame_id=173} + %broadcast_in_dim.1112 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1216), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention_layernorm/broadcast_in_dim" stack_frame_id=173} + %mul.2264 = f32[1,1,4096]{2,1,0} multiply(%mul.2263, %broadcast_in_dim.1112), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention_layernorm/mul" stack_frame_id=173} + %convert_element_type.1217 = bf16[1,1,4096]{2,1,0} convert(%mul.2264), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention_layernorm/convert_element_type" stack_frame_id=177} + %dot_general.755 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1217, %arg_tuple.5#56), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=265} + %lt.461 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/lt" stack_frame_id=329} + %add.773 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/add" stack_frame_id=329} + %select_n.102 = s32[] select(%lt.461, %add.773, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/select_n" stack_frame_id=329} + %dynamic_update_slice.76 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.117, %dot_general.755, %constant.127, %select_n.102, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/dynamic_update_slice" stack_frame_id=329} + %broadcast_in_dim.1114 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.76), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/broadcast_in_dim" stack_frame_id=338} + %reshape.514 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1114), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/reshape" stack_frame_id=338} + %convert_element_type.1212 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/convert_element_type" stack_frame_id=139} + %add.764 = f32[1,1]{1,0} reshape(%convert_element_type.1212), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_5/add"} + %ge.373 = f32[1,1]{1,0} broadcast(%add.764), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/ge" stack_frame_id=143} + %ge.374 = f32[1]{0} reshape(%ge.373), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/ge" stack_frame_id=143} + %ge.375 = f32[1,41]{1,0} broadcast(%ge.374), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/ge" stack_frame_id=143} + %iota.402 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/iota" stack_frame_id=142} + %broadcast_in_dim.1107 = f32[1,41]{1,0} reshape(%iota.402), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/broadcast_in_dim" stack_frame_id=143} + %ge.376 = pred[1,41]{1,0} compare(%ge.375, %broadcast_in_dim.1107), direction=GE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/ge" stack_frame_id=143} + %broadcast_in_dim.1108 = pred[1,1,41]{2,1,0} reshape(%ge.376), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/broadcast_in_dim" stack_frame_id=146} + %convert_element_type.1214 = s32[1,1,41]{2,1,0} convert(%broadcast_in_dim.1108), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/convert_element_type" stack_frame_id=160} + %add.765 = s32[1,1]{1,0} reshape(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/add" stack_frame_id=147} + %lt.456 = s32[1,1]{1,0} broadcast(%add.765), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/lt" stack_frame_id=152} + %lt.457 = s32[1]{0} reshape(%lt.456), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/lt" stack_frame_id=152} + %lt.458 = s32[1,41]{1,0} broadcast(%lt.457), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/lt" stack_frame_id=152} + %iota.403 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/iota" stack_frame_id=150} + %broadcast_in_dim.1109 = s32[1,41]{1,0} reshape(%iota.403), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/broadcast_in_dim" stack_frame_id=148} + %add.766 = s32[1,41]{1,0} add(%broadcast_in_dim.1109, %broadcast.43), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/add" stack_frame_id=151} + %lt.459 = pred[1,41]{1,0} compare(%lt.458, %add.766), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/lt" stack_frame_id=152} + %convert_element_type.1213 = s32[1,41]{1,0} convert(%lt.459), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/convert_element_type" stack_frame_id=156} + %broadcast_in_dim.1110 = s32[1,1,41]{2,1,0} reshape(%convert_element_type.1213), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/broadcast_in_dim" stack_frame_id=160} + %min.68 = s32[1,1,41]{2,1,0} minimum(%convert_element_type.1214, %broadcast_in_dim.1110), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/min" stack_frame_id=160} + %broadcast_in_dim.1115 = s32[1,1,1,41]{3,2,1,0} reshape(%min.68), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/broadcast_in_dim" stack_frame_id=341} + %convert_element_type.1224 = pred[1,1,1,41]{3,2,1,0} compare(%broadcast_in_dim.1115, %broadcast.37), direction=NE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/convert_element_type" stack_frame_id=345} + %broadcast_in_dim.1116 = pred[1,1,1,1,41]{4,3,2,1,0} reshape(%convert_element_type.1224), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/broadcast_in_dim" stack_frame_id=83} + %and.212 = pred[1,1,1,1,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1116), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/vmap()/and" stack_frame_id=83} + %and.213 = pred[1,1,1,41]{3,2,1,0} reshape(%and.212), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/vmap()/and" stack_frame_id=83} + %and.214 = pred[1,1,32,1,41]{4,3,2,1,0} broadcast(%and.213), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/vmap()/and" stack_frame_id=83} + %dot_general.752 = bf16[1,1,32,128]{3,2,1,0} dot(%convert_element_type.1217, %arg_tuple.5#54), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=186} + %convert_element_type.1218 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/convert_element_type" stack_frame_id=208} + %add.768 = f32[1,1]{1,0} reshape(%convert_element_type.1218), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/add"} + %iota.404 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/iota" stack_frame_id=197} + %mul.2265 = f32[64]{0} multiply(%iota.404, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=197} + %div.694 = f32[64]{0} divide(%mul.2265, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/div" stack_frame_id=198} + %neg.282 = f32[64]{0} negate(%div.694), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/neg" stack_frame_id=199} + %pow.276 = f32[64]{0} power(%broadcast.39, %neg.282), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/pow" stack_frame_id=202} + %div.695 = f32[64]{0} divide(%pow.276, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/div" stack_frame_id=203} + %dot_general.753 = f32[1,1,64]{2,1,0} dot(%add.768, %div.695), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/bi,j->bij/dot_general" stack_frame_id=211} + %stack.1095 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.753), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/stack" stack_frame_id=214} + %stack.1096 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.753), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/stack" stack_frame_id=214} + %stack.1097 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1095, %stack.1096), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/stack" stack_frame_id=214} + %reshape.509 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1097), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/reshape"} + %cos.137 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.509), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/cos" stack_frame_id=218} + %convert_element_type.1219 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.137), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/convert_element_type" stack_frame_id=222} + %mul.2266 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1219), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=242} + %mul.2267 = bf16[1,1,128]{2,1,0} reshape(%mul.2266), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=242} + %mul.2268 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2267), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=242} + %mul.2269 = bf16[1,1,32,128]{3,2,1,0} multiply(%dot_general.752, %mul.2268), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=242} + %split.275 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.752), slice={[0:1], [0:1], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/split" stack_frame_id=234} + %neg.283 = bf16[1,1,32,64]{3,2,1,0} negate(%split.275), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/neg" stack_frame_id=235} + %stack.1098 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%neg.283), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/stack" stack_frame_id=238} + %split.274 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.752), slice={[0:1], [0:1], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/split" stack_frame_id=234} + %stack.1099 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%split.274), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/stack" stack_frame_id=238} + %stack.1100 = bf16[1,1,32,2,64]{4,3,2,1,0} concatenate(%stack.1098, %stack.1099), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/stack" stack_frame_id=238} + %reshape.510 = bf16[1,1,32,128]{3,2,1,0} reshape(%stack.1100), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/reshape" stack_frame_id=241} + %sin.137 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.509), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/sin" stack_frame_id=226} + %convert_element_type.1220 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.137), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/convert_element_type" stack_frame_id=230} + %mul.2270 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1220), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=243} + %mul.2271 = bf16[1,1,128]{2,1,0} reshape(%mul.2270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=243} + %mul.2272 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2271), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=243} + %mul.2273 = bf16[1,1,32,128]{3,2,1,0} multiply(%reshape.510, %mul.2272), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=243} + %add.769 = bf16[1,1,32,128]{3,2,1,0} add(%mul.2269, %mul.2273), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/add" stack_frame_id=244} + %reshape.515 = bf16[1,1,32,1,128]{4,3,2,1,0} reshape(%add.769), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/reshape" stack_frame_id=83} + %slice.112 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.115), slice={[0:1], [0:1], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/slice" stack_frame_id=245} + %squeeze.116 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.112), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/squeeze" stack_frame_id=245} + %dot_general.754 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1217, %arg_tuple.5#55), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=256} + %convert_element_type.1221 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/convert_element_type" stack_frame_id=287} + %add.770 = f32[1,1]{1,0} reshape(%convert_element_type.1221), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/add"} + %iota.405 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/iota" stack_frame_id=276} + %mul.2274 = f32[64]{0} multiply(%iota.405, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=276} + %div.696 = f32[64]{0} divide(%mul.2274, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/div" stack_frame_id=277} + %neg.284 = f32[64]{0} negate(%div.696), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/neg" stack_frame_id=278} + %pow.277 = f32[64]{0} power(%broadcast.39, %neg.284), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/pow" stack_frame_id=281} + %div.697 = f32[64]{0} divide(%pow.277, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/div" stack_frame_id=282} + %dot_general.756 = f32[1,1,64]{2,1,0} dot(%add.770, %div.697), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/bi,j->bij/dot_general" stack_frame_id=290} + %stack.1101 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.756), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/stack" stack_frame_id=293} + %stack.1102 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.756), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/stack" stack_frame_id=293} + %stack.1103 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1101, %stack.1102), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/stack" stack_frame_id=293} + %reshape.511 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1103), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/reshape"} + %cos.138 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.511), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/cos" stack_frame_id=297} + %convert_element_type.1222 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.138), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/convert_element_type" stack_frame_id=301} + %mul.2275 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1222), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=321} + %mul.2276 = bf16[1,1,128]{2,1,0} reshape(%mul.2275), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=321} + %mul.2277 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2276), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=321} + %mul.2278 = bf16[1,1,8,128]{3,2,1,0} multiply(%dot_general.754, %mul.2277), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=321} + %split.277 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.754), slice={[0:1], [0:1], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/split" stack_frame_id=313} + %neg.285 = bf16[1,1,8,64]{3,2,1,0} negate(%split.277), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/neg" stack_frame_id=314} + %stack.1104 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%neg.285), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/stack" stack_frame_id=317} + %split.276 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.754), slice={[0:1], [0:1], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/split" stack_frame_id=313} + %stack.1105 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%split.276), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/stack" stack_frame_id=317} + %stack.1106 = bf16[1,1,8,2,64]{4,3,2,1,0} concatenate(%stack.1104, %stack.1105), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/stack" stack_frame_id=317} + %reshape.512 = bf16[1,1,8,128]{3,2,1,0} reshape(%stack.1106), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/reshape" stack_frame_id=320} + %sin.138 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.511), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/sin" stack_frame_id=305} + %convert_element_type.1223 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.138), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/convert_element_type" stack_frame_id=309} + %mul.2279 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1223), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=322} + %mul.2280 = bf16[1,1,128]{2,1,0} reshape(%mul.2279), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=322} + %mul.2281 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2280), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=322} + %mul.2282 = bf16[1,1,8,128]{3,2,1,0} multiply(%reshape.512, %mul.2281), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=322} + %add.771 = bf16[1,1,8,128]{3,2,1,0} add(%mul.2278, %mul.2282), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/rotary_embedding_5/add" stack_frame_id=323} + %lt.460 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/lt" stack_frame_id=326} + %add.772 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/add" stack_frame_id=326} + %select_n.101 = s32[] select(%lt.460, %add.772, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/select_n" stack_frame_id=326} + %dynamic_update_slice.75 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.116, %add.771, %constant.127, %select_n.101, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/dynamic_update_slice" stack_frame_id=326} + %broadcast_in_dim.1113 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.75), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/broadcast_in_dim" stack_frame_id=335} + %reshape.513 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1113), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/reshape" stack_frame_id=335} + %dot_general.757 = f32[1,32,1,1,41]{4,3,2,1,0} dot(%reshape.515, %reshape.513), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=83} + %mul.2283 = f32[1,32,1,1,41]{4,3,2,1,0} multiply(%dot_general.757, %broadcast.36), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/vmap()/mul" stack_frame_id=83} + %vmap_jit__where__.68 = f32[1,1,32,1,41]{4,3,2,1,0} call(%and.214, %mul.2283, %constant.114), to_apply=%_where_3.137, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/vmap(jit(_where))" stack_frame_id=83} + %reduce_max.453 = f32[1,1,32,1]{3,2,1,0} reduce(%vmap_jit__where__.68, %constant.113), dimensions={4}, to_apply=%region_151.159, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/vmap()/reduce_max" stack_frame_id=83} + %max.77 = f32[1,1,32,1]{3,2,1,0} maximum(%reduce_max.453, %broadcast.35), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/vmap()/max" stack_frame_id=83} + %broadcast_in_dim.1117 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%max.77), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %sub.291 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1117), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/vmap()/sub" stack_frame_id=83} + %sub.292 = f32[1,1,32,1]{3,2,1,0} reshape(%sub.291), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/vmap()/sub" stack_frame_id=83} + %sub.293 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%sub.292), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/vmap()/sub" stack_frame_id=83} + %sub.294 = f32[1,1,32,1,41]{4,3,2,1,0} subtract(%vmap_jit__where__.68, %sub.293), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/vmap()/sub" stack_frame_id=83} + %exp.73 = f32[1,1,32,1,41]{4,3,2,1,0} exponential(%sub.294), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/vmap()/exp" stack_frame_id=83} + %reduce_sum.1367 = f32[1,1,32,1]{3,2,1,0} reduce(%exp.73, %constant.123), dimensions={4}, to_apply=%region_152.160, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/vmap()/reduce_sum" stack_frame_id=83} + %broadcast_in_dim.1118 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%reduce_sum.1367), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %div.698 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1118), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/vmap()/div" stack_frame_id=83} + %div.699 = f32[1,1,32,1]{3,2,1,0} reshape(%div.698), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/vmap()/div" stack_frame_id=83} + %div.700 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%div.699), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/vmap()/div" stack_frame_id=83} + %div.701 = f32[1,1,32,1,41]{4,3,2,1,0} divide(%exp.73, %div.700), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/vmap()/div" stack_frame_id=83} + %convert_element_type.1225 = bf16[1,1,32,1,41]{4,3,2,1,0} convert(%div.701), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/vmap()/convert_element_type" stack_frame_id=83} + %dot_general.758 = bf16[1,32,128,1,1]{4,3,2,1,0} dot(%reshape.514, %convert_element_type.1225), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=83} + %transpose.72 = bf16[1,1,32,128]{3,2,1,0} reshape(%dot_general.758), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/reshape;jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/vmap()/transpose;jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %dot_general.759 = bf16[1,1,4096]{2,1,0} dot(%transpose.72, %arg_tuple.5#57), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=354} + %add.774 = bf16[1,1,4096]{2,1,0} add(%dot_general.759, %add.763), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/add" stack_frame_id=355} + %convert_element_type.1226 = f32[1,1,4096]{2,1,0} convert(%add.774), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/feedforward_layernorm/convert_element_type" stack_frame_id=359} + %pow.278 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1226, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/feedforward_layernorm/pow" stack_frame_id=362} + %reduce_sum.1368 = f32[1,1]{1,0} reduce(%pow.278, %constant.123), dimensions={2}, to_apply=%region_153.161, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/feedforward_layernorm/reduce_sum" stack_frame_id=92} + %broadcast_in_dim.1119 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1368), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/feedforward_layernorm/broadcast_in_dim" stack_frame_id=92} + %div.702 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1119, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/feedforward_layernorm/div" stack_frame_id=92} + %add.775 = f32[1,1,1]{2,1,0} add(%div.702, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/feedforward_layernorm/add" stack_frame_id=363} + %rsqrt.139 = f32[1,1,1]{2,1,0} rsqrt(%add.775), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/feedforward_layernorm/rsqrt" stack_frame_id=366} + %mul.2284 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.139), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2285 = f32[1,1]{1,0} reshape(%mul.2284), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2286 = f32[1,1,4096]{2,1,0} broadcast(%mul.2285), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2287 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1226, %mul.2286), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/feedforward_layernorm/mul" stack_frame_id=367} + %convert_element_type.1227 = f32[4096]{0} convert(%arg_tuple.5#58), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/feedforward_layernorm/convert_element_type" stack_frame_id=368} + %broadcast_in_dim.1120 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1227), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/feedforward_layernorm/broadcast_in_dim" stack_frame_id=368} + %mul.2288 = f32[1,1,4096]{2,1,0} multiply(%mul.2287, %broadcast_in_dim.1120), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/feedforward_layernorm/mul" stack_frame_id=368} + %convert_element_type.1228 = bf16[1,1,4096]{2,1,0} convert(%mul.2288), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/feedforward_layernorm/convert_element_type" stack_frame_id=372} + %dot_general.761 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1228, %arg_tuple.5#60), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/feedforward_intermediate_dense/dot_general" stack_frame_id=401} + %dot_general.760 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1228, %arg_tuple.5#59), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/feedforward_gate_dense/dot_general" stack_frame_id=381} + %convert_element_type.1229 = f32[1,1,14336]{2,1,0} convert(%dot_general.760), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/convert_element_type" stack_frame_id=385} + %jit_silu_.68 = f32[1,1,14336]{2,1,0} call(%convert_element_type.1229), to_apply=%silu_4.141, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/jit(silu)" stack_frame_id=388} + %convert_element_type.1230 = bf16[1,1,14336]{2,1,0} convert(%jit_silu_.68), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/convert_element_type" stack_frame_id=392} + %mul.2289 = bf16[1,1,14336]{2,1,0} multiply(%dot_general.761, %convert_element_type.1230), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/mul" stack_frame_id=404} + %dot_general.762 = bf16[1,1,4096]{2,1,0} dot(%mul.2289, %arg_tuple.5#61), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/feedforward_output_dense/dot_general" stack_frame_id=413} + %add.776 = bf16[1,1,4096]{2,1,0} add(%dot_general.762, %add.774), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/add" stack_frame_id=414} + %convert_element_type.1234 = f32[1,1,4096]{2,1,0} convert(%add.776), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention_layernorm/convert_element_type" stack_frame_id=164} + %pow.279 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1234, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention_layernorm/pow" stack_frame_id=167} + %reduce_sum.1369 = f32[1,1]{1,0} reduce(%pow.279, %constant.123), dimensions={2}, to_apply=%region_154.162, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention_layernorm/reduce_sum" stack_frame_id=73} + %broadcast_in_dim.1125 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1369), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention_layernorm/broadcast_in_dim" stack_frame_id=73} + %div.703 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1125, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention_layernorm/div" stack_frame_id=73} + %add.780 = f32[1,1,1]{2,1,0} add(%div.703, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention_layernorm/add" stack_frame_id=168} + %rsqrt.140 = f32[1,1,1]{2,1,0} rsqrt(%add.780), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention_layernorm/rsqrt" stack_frame_id=171} + %mul.2290 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.140), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2291 = f32[1,1]{1,0} reshape(%mul.2290), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2292 = f32[1,1,4096]{2,1,0} broadcast(%mul.2291), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2293 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1234, %mul.2292), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention_layernorm/mul" stack_frame_id=172} + %convert_element_type.1235 = f32[4096]{0} convert(%arg_tuple.5#62), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention_layernorm/convert_element_type" stack_frame_id=173} + %broadcast_in_dim.1126 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1235), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention_layernorm/broadcast_in_dim" stack_frame_id=173} + %mul.2294 = f32[1,1,4096]{2,1,0} multiply(%mul.2293, %broadcast_in_dim.1126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention_layernorm/mul" stack_frame_id=173} + %convert_element_type.1236 = bf16[1,1,4096]{2,1,0} convert(%mul.2294), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention_layernorm/convert_element_type" stack_frame_id=177} + %dot_general.766 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1236, %arg_tuple.5#65), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=265} + %lt.467 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/lt" stack_frame_id=329} + %add.786 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/add" stack_frame_id=329} + %select_n.104 = s32[] select(%lt.467, %add.786, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/select_n" stack_frame_id=329} + %dynamic_update_slice.78 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.120, %dot_general.766, %constant.127, %select_n.104, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/dynamic_update_slice" stack_frame_id=329} + %broadcast_in_dim.1128 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.78), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/broadcast_in_dim" stack_frame_id=338} + %reshape.521 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1128), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/reshape" stack_frame_id=338} + %convert_element_type.1231 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/convert_element_type" stack_frame_id=139} + %add.777 = f32[1,1]{1,0} reshape(%convert_element_type.1231), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_6/add"} + %ge.377 = f32[1,1]{1,0} broadcast(%add.777), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/ge" stack_frame_id=143} + %ge.378 = f32[1]{0} reshape(%ge.377), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/ge" stack_frame_id=143} + %ge.379 = f32[1,41]{1,0} broadcast(%ge.378), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/ge" stack_frame_id=143} + %iota.406 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/iota" stack_frame_id=142} + %broadcast_in_dim.1121 = f32[1,41]{1,0} reshape(%iota.406), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/broadcast_in_dim" stack_frame_id=143} + %ge.380 = pred[1,41]{1,0} compare(%ge.379, %broadcast_in_dim.1121), direction=GE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/ge" stack_frame_id=143} + %broadcast_in_dim.1122 = pred[1,1,41]{2,1,0} reshape(%ge.380), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/broadcast_in_dim" stack_frame_id=146} + %convert_element_type.1233 = s32[1,1,41]{2,1,0} convert(%broadcast_in_dim.1122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/convert_element_type" stack_frame_id=160} + %add.778 = s32[1,1]{1,0} reshape(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/add" stack_frame_id=147} + %lt.462 = s32[1,1]{1,0} broadcast(%add.778), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/lt" stack_frame_id=152} + %lt.463 = s32[1]{0} reshape(%lt.462), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/lt" stack_frame_id=152} + %lt.464 = s32[1,41]{1,0} broadcast(%lt.463), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/lt" stack_frame_id=152} + %iota.407 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/iota" stack_frame_id=150} + %broadcast_in_dim.1123 = s32[1,41]{1,0} reshape(%iota.407), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/broadcast_in_dim" stack_frame_id=148} + %add.779 = s32[1,41]{1,0} add(%broadcast_in_dim.1123, %broadcast.43), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/add" stack_frame_id=151} + %lt.465 = pred[1,41]{1,0} compare(%lt.464, %add.779), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/lt" stack_frame_id=152} + %convert_element_type.1232 = s32[1,41]{1,0} convert(%lt.465), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/convert_element_type" stack_frame_id=156} + %broadcast_in_dim.1124 = s32[1,1,41]{2,1,0} reshape(%convert_element_type.1232), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/broadcast_in_dim" stack_frame_id=160} + %min.69 = s32[1,1,41]{2,1,0} minimum(%convert_element_type.1233, %broadcast_in_dim.1124), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/min" stack_frame_id=160} + %broadcast_in_dim.1129 = s32[1,1,1,41]{3,2,1,0} reshape(%min.69), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/broadcast_in_dim" stack_frame_id=341} + %convert_element_type.1243 = pred[1,1,1,41]{3,2,1,0} compare(%broadcast_in_dim.1129, %broadcast.37), direction=NE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/convert_element_type" stack_frame_id=345} + %broadcast_in_dim.1130 = pred[1,1,1,1,41]{4,3,2,1,0} reshape(%convert_element_type.1243), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/broadcast_in_dim" stack_frame_id=83} + %and.215 = pred[1,1,1,1,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1130), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/vmap()/and" stack_frame_id=83} + %and.216 = pred[1,1,1,41]{3,2,1,0} reshape(%and.215), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/vmap()/and" stack_frame_id=83} + %and.217 = pred[1,1,32,1,41]{4,3,2,1,0} broadcast(%and.216), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/vmap()/and" stack_frame_id=83} + %dot_general.763 = bf16[1,1,32,128]{3,2,1,0} dot(%convert_element_type.1236, %arg_tuple.5#63), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=186} + %convert_element_type.1237 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/convert_element_type" stack_frame_id=208} + %add.781 = f32[1,1]{1,0} reshape(%convert_element_type.1237), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/add"} + %iota.408 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/iota" stack_frame_id=197} + %mul.2295 = f32[64]{0} multiply(%iota.408, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=197} + %div.704 = f32[64]{0} divide(%mul.2295, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/div" stack_frame_id=198} + %neg.286 = f32[64]{0} negate(%div.704), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/neg" stack_frame_id=199} + %pow.280 = f32[64]{0} power(%broadcast.39, %neg.286), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/pow" stack_frame_id=202} + %div.705 = f32[64]{0} divide(%pow.280, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/div" stack_frame_id=203} + %dot_general.764 = f32[1,1,64]{2,1,0} dot(%add.781, %div.705), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/bi,j->bij/dot_general" stack_frame_id=211} + %stack.1110 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.764), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/stack" stack_frame_id=214} + %stack.1111 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.764), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/stack" stack_frame_id=214} + %stack.1112 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1110, %stack.1111), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/stack" stack_frame_id=214} + %reshape.516 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1112), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/reshape"} + %cos.139 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.516), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/cos" stack_frame_id=218} + %convert_element_type.1238 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.139), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/convert_element_type" stack_frame_id=222} + %mul.2296 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1238), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=242} + %mul.2297 = bf16[1,1,128]{2,1,0} reshape(%mul.2296), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=242} + %mul.2298 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2297), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=242} + %mul.2299 = bf16[1,1,32,128]{3,2,1,0} multiply(%dot_general.763, %mul.2298), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=242} + %split.279 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.763), slice={[0:1], [0:1], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/split" stack_frame_id=234} + %neg.287 = bf16[1,1,32,64]{3,2,1,0} negate(%split.279), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/neg" stack_frame_id=235} + %stack.1113 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%neg.287), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/stack" stack_frame_id=238} + %split.278 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.763), slice={[0:1], [0:1], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/split" stack_frame_id=234} + %stack.1114 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%split.278), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/stack" stack_frame_id=238} + %stack.1115 = bf16[1,1,32,2,64]{4,3,2,1,0} concatenate(%stack.1113, %stack.1114), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/stack" stack_frame_id=238} + %reshape.517 = bf16[1,1,32,128]{3,2,1,0} reshape(%stack.1115), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/reshape" stack_frame_id=241} + %sin.139 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.516), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/sin" stack_frame_id=226} + %convert_element_type.1239 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.139), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/convert_element_type" stack_frame_id=230} + %mul.2300 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1239), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=243} + %mul.2301 = bf16[1,1,128]{2,1,0} reshape(%mul.2300), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=243} + %mul.2302 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2301), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=243} + %mul.2303 = bf16[1,1,32,128]{3,2,1,0} multiply(%reshape.517, %mul.2302), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=243} + %add.782 = bf16[1,1,32,128]{3,2,1,0} add(%mul.2299, %mul.2303), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/add" stack_frame_id=244} + %reshape.522 = bf16[1,1,32,1,128]{4,3,2,1,0} reshape(%add.782), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/reshape" stack_frame_id=83} + %slice.115 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.118), slice={[0:1], [0:1], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/slice" stack_frame_id=245} + %squeeze.119 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.115), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/squeeze" stack_frame_id=245} + %dot_general.765 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1236, %arg_tuple.5#64), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=256} + %convert_element_type.1240 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/convert_element_type" stack_frame_id=287} + %add.783 = f32[1,1]{1,0} reshape(%convert_element_type.1240), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/add"} + %iota.409 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/iota" stack_frame_id=276} + %mul.2304 = f32[64]{0} multiply(%iota.409, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=276} + %div.706 = f32[64]{0} divide(%mul.2304, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/div" stack_frame_id=277} + %neg.288 = f32[64]{0} negate(%div.706), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/neg" stack_frame_id=278} + %pow.281 = f32[64]{0} power(%broadcast.39, %neg.288), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/pow" stack_frame_id=281} + %div.707 = f32[64]{0} divide(%pow.281, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/div" stack_frame_id=282} + %dot_general.767 = f32[1,1,64]{2,1,0} dot(%add.783, %div.707), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/bi,j->bij/dot_general" stack_frame_id=290} + %stack.1116 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.767), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/stack" stack_frame_id=293} + %stack.1117 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.767), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/stack" stack_frame_id=293} + %stack.1118 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1116, %stack.1117), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/stack" stack_frame_id=293} + %reshape.518 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1118), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/reshape"} + %cos.140 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.518), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/cos" stack_frame_id=297} + %convert_element_type.1241 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.140), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/convert_element_type" stack_frame_id=301} + %mul.2305 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1241), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=321} + %mul.2306 = bf16[1,1,128]{2,1,0} reshape(%mul.2305), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=321} + %mul.2307 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2306), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=321} + %mul.2308 = bf16[1,1,8,128]{3,2,1,0} multiply(%dot_general.765, %mul.2307), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=321} + %split.281 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.765), slice={[0:1], [0:1], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/split" stack_frame_id=313} + %neg.289 = bf16[1,1,8,64]{3,2,1,0} negate(%split.281), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/neg" stack_frame_id=314} + %stack.1119 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%neg.289), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/stack" stack_frame_id=317} + %split.280 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.765), slice={[0:1], [0:1], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/split" stack_frame_id=313} + %stack.1120 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%split.280), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/stack" stack_frame_id=317} + %stack.1121 = bf16[1,1,8,2,64]{4,3,2,1,0} concatenate(%stack.1119, %stack.1120), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/stack" stack_frame_id=317} + %reshape.519 = bf16[1,1,8,128]{3,2,1,0} reshape(%stack.1121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/reshape" stack_frame_id=320} + %sin.140 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.518), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/sin" stack_frame_id=305} + %convert_element_type.1242 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.140), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/convert_element_type" stack_frame_id=309} + %mul.2309 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1242), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=322} + %mul.2310 = bf16[1,1,128]{2,1,0} reshape(%mul.2309), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=322} + %mul.2311 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2310), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=322} + %mul.2312 = bf16[1,1,8,128]{3,2,1,0} multiply(%reshape.519, %mul.2311), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=322} + %add.784 = bf16[1,1,8,128]{3,2,1,0} add(%mul.2308, %mul.2312), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/rotary_embedding_6/add" stack_frame_id=323} + %lt.466 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/lt" stack_frame_id=326} + %add.785 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/add" stack_frame_id=326} + %select_n.103 = s32[] select(%lt.466, %add.785, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/select_n" stack_frame_id=326} + %dynamic_update_slice.77 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.119, %add.784, %constant.127, %select_n.103, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/dynamic_update_slice" stack_frame_id=326} + %broadcast_in_dim.1127 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.77), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/broadcast_in_dim" stack_frame_id=335} + %reshape.520 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/reshape" stack_frame_id=335} + %dot_general.768 = f32[1,32,1,1,41]{4,3,2,1,0} dot(%reshape.522, %reshape.520), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=83} + %mul.2313 = f32[1,32,1,1,41]{4,3,2,1,0} multiply(%dot_general.768, %broadcast.36), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/vmap()/mul" stack_frame_id=83} + %vmap_jit__where__.69 = f32[1,1,32,1,41]{4,3,2,1,0} call(%and.217, %mul.2313, %constant.114), to_apply=%_where_3.137, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/vmap(jit(_where))" stack_frame_id=83} + %reduce_max.454 = f32[1,1,32,1]{3,2,1,0} reduce(%vmap_jit__where__.69, %constant.113), dimensions={4}, to_apply=%region_155.163, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/vmap()/reduce_max" stack_frame_id=83} + %max.78 = f32[1,1,32,1]{3,2,1,0} maximum(%reduce_max.454, %broadcast.35), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/vmap()/max" stack_frame_id=83} + %broadcast_in_dim.1131 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%max.78), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %sub.295 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1131), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/vmap()/sub" stack_frame_id=83} + %sub.296 = f32[1,1,32,1]{3,2,1,0} reshape(%sub.295), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/vmap()/sub" stack_frame_id=83} + %sub.297 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%sub.296), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/vmap()/sub" stack_frame_id=83} + %sub.298 = f32[1,1,32,1,41]{4,3,2,1,0} subtract(%vmap_jit__where__.69, %sub.297), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/vmap()/sub" stack_frame_id=83} + %exp.74 = f32[1,1,32,1,41]{4,3,2,1,0} exponential(%sub.298), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/vmap()/exp" stack_frame_id=83} + %reduce_sum.1370 = f32[1,1,32,1]{3,2,1,0} reduce(%exp.74, %constant.123), dimensions={4}, to_apply=%region_156.164, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/vmap()/reduce_sum" stack_frame_id=83} + %broadcast_in_dim.1132 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%reduce_sum.1370), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %div.708 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1132), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/vmap()/div" stack_frame_id=83} + %div.709 = f32[1,1,32,1]{3,2,1,0} reshape(%div.708), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/vmap()/div" stack_frame_id=83} + %div.710 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%div.709), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/vmap()/div" stack_frame_id=83} + %div.711 = f32[1,1,32,1,41]{4,3,2,1,0} divide(%exp.74, %div.710), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/vmap()/div" stack_frame_id=83} + %convert_element_type.1244 = bf16[1,1,32,1,41]{4,3,2,1,0} convert(%div.711), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/vmap()/convert_element_type" stack_frame_id=83} + %dot_general.769 = bf16[1,32,128,1,1]{4,3,2,1,0} dot(%reshape.521, %convert_element_type.1244), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=83} + %transpose.73 = bf16[1,1,32,128]{3,2,1,0} reshape(%dot_general.769), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/reshape;jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/vmap()/transpose;jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %dot_general.770 = bf16[1,1,4096]{2,1,0} dot(%transpose.73, %arg_tuple.5#66), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=354} + %add.787 = bf16[1,1,4096]{2,1,0} add(%dot_general.770, %add.776), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/add" stack_frame_id=355} + %convert_element_type.1245 = f32[1,1,4096]{2,1,0} convert(%add.787), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/feedforward_layernorm/convert_element_type" stack_frame_id=359} + %pow.282 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1245, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/feedforward_layernorm/pow" stack_frame_id=362} + %reduce_sum.1371 = f32[1,1]{1,0} reduce(%pow.282, %constant.123), dimensions={2}, to_apply=%region_157.165, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/feedforward_layernorm/reduce_sum" stack_frame_id=92} + %broadcast_in_dim.1133 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1371), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/feedforward_layernorm/broadcast_in_dim" stack_frame_id=92} + %div.712 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1133, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/feedforward_layernorm/div" stack_frame_id=92} + %add.788 = f32[1,1,1]{2,1,0} add(%div.712, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/feedforward_layernorm/add" stack_frame_id=363} + %rsqrt.141 = f32[1,1,1]{2,1,0} rsqrt(%add.788), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/feedforward_layernorm/rsqrt" stack_frame_id=366} + %mul.2314 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.141), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2315 = f32[1,1]{1,0} reshape(%mul.2314), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2316 = f32[1,1,4096]{2,1,0} broadcast(%mul.2315), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2317 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1245, %mul.2316), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/feedforward_layernorm/mul" stack_frame_id=367} + %convert_element_type.1246 = f32[4096]{0} convert(%arg_tuple.5#67), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/feedforward_layernorm/convert_element_type" stack_frame_id=368} + %broadcast_in_dim.1134 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1246), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/feedforward_layernorm/broadcast_in_dim" stack_frame_id=368} + %mul.2318 = f32[1,1,4096]{2,1,0} multiply(%mul.2317, %broadcast_in_dim.1134), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/feedforward_layernorm/mul" stack_frame_id=368} + %convert_element_type.1247 = bf16[1,1,4096]{2,1,0} convert(%mul.2318), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/feedforward_layernorm/convert_element_type" stack_frame_id=372} + %dot_general.772 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1247, %arg_tuple.5#69), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/feedforward_intermediate_dense/dot_general" stack_frame_id=401} + %dot_general.771 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1247, %arg_tuple.5#68), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/feedforward_gate_dense/dot_general" stack_frame_id=381} + %convert_element_type.1248 = f32[1,1,14336]{2,1,0} convert(%dot_general.771), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/convert_element_type" stack_frame_id=385} + %jit_silu_.69 = f32[1,1,14336]{2,1,0} call(%convert_element_type.1248), to_apply=%silu_4.141, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/jit(silu)" stack_frame_id=388} + %convert_element_type.1249 = bf16[1,1,14336]{2,1,0} convert(%jit_silu_.69), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/convert_element_type" stack_frame_id=392} + %mul.2319 = bf16[1,1,14336]{2,1,0} multiply(%dot_general.772, %convert_element_type.1249), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/mul" stack_frame_id=404} + %dot_general.773 = bf16[1,1,4096]{2,1,0} dot(%mul.2319, %arg_tuple.5#70), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/feedforward_output_dense/dot_general" stack_frame_id=413} + %add.789 = bf16[1,1,4096]{2,1,0} add(%dot_general.773, %add.787), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/add" stack_frame_id=414} + %convert_element_type.1253 = f32[1,1,4096]{2,1,0} convert(%add.789), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention_layernorm/convert_element_type" stack_frame_id=164} + %pow.283 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1253, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention_layernorm/pow" stack_frame_id=167} + %reduce_sum.1372 = f32[1,1]{1,0} reduce(%pow.283, %constant.123), dimensions={2}, to_apply=%region_158.166, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention_layernorm/reduce_sum" stack_frame_id=73} + %broadcast_in_dim.1139 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1372), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention_layernorm/broadcast_in_dim" stack_frame_id=73} + %div.713 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1139, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention_layernorm/div" stack_frame_id=73} + %add.793 = f32[1,1,1]{2,1,0} add(%div.713, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention_layernorm/add" stack_frame_id=168} + %rsqrt.142 = f32[1,1,1]{2,1,0} rsqrt(%add.793), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention_layernorm/rsqrt" stack_frame_id=171} + %mul.2320 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.142), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2321 = f32[1,1]{1,0} reshape(%mul.2320), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2322 = f32[1,1,4096]{2,1,0} broadcast(%mul.2321), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2323 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1253, %mul.2322), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention_layernorm/mul" stack_frame_id=172} + %convert_element_type.1254 = f32[4096]{0} convert(%arg_tuple.5#71), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention_layernorm/convert_element_type" stack_frame_id=173} + %broadcast_in_dim.1140 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1254), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention_layernorm/broadcast_in_dim" stack_frame_id=173} + %mul.2324 = f32[1,1,4096]{2,1,0} multiply(%mul.2323, %broadcast_in_dim.1140), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention_layernorm/mul" stack_frame_id=173} + %convert_element_type.1255 = bf16[1,1,4096]{2,1,0} convert(%mul.2324), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention_layernorm/convert_element_type" stack_frame_id=177} + %dot_general.777 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1255, %arg_tuple.5#74), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=265} + %lt.473 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/lt" stack_frame_id=329} + %add.799 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/add" stack_frame_id=329} + %select_n.106 = s32[] select(%lt.473, %add.799, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/select_n" stack_frame_id=329} + %dynamic_update_slice.80 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.123, %dot_general.777, %constant.127, %select_n.106, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/dynamic_update_slice" stack_frame_id=329} + %broadcast_in_dim.1142 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.80), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/broadcast_in_dim" stack_frame_id=338} + %reshape.528 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1142), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/reshape" stack_frame_id=338} + %convert_element_type.1250 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/convert_element_type" stack_frame_id=139} + %add.790 = f32[1,1]{1,0} reshape(%convert_element_type.1250), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_7/add"} + %ge.381 = f32[1,1]{1,0} broadcast(%add.790), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/ge" stack_frame_id=143} + %ge.382 = f32[1]{0} reshape(%ge.381), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/ge" stack_frame_id=143} + %ge.383 = f32[1,41]{1,0} broadcast(%ge.382), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/ge" stack_frame_id=143} + %iota.410 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/iota" stack_frame_id=142} + %broadcast_in_dim.1135 = f32[1,41]{1,0} reshape(%iota.410), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/broadcast_in_dim" stack_frame_id=143} + %ge.384 = pred[1,41]{1,0} compare(%ge.383, %broadcast_in_dim.1135), direction=GE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/ge" stack_frame_id=143} + %broadcast_in_dim.1136 = pred[1,1,41]{2,1,0} reshape(%ge.384), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/broadcast_in_dim" stack_frame_id=146} + %convert_element_type.1252 = s32[1,1,41]{2,1,0} convert(%broadcast_in_dim.1136), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/convert_element_type" stack_frame_id=160} + %add.791 = s32[1,1]{1,0} reshape(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/add" stack_frame_id=147} + %lt.468 = s32[1,1]{1,0} broadcast(%add.791), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/lt" stack_frame_id=152} + %lt.469 = s32[1]{0} reshape(%lt.468), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/lt" stack_frame_id=152} + %lt.470 = s32[1,41]{1,0} broadcast(%lt.469), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/lt" stack_frame_id=152} + %iota.411 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/iota" stack_frame_id=150} + %broadcast_in_dim.1137 = s32[1,41]{1,0} reshape(%iota.411), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/broadcast_in_dim" stack_frame_id=148} + %add.792 = s32[1,41]{1,0} add(%broadcast_in_dim.1137, %broadcast.43), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/add" stack_frame_id=151} + %lt.471 = pred[1,41]{1,0} compare(%lt.470, %add.792), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/lt" stack_frame_id=152} + %convert_element_type.1251 = s32[1,41]{1,0} convert(%lt.471), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/convert_element_type" stack_frame_id=156} + %broadcast_in_dim.1138 = s32[1,1,41]{2,1,0} reshape(%convert_element_type.1251), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/broadcast_in_dim" stack_frame_id=160} + %min.70 = s32[1,1,41]{2,1,0} minimum(%convert_element_type.1252, %broadcast_in_dim.1138), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/min" stack_frame_id=160} + %broadcast_in_dim.1143 = s32[1,1,1,41]{3,2,1,0} reshape(%min.70), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/broadcast_in_dim" stack_frame_id=341} + %convert_element_type.1262 = pred[1,1,1,41]{3,2,1,0} compare(%broadcast_in_dim.1143, %broadcast.37), direction=NE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/convert_element_type" stack_frame_id=345} + %broadcast_in_dim.1144 = pred[1,1,1,1,41]{4,3,2,1,0} reshape(%convert_element_type.1262), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/broadcast_in_dim" stack_frame_id=83} + %and.218 = pred[1,1,1,1,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1144), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/vmap()/and" stack_frame_id=83} + %and.219 = pred[1,1,1,41]{3,2,1,0} reshape(%and.218), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/vmap()/and" stack_frame_id=83} + %and.220 = pred[1,1,32,1,41]{4,3,2,1,0} broadcast(%and.219), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/vmap()/and" stack_frame_id=83} + %dot_general.774 = bf16[1,1,32,128]{3,2,1,0} dot(%convert_element_type.1255, %arg_tuple.5#72), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=186} + %convert_element_type.1256 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/convert_element_type" stack_frame_id=208} + %add.794 = f32[1,1]{1,0} reshape(%convert_element_type.1256), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/add"} + %iota.412 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/iota" stack_frame_id=197} + %mul.2325 = f32[64]{0} multiply(%iota.412, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=197} + %div.714 = f32[64]{0} divide(%mul.2325, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/div" stack_frame_id=198} + %neg.290 = f32[64]{0} negate(%div.714), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/neg" stack_frame_id=199} + %pow.284 = f32[64]{0} power(%broadcast.39, %neg.290), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/pow" stack_frame_id=202} + %div.715 = f32[64]{0} divide(%pow.284, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/div" stack_frame_id=203} + %dot_general.775 = f32[1,1,64]{2,1,0} dot(%add.794, %div.715), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/bi,j->bij/dot_general" stack_frame_id=211} + %stack.1125 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.775), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/stack" stack_frame_id=214} + %stack.1126 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.775), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/stack" stack_frame_id=214} + %stack.1127 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1125, %stack.1126), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/stack" stack_frame_id=214} + %reshape.523 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/reshape"} + %cos.141 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.523), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/cos" stack_frame_id=218} + %convert_element_type.1257 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.141), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/convert_element_type" stack_frame_id=222} + %mul.2326 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1257), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=242} + %mul.2327 = bf16[1,1,128]{2,1,0} reshape(%mul.2326), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=242} + %mul.2328 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2327), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=242} + %mul.2329 = bf16[1,1,32,128]{3,2,1,0} multiply(%dot_general.774, %mul.2328), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=242} + %split.283 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.774), slice={[0:1], [0:1], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/split" stack_frame_id=234} + %neg.291 = bf16[1,1,32,64]{3,2,1,0} negate(%split.283), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/neg" stack_frame_id=235} + %stack.1128 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%neg.291), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/stack" stack_frame_id=238} + %split.282 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.774), slice={[0:1], [0:1], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/split" stack_frame_id=234} + %stack.1129 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%split.282), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/stack" stack_frame_id=238} + %stack.1130 = bf16[1,1,32,2,64]{4,3,2,1,0} concatenate(%stack.1128, %stack.1129), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/stack" stack_frame_id=238} + %reshape.524 = bf16[1,1,32,128]{3,2,1,0} reshape(%stack.1130), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/reshape" stack_frame_id=241} + %sin.141 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.523), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/sin" stack_frame_id=226} + %convert_element_type.1258 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.141), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/convert_element_type" stack_frame_id=230} + %mul.2330 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1258), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=243} + %mul.2331 = bf16[1,1,128]{2,1,0} reshape(%mul.2330), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=243} + %mul.2332 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2331), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=243} + %mul.2333 = bf16[1,1,32,128]{3,2,1,0} multiply(%reshape.524, %mul.2332), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=243} + %add.795 = bf16[1,1,32,128]{3,2,1,0} add(%mul.2329, %mul.2333), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/add" stack_frame_id=244} + %reshape.529 = bf16[1,1,32,1,128]{4,3,2,1,0} reshape(%add.795), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/reshape" stack_frame_id=83} + %slice.118 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.121), slice={[0:1], [0:1], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/slice" stack_frame_id=245} + %squeeze.122 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.118), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/squeeze" stack_frame_id=245} + %dot_general.776 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1255, %arg_tuple.5#73), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=256} + %convert_element_type.1259 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/convert_element_type" stack_frame_id=287} + %add.796 = f32[1,1]{1,0} reshape(%convert_element_type.1259), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/add"} + %iota.413 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/iota" stack_frame_id=276} + %mul.2334 = f32[64]{0} multiply(%iota.413, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=276} + %div.716 = f32[64]{0} divide(%mul.2334, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/div" stack_frame_id=277} + %neg.292 = f32[64]{0} negate(%div.716), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/neg" stack_frame_id=278} + %pow.285 = f32[64]{0} power(%broadcast.39, %neg.292), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/pow" stack_frame_id=281} + %div.717 = f32[64]{0} divide(%pow.285, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/div" stack_frame_id=282} + %dot_general.778 = f32[1,1,64]{2,1,0} dot(%add.796, %div.717), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/bi,j->bij/dot_general" stack_frame_id=290} + %stack.1131 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.778), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/stack" stack_frame_id=293} + %stack.1132 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.778), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/stack" stack_frame_id=293} + %stack.1133 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1131, %stack.1132), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/stack" stack_frame_id=293} + %reshape.525 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1133), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/reshape"} + %cos.142 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.525), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/cos" stack_frame_id=297} + %convert_element_type.1260 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.142), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/convert_element_type" stack_frame_id=301} + %mul.2335 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1260), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=321} + %mul.2336 = bf16[1,1,128]{2,1,0} reshape(%mul.2335), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=321} + %mul.2337 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2336), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=321} + %mul.2338 = bf16[1,1,8,128]{3,2,1,0} multiply(%dot_general.776, %mul.2337), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=321} + %split.285 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.776), slice={[0:1], [0:1], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/split" stack_frame_id=313} + %neg.293 = bf16[1,1,8,64]{3,2,1,0} negate(%split.285), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/neg" stack_frame_id=314} + %stack.1134 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%neg.293), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/stack" stack_frame_id=317} + %split.284 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.776), slice={[0:1], [0:1], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/split" stack_frame_id=313} + %stack.1135 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%split.284), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/stack" stack_frame_id=317} + %stack.1136 = bf16[1,1,8,2,64]{4,3,2,1,0} concatenate(%stack.1134, %stack.1135), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/stack" stack_frame_id=317} + %reshape.526 = bf16[1,1,8,128]{3,2,1,0} reshape(%stack.1136), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/reshape" stack_frame_id=320} + %sin.142 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.525), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/sin" stack_frame_id=305} + %convert_element_type.1261 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.142), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/convert_element_type" stack_frame_id=309} + %mul.2339 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1261), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=322} + %mul.2340 = bf16[1,1,128]{2,1,0} reshape(%mul.2339), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=322} + %mul.2341 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2340), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=322} + %mul.2342 = bf16[1,1,8,128]{3,2,1,0} multiply(%reshape.526, %mul.2341), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=322} + %add.797 = bf16[1,1,8,128]{3,2,1,0} add(%mul.2338, %mul.2342), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/rotary_embedding_7/add" stack_frame_id=323} + %lt.472 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/lt" stack_frame_id=326} + %add.798 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/add" stack_frame_id=326} + %select_n.105 = s32[] select(%lt.472, %add.798, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/select_n" stack_frame_id=326} + %dynamic_update_slice.79 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.122, %add.797, %constant.127, %select_n.105, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/dynamic_update_slice" stack_frame_id=326} + %broadcast_in_dim.1141 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.79), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/broadcast_in_dim" stack_frame_id=335} + %reshape.527 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1141), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/reshape" stack_frame_id=335} + %dot_general.779 = f32[1,32,1,1,41]{4,3,2,1,0} dot(%reshape.529, %reshape.527), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=83} + %mul.2343 = f32[1,32,1,1,41]{4,3,2,1,0} multiply(%dot_general.779, %broadcast.36), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/vmap()/mul" stack_frame_id=83} + %vmap_jit__where__.70 = f32[1,1,32,1,41]{4,3,2,1,0} call(%and.220, %mul.2343, %constant.114), to_apply=%_where_3.137, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/vmap(jit(_where))" stack_frame_id=83} + %reduce_max.455 = f32[1,1,32,1]{3,2,1,0} reduce(%vmap_jit__where__.70, %constant.113), dimensions={4}, to_apply=%region_159.167, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/vmap()/reduce_max" stack_frame_id=83} + %max.79 = f32[1,1,32,1]{3,2,1,0} maximum(%reduce_max.455, %broadcast.35), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/vmap()/max" stack_frame_id=83} + %broadcast_in_dim.1145 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%max.79), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %sub.299 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1145), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/vmap()/sub" stack_frame_id=83} + %sub.300 = f32[1,1,32,1]{3,2,1,0} reshape(%sub.299), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/vmap()/sub" stack_frame_id=83} + %sub.301 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%sub.300), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/vmap()/sub" stack_frame_id=83} + %sub.302 = f32[1,1,32,1,41]{4,3,2,1,0} subtract(%vmap_jit__where__.70, %sub.301), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/vmap()/sub" stack_frame_id=83} + %exp.75 = f32[1,1,32,1,41]{4,3,2,1,0} exponential(%sub.302), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/vmap()/exp" stack_frame_id=83} + %reduce_sum.1373 = f32[1,1,32,1]{3,2,1,0} reduce(%exp.75, %constant.123), dimensions={4}, to_apply=%region_160.168, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/vmap()/reduce_sum" stack_frame_id=83} + %broadcast_in_dim.1146 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%reduce_sum.1373), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %div.718 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1146), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/vmap()/div" stack_frame_id=83} + %div.719 = f32[1,1,32,1]{3,2,1,0} reshape(%div.718), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/vmap()/div" stack_frame_id=83} + %div.720 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%div.719), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/vmap()/div" stack_frame_id=83} + %div.721 = f32[1,1,32,1,41]{4,3,2,1,0} divide(%exp.75, %div.720), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/vmap()/div" stack_frame_id=83} + %convert_element_type.1263 = bf16[1,1,32,1,41]{4,3,2,1,0} convert(%div.721), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/vmap()/convert_element_type" stack_frame_id=83} + %dot_general.780 = bf16[1,32,128,1,1]{4,3,2,1,0} dot(%reshape.528, %convert_element_type.1263), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=83} + %transpose.74 = bf16[1,1,32,128]{3,2,1,0} reshape(%dot_general.780), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/reshape;jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/vmap()/transpose;jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %dot_general.781 = bf16[1,1,4096]{2,1,0} dot(%transpose.74, %arg_tuple.5#75), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=354} + %add.800 = bf16[1,1,4096]{2,1,0} add(%dot_general.781, %add.789), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/add" stack_frame_id=355} + %convert_element_type.1264 = f32[1,1,4096]{2,1,0} convert(%add.800), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/feedforward_layernorm/convert_element_type" stack_frame_id=359} + %pow.286 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1264, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/feedforward_layernorm/pow" stack_frame_id=362} + %reduce_sum.1374 = f32[1,1]{1,0} reduce(%pow.286, %constant.123), dimensions={2}, to_apply=%region_161.169, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/feedforward_layernorm/reduce_sum" stack_frame_id=92} + %broadcast_in_dim.1147 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1374), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/feedforward_layernorm/broadcast_in_dim" stack_frame_id=92} + %div.722 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1147, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/feedforward_layernorm/div" stack_frame_id=92} + %add.801 = f32[1,1,1]{2,1,0} add(%div.722, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/feedforward_layernorm/add" stack_frame_id=363} + %rsqrt.143 = f32[1,1,1]{2,1,0} rsqrt(%add.801), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/feedforward_layernorm/rsqrt" stack_frame_id=366} + %mul.2344 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.143), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2345 = f32[1,1]{1,0} reshape(%mul.2344), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2346 = f32[1,1,4096]{2,1,0} broadcast(%mul.2345), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2347 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1264, %mul.2346), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/feedforward_layernorm/mul" stack_frame_id=367} + %convert_element_type.1265 = f32[4096]{0} convert(%arg_tuple.5#76), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/feedforward_layernorm/convert_element_type" stack_frame_id=368} + %broadcast_in_dim.1148 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1265), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/feedforward_layernorm/broadcast_in_dim" stack_frame_id=368} + %mul.2348 = f32[1,1,4096]{2,1,0} multiply(%mul.2347, %broadcast_in_dim.1148), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/feedforward_layernorm/mul" stack_frame_id=368} + %convert_element_type.1266 = bf16[1,1,4096]{2,1,0} convert(%mul.2348), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/feedforward_layernorm/convert_element_type" stack_frame_id=372} + %dot_general.783 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1266, %arg_tuple.5#78), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/feedforward_intermediate_dense/dot_general" stack_frame_id=401} + %dot_general.782 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1266, %arg_tuple.5#77), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/feedforward_gate_dense/dot_general" stack_frame_id=381} + %convert_element_type.1267 = f32[1,1,14336]{2,1,0} convert(%dot_general.782), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/convert_element_type" stack_frame_id=385} + %jit_silu_.70 = f32[1,1,14336]{2,1,0} call(%convert_element_type.1267), to_apply=%silu_4.141, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/jit(silu)" stack_frame_id=388} + %convert_element_type.1268 = bf16[1,1,14336]{2,1,0} convert(%jit_silu_.70), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/convert_element_type" stack_frame_id=392} + %mul.2349 = bf16[1,1,14336]{2,1,0} multiply(%dot_general.783, %convert_element_type.1268), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/mul" stack_frame_id=404} + %dot_general.784 = bf16[1,1,4096]{2,1,0} dot(%mul.2349, %arg_tuple.5#79), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/feedforward_output_dense/dot_general" stack_frame_id=413} + %add.802 = bf16[1,1,4096]{2,1,0} add(%dot_general.784, %add.800), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/add" stack_frame_id=414} + %convert_element_type.1272 = f32[1,1,4096]{2,1,0} convert(%add.802), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention_layernorm/convert_element_type" stack_frame_id=164} + %pow.287 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1272, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention_layernorm/pow" stack_frame_id=167} + %reduce_sum.1375 = f32[1,1]{1,0} reduce(%pow.287, %constant.123), dimensions={2}, to_apply=%region_162.170, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention_layernorm/reduce_sum" stack_frame_id=73} + %broadcast_in_dim.1153 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1375), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention_layernorm/broadcast_in_dim" stack_frame_id=73} + %div.723 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1153, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention_layernorm/div" stack_frame_id=73} + %add.806 = f32[1,1,1]{2,1,0} add(%div.723, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention_layernorm/add" stack_frame_id=168} + %rsqrt.144 = f32[1,1,1]{2,1,0} rsqrt(%add.806), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention_layernorm/rsqrt" stack_frame_id=171} + %mul.2350 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.144), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2351 = f32[1,1]{1,0} reshape(%mul.2350), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2352 = f32[1,1,4096]{2,1,0} broadcast(%mul.2351), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2353 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1272, %mul.2352), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention_layernorm/mul" stack_frame_id=172} + %convert_element_type.1273 = f32[4096]{0} convert(%arg_tuple.5#80), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention_layernorm/convert_element_type" stack_frame_id=173} + %broadcast_in_dim.1154 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1273), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention_layernorm/broadcast_in_dim" stack_frame_id=173} + %mul.2354 = f32[1,1,4096]{2,1,0} multiply(%mul.2353, %broadcast_in_dim.1154), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention_layernorm/mul" stack_frame_id=173} + %convert_element_type.1274 = bf16[1,1,4096]{2,1,0} convert(%mul.2354), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention_layernorm/convert_element_type" stack_frame_id=177} + %dot_general.788 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1274, %arg_tuple.5#83), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=265} + %lt.479 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/lt" stack_frame_id=329} + %add.812 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/add" stack_frame_id=329} + %select_n.108 = s32[] select(%lt.479, %add.812, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/select_n" stack_frame_id=329} + %dynamic_update_slice.82 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.126, %dot_general.788, %constant.127, %select_n.108, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/dynamic_update_slice" stack_frame_id=329} + %broadcast_in_dim.1156 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.82), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/broadcast_in_dim" stack_frame_id=338} + %reshape.535 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1156), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/reshape" stack_frame_id=338} + %convert_element_type.1269 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/convert_element_type" stack_frame_id=139} + %add.803 = f32[1,1]{1,0} reshape(%convert_element_type.1269), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_8/add"} + %ge.385 = f32[1,1]{1,0} broadcast(%add.803), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/ge" stack_frame_id=143} + %ge.386 = f32[1]{0} reshape(%ge.385), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/ge" stack_frame_id=143} + %ge.387 = f32[1,41]{1,0} broadcast(%ge.386), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/ge" stack_frame_id=143} + %iota.414 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/iota" stack_frame_id=142} + %broadcast_in_dim.1149 = f32[1,41]{1,0} reshape(%iota.414), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/broadcast_in_dim" stack_frame_id=143} + %ge.388 = pred[1,41]{1,0} compare(%ge.387, %broadcast_in_dim.1149), direction=GE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/ge" stack_frame_id=143} + %broadcast_in_dim.1150 = pred[1,1,41]{2,1,0} reshape(%ge.388), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/broadcast_in_dim" stack_frame_id=146} + %convert_element_type.1271 = s32[1,1,41]{2,1,0} convert(%broadcast_in_dim.1150), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/convert_element_type" stack_frame_id=160} + %add.804 = s32[1,1]{1,0} reshape(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/add" stack_frame_id=147} + %lt.474 = s32[1,1]{1,0} broadcast(%add.804), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/lt" stack_frame_id=152} + %lt.475 = s32[1]{0} reshape(%lt.474), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/lt" stack_frame_id=152} + %lt.476 = s32[1,41]{1,0} broadcast(%lt.475), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/lt" stack_frame_id=152} + %iota.415 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/iota" stack_frame_id=150} + %broadcast_in_dim.1151 = s32[1,41]{1,0} reshape(%iota.415), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/broadcast_in_dim" stack_frame_id=148} + %add.805 = s32[1,41]{1,0} add(%broadcast_in_dim.1151, %broadcast.43), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/add" stack_frame_id=151} + %lt.477 = pred[1,41]{1,0} compare(%lt.476, %add.805), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/lt" stack_frame_id=152} + %convert_element_type.1270 = s32[1,41]{1,0} convert(%lt.477), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/convert_element_type" stack_frame_id=156} + %broadcast_in_dim.1152 = s32[1,1,41]{2,1,0} reshape(%convert_element_type.1270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/broadcast_in_dim" stack_frame_id=160} + %min.71 = s32[1,1,41]{2,1,0} minimum(%convert_element_type.1271, %broadcast_in_dim.1152), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/min" stack_frame_id=160} + %broadcast_in_dim.1157 = s32[1,1,1,41]{3,2,1,0} reshape(%min.71), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/broadcast_in_dim" stack_frame_id=341} + %convert_element_type.1281 = pred[1,1,1,41]{3,2,1,0} compare(%broadcast_in_dim.1157, %broadcast.37), direction=NE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/convert_element_type" stack_frame_id=345} + %broadcast_in_dim.1158 = pred[1,1,1,1,41]{4,3,2,1,0} reshape(%convert_element_type.1281), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/broadcast_in_dim" stack_frame_id=83} + %and.221 = pred[1,1,1,1,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1158), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/vmap()/and" stack_frame_id=83} + %and.222 = pred[1,1,1,41]{3,2,1,0} reshape(%and.221), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/vmap()/and" stack_frame_id=83} + %and.223 = pred[1,1,32,1,41]{4,3,2,1,0} broadcast(%and.222), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/vmap()/and" stack_frame_id=83} + %dot_general.785 = bf16[1,1,32,128]{3,2,1,0} dot(%convert_element_type.1274, %arg_tuple.5#81), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=186} + %convert_element_type.1275 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/convert_element_type" stack_frame_id=208} + %add.807 = f32[1,1]{1,0} reshape(%convert_element_type.1275), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/add"} + %iota.416 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/iota" stack_frame_id=197} + %mul.2355 = f32[64]{0} multiply(%iota.416, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=197} + %div.724 = f32[64]{0} divide(%mul.2355, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/div" stack_frame_id=198} + %neg.294 = f32[64]{0} negate(%div.724), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/neg" stack_frame_id=199} + %pow.288 = f32[64]{0} power(%broadcast.39, %neg.294), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/pow" stack_frame_id=202} + %div.725 = f32[64]{0} divide(%pow.288, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/div" stack_frame_id=203} + %dot_general.786 = f32[1,1,64]{2,1,0} dot(%add.807, %div.725), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/bi,j->bij/dot_general" stack_frame_id=211} + %stack.1140 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.786), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/stack" stack_frame_id=214} + %stack.1141 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.786), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/stack" stack_frame_id=214} + %stack.1142 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1140, %stack.1141), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/stack" stack_frame_id=214} + %reshape.530 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1142), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/reshape"} + %cos.143 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.530), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/cos" stack_frame_id=218} + %convert_element_type.1276 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.143), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/convert_element_type" stack_frame_id=222} + %mul.2356 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1276), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=242} + %mul.2357 = bf16[1,1,128]{2,1,0} reshape(%mul.2356), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=242} + %mul.2358 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2357), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=242} + %mul.2359 = bf16[1,1,32,128]{3,2,1,0} multiply(%dot_general.785, %mul.2358), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=242} + %split.287 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.785), slice={[0:1], [0:1], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/split" stack_frame_id=234} + %neg.295 = bf16[1,1,32,64]{3,2,1,0} negate(%split.287), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/neg" stack_frame_id=235} + %stack.1143 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%neg.295), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/stack" stack_frame_id=238} + %split.286 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.785), slice={[0:1], [0:1], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/split" stack_frame_id=234} + %stack.1144 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%split.286), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/stack" stack_frame_id=238} + %stack.1145 = bf16[1,1,32,2,64]{4,3,2,1,0} concatenate(%stack.1143, %stack.1144), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/stack" stack_frame_id=238} + %reshape.531 = bf16[1,1,32,128]{3,2,1,0} reshape(%stack.1145), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/reshape" stack_frame_id=241} + %sin.143 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.530), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/sin" stack_frame_id=226} + %convert_element_type.1277 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.143), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/convert_element_type" stack_frame_id=230} + %mul.2360 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1277), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=243} + %mul.2361 = bf16[1,1,128]{2,1,0} reshape(%mul.2360), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=243} + %mul.2362 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2361), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=243} + %mul.2363 = bf16[1,1,32,128]{3,2,1,0} multiply(%reshape.531, %mul.2362), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=243} + %add.808 = bf16[1,1,32,128]{3,2,1,0} add(%mul.2359, %mul.2363), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/add" stack_frame_id=244} + %reshape.536 = bf16[1,1,32,1,128]{4,3,2,1,0} reshape(%add.808), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/reshape" stack_frame_id=83} + %slice.121 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.124), slice={[0:1], [0:1], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/slice" stack_frame_id=245} + %squeeze.125 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/squeeze" stack_frame_id=245} + %dot_general.787 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1274, %arg_tuple.5#82), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=256} + %convert_element_type.1278 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/convert_element_type" stack_frame_id=287} + %add.809 = f32[1,1]{1,0} reshape(%convert_element_type.1278), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/add"} + %iota.417 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/iota" stack_frame_id=276} + %mul.2364 = f32[64]{0} multiply(%iota.417, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=276} + %div.726 = f32[64]{0} divide(%mul.2364, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/div" stack_frame_id=277} + %neg.296 = f32[64]{0} negate(%div.726), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/neg" stack_frame_id=278} + %pow.289 = f32[64]{0} power(%broadcast.39, %neg.296), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/pow" stack_frame_id=281} + %div.727 = f32[64]{0} divide(%pow.289, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/div" stack_frame_id=282} + %dot_general.789 = f32[1,1,64]{2,1,0} dot(%add.809, %div.727), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/bi,j->bij/dot_general" stack_frame_id=290} + %stack.1146 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.789), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/stack" stack_frame_id=293} + %stack.1147 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.789), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/stack" stack_frame_id=293} + %stack.1148 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1146, %stack.1147), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/stack" stack_frame_id=293} + %reshape.532 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1148), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/reshape"} + %cos.144 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.532), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/cos" stack_frame_id=297} + %convert_element_type.1279 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.144), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/convert_element_type" stack_frame_id=301} + %mul.2365 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1279), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=321} + %mul.2366 = bf16[1,1,128]{2,1,0} reshape(%mul.2365), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=321} + %mul.2367 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2366), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=321} + %mul.2368 = bf16[1,1,8,128]{3,2,1,0} multiply(%dot_general.787, %mul.2367), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=321} + %split.289 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.787), slice={[0:1], [0:1], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/split" stack_frame_id=313} + %neg.297 = bf16[1,1,8,64]{3,2,1,0} negate(%split.289), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/neg" stack_frame_id=314} + %stack.1149 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%neg.297), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/stack" stack_frame_id=317} + %split.288 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.787), slice={[0:1], [0:1], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/split" stack_frame_id=313} + %stack.1150 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%split.288), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/stack" stack_frame_id=317} + %stack.1151 = bf16[1,1,8,2,64]{4,3,2,1,0} concatenate(%stack.1149, %stack.1150), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/stack" stack_frame_id=317} + %reshape.533 = bf16[1,1,8,128]{3,2,1,0} reshape(%stack.1151), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/reshape" stack_frame_id=320} + %sin.144 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.532), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/sin" stack_frame_id=305} + %convert_element_type.1280 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.144), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/convert_element_type" stack_frame_id=309} + %mul.2369 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1280), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=322} + %mul.2370 = bf16[1,1,128]{2,1,0} reshape(%mul.2369), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=322} + %mul.2371 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2370), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=322} + %mul.2372 = bf16[1,1,8,128]{3,2,1,0} multiply(%reshape.533, %mul.2371), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=322} + %add.810 = bf16[1,1,8,128]{3,2,1,0} add(%mul.2368, %mul.2372), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/rotary_embedding_8/add" stack_frame_id=323} + %lt.478 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/lt" stack_frame_id=326} + %add.811 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/add" stack_frame_id=326} + %select_n.107 = s32[] select(%lt.478, %add.811, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/select_n" stack_frame_id=326} + %dynamic_update_slice.81 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.125, %add.810, %constant.127, %select_n.107, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/dynamic_update_slice" stack_frame_id=326} + %broadcast_in_dim.1155 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.81), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/broadcast_in_dim" stack_frame_id=335} + %reshape.534 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1155), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/reshape" stack_frame_id=335} + %dot_general.790 = f32[1,32,1,1,41]{4,3,2,1,0} dot(%reshape.536, %reshape.534), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=83} + %mul.2373 = f32[1,32,1,1,41]{4,3,2,1,0} multiply(%dot_general.790, %broadcast.36), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/vmap()/mul" stack_frame_id=83} + %vmap_jit__where__.71 = f32[1,1,32,1,41]{4,3,2,1,0} call(%and.223, %mul.2373, %constant.114), to_apply=%_where_3.137, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/vmap(jit(_where))" stack_frame_id=83} + %reduce_max.456 = f32[1,1,32,1]{3,2,1,0} reduce(%vmap_jit__where__.71, %constant.113), dimensions={4}, to_apply=%region_163.171, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/vmap()/reduce_max" stack_frame_id=83} + %max.80 = f32[1,1,32,1]{3,2,1,0} maximum(%reduce_max.456, %broadcast.35), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/vmap()/max" stack_frame_id=83} + %broadcast_in_dim.1159 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%max.80), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %sub.303 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1159), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/vmap()/sub" stack_frame_id=83} + %sub.304 = f32[1,1,32,1]{3,2,1,0} reshape(%sub.303), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/vmap()/sub" stack_frame_id=83} + %sub.305 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%sub.304), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/vmap()/sub" stack_frame_id=83} + %sub.306 = f32[1,1,32,1,41]{4,3,2,1,0} subtract(%vmap_jit__where__.71, %sub.305), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/vmap()/sub" stack_frame_id=83} + %exp.76 = f32[1,1,32,1,41]{4,3,2,1,0} exponential(%sub.306), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/vmap()/exp" stack_frame_id=83} + %reduce_sum.1376 = f32[1,1,32,1]{3,2,1,0} reduce(%exp.76, %constant.123), dimensions={4}, to_apply=%region_164.172, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/vmap()/reduce_sum" stack_frame_id=83} + %broadcast_in_dim.1160 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%reduce_sum.1376), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %div.728 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1160), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/vmap()/div" stack_frame_id=83} + %div.729 = f32[1,1,32,1]{3,2,1,0} reshape(%div.728), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/vmap()/div" stack_frame_id=83} + %div.730 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%div.729), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/vmap()/div" stack_frame_id=83} + %div.731 = f32[1,1,32,1,41]{4,3,2,1,0} divide(%exp.76, %div.730), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/vmap()/div" stack_frame_id=83} + %convert_element_type.1282 = bf16[1,1,32,1,41]{4,3,2,1,0} convert(%div.731), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/vmap()/convert_element_type" stack_frame_id=83} + %dot_general.791 = bf16[1,32,128,1,1]{4,3,2,1,0} dot(%reshape.535, %convert_element_type.1282), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=83} + %transpose.75 = bf16[1,1,32,128]{3,2,1,0} reshape(%dot_general.791), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/reshape;jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/vmap()/transpose;jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %dot_general.792 = bf16[1,1,4096]{2,1,0} dot(%transpose.75, %arg_tuple.5#84), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=354} + %add.813 = bf16[1,1,4096]{2,1,0} add(%dot_general.792, %add.802), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/add" stack_frame_id=355} + %convert_element_type.1283 = f32[1,1,4096]{2,1,0} convert(%add.813), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/feedforward_layernorm/convert_element_type" stack_frame_id=359} + %pow.290 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1283, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/feedforward_layernorm/pow" stack_frame_id=362} + %reduce_sum.1377 = f32[1,1]{1,0} reduce(%pow.290, %constant.123), dimensions={2}, to_apply=%region_165.173, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/feedforward_layernorm/reduce_sum" stack_frame_id=92} + %broadcast_in_dim.1161 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1377), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/feedforward_layernorm/broadcast_in_dim" stack_frame_id=92} + %div.732 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1161, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/feedforward_layernorm/div" stack_frame_id=92} + %add.814 = f32[1,1,1]{2,1,0} add(%div.732, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/feedforward_layernorm/add" stack_frame_id=363} + %rsqrt.145 = f32[1,1,1]{2,1,0} rsqrt(%add.814), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/feedforward_layernorm/rsqrt" stack_frame_id=366} + %mul.2374 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.145), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2375 = f32[1,1]{1,0} reshape(%mul.2374), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2376 = f32[1,1,4096]{2,1,0} broadcast(%mul.2375), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2377 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1283, %mul.2376), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/feedforward_layernorm/mul" stack_frame_id=367} + %convert_element_type.1284 = f32[4096]{0} convert(%arg_tuple.5#85), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/feedforward_layernorm/convert_element_type" stack_frame_id=368} + %broadcast_in_dim.1162 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1284), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/feedforward_layernorm/broadcast_in_dim" stack_frame_id=368} + %mul.2378 = f32[1,1,4096]{2,1,0} multiply(%mul.2377, %broadcast_in_dim.1162), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/feedforward_layernorm/mul" stack_frame_id=368} + %convert_element_type.1285 = bf16[1,1,4096]{2,1,0} convert(%mul.2378), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/feedforward_layernorm/convert_element_type" stack_frame_id=372} + %dot_general.794 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1285, %arg_tuple.5#87), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/feedforward_intermediate_dense/dot_general" stack_frame_id=401} + %dot_general.793 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1285, %arg_tuple.5#86), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/feedforward_gate_dense/dot_general" stack_frame_id=381} + %convert_element_type.1286 = f32[1,1,14336]{2,1,0} convert(%dot_general.793), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/convert_element_type" stack_frame_id=385} + %jit_silu_.71 = f32[1,1,14336]{2,1,0} call(%convert_element_type.1286), to_apply=%silu_4.141, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/jit(silu)" stack_frame_id=388} + %convert_element_type.1287 = bf16[1,1,14336]{2,1,0} convert(%jit_silu_.71), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/convert_element_type" stack_frame_id=392} + %mul.2379 = bf16[1,1,14336]{2,1,0} multiply(%dot_general.794, %convert_element_type.1287), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/mul" stack_frame_id=404} + %dot_general.795 = bf16[1,1,4096]{2,1,0} dot(%mul.2379, %arg_tuple.5#88), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/feedforward_output_dense/dot_general" stack_frame_id=413} + %add.815 = bf16[1,1,4096]{2,1,0} add(%dot_general.795, %add.813), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/add" stack_frame_id=414} + %convert_element_type.1291 = f32[1,1,4096]{2,1,0} convert(%add.815), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention_layernorm/convert_element_type" stack_frame_id=164} + %pow.291 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1291, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention_layernorm/pow" stack_frame_id=167} + %reduce_sum.1378 = f32[1,1]{1,0} reduce(%pow.291, %constant.123), dimensions={2}, to_apply=%region_166.174, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention_layernorm/reduce_sum" stack_frame_id=73} + %broadcast_in_dim.1167 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1378), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention_layernorm/broadcast_in_dim" stack_frame_id=73} + %div.733 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1167, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention_layernorm/div" stack_frame_id=73} + %add.819 = f32[1,1,1]{2,1,0} add(%div.733, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention_layernorm/add" stack_frame_id=168} + %rsqrt.146 = f32[1,1,1]{2,1,0} rsqrt(%add.819), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention_layernorm/rsqrt" stack_frame_id=171} + %mul.2380 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.146), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2381 = f32[1,1]{1,0} reshape(%mul.2380), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2382 = f32[1,1,4096]{2,1,0} broadcast(%mul.2381), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2383 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1291, %mul.2382), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention_layernorm/mul" stack_frame_id=172} + %convert_element_type.1292 = f32[4096]{0} convert(%arg_tuple.5#89), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention_layernorm/convert_element_type" stack_frame_id=173} + %broadcast_in_dim.1168 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1292), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention_layernorm/broadcast_in_dim" stack_frame_id=173} + %mul.2384 = f32[1,1,4096]{2,1,0} multiply(%mul.2383, %broadcast_in_dim.1168), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention_layernorm/mul" stack_frame_id=173} + %convert_element_type.1293 = bf16[1,1,4096]{2,1,0} convert(%mul.2384), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention_layernorm/convert_element_type" stack_frame_id=177} + %dot_general.799 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1293, %arg_tuple.5#92), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=265} + %lt.485 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/lt" stack_frame_id=329} + %add.825 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/add" stack_frame_id=329} + %select_n.110 = s32[] select(%lt.485, %add.825, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/select_n" stack_frame_id=329} + %dynamic_update_slice.84 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.129, %dot_general.799, %constant.127, %select_n.110, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/dynamic_update_slice" stack_frame_id=329} + %broadcast_in_dim.1170 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.84), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/broadcast_in_dim" stack_frame_id=338} + %reshape.542 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1170), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/reshape" stack_frame_id=338} + %convert_element_type.1288 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/convert_element_type" stack_frame_id=139} + %add.816 = f32[1,1]{1,0} reshape(%convert_element_type.1288), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_9/add"} + %ge.389 = f32[1,1]{1,0} broadcast(%add.816), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/ge" stack_frame_id=143} + %ge.390 = f32[1]{0} reshape(%ge.389), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/ge" stack_frame_id=143} + %ge.391 = f32[1,41]{1,0} broadcast(%ge.390), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/ge" stack_frame_id=143} + %iota.418 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/iota" stack_frame_id=142} + %broadcast_in_dim.1163 = f32[1,41]{1,0} reshape(%iota.418), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/broadcast_in_dim" stack_frame_id=143} + %ge.392 = pred[1,41]{1,0} compare(%ge.391, %broadcast_in_dim.1163), direction=GE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/ge" stack_frame_id=143} + %broadcast_in_dim.1164 = pred[1,1,41]{2,1,0} reshape(%ge.392), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/broadcast_in_dim" stack_frame_id=146} + %convert_element_type.1290 = s32[1,1,41]{2,1,0} convert(%broadcast_in_dim.1164), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/convert_element_type" stack_frame_id=160} + %add.817 = s32[1,1]{1,0} reshape(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/add" stack_frame_id=147} + %lt.480 = s32[1,1]{1,0} broadcast(%add.817), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/lt" stack_frame_id=152} + %lt.481 = s32[1]{0} reshape(%lt.480), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/lt" stack_frame_id=152} + %lt.482 = s32[1,41]{1,0} broadcast(%lt.481), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/lt" stack_frame_id=152} + %iota.419 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/iota" stack_frame_id=150} + %broadcast_in_dim.1165 = s32[1,41]{1,0} reshape(%iota.419), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/broadcast_in_dim" stack_frame_id=148} + %add.818 = s32[1,41]{1,0} add(%broadcast_in_dim.1165, %broadcast.43), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/add" stack_frame_id=151} + %lt.483 = pred[1,41]{1,0} compare(%lt.482, %add.818), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/lt" stack_frame_id=152} + %convert_element_type.1289 = s32[1,41]{1,0} convert(%lt.483), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/convert_element_type" stack_frame_id=156} + %broadcast_in_dim.1166 = s32[1,1,41]{2,1,0} reshape(%convert_element_type.1289), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/broadcast_in_dim" stack_frame_id=160} + %min.72 = s32[1,1,41]{2,1,0} minimum(%convert_element_type.1290, %broadcast_in_dim.1166), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/min" stack_frame_id=160} + %broadcast_in_dim.1171 = s32[1,1,1,41]{3,2,1,0} reshape(%min.72), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/broadcast_in_dim" stack_frame_id=341} + %convert_element_type.1300 = pred[1,1,1,41]{3,2,1,0} compare(%broadcast_in_dim.1171, %broadcast.37), direction=NE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/convert_element_type" stack_frame_id=345} + %broadcast_in_dim.1172 = pred[1,1,1,1,41]{4,3,2,1,0} reshape(%convert_element_type.1300), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/broadcast_in_dim" stack_frame_id=83} + %and.224 = pred[1,1,1,1,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1172), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/vmap()/and" stack_frame_id=83} + %and.225 = pred[1,1,1,41]{3,2,1,0} reshape(%and.224), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/vmap()/and" stack_frame_id=83} + %and.226 = pred[1,1,32,1,41]{4,3,2,1,0} broadcast(%and.225), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/vmap()/and" stack_frame_id=83} + %dot_general.796 = bf16[1,1,32,128]{3,2,1,0} dot(%convert_element_type.1293, %arg_tuple.5#90), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=186} + %convert_element_type.1294 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/convert_element_type" stack_frame_id=208} + %add.820 = f32[1,1]{1,0} reshape(%convert_element_type.1294), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/add"} + %iota.420 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/iota" stack_frame_id=197} + %mul.2385 = f32[64]{0} multiply(%iota.420, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=197} + %div.734 = f32[64]{0} divide(%mul.2385, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/div" stack_frame_id=198} + %neg.298 = f32[64]{0} negate(%div.734), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/neg" stack_frame_id=199} + %pow.292 = f32[64]{0} power(%broadcast.39, %neg.298), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/pow" stack_frame_id=202} + %div.735 = f32[64]{0} divide(%pow.292, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/div" stack_frame_id=203} + %dot_general.797 = f32[1,1,64]{2,1,0} dot(%add.820, %div.735), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/bi,j->bij/dot_general" stack_frame_id=211} + %stack.1155 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.797), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/stack" stack_frame_id=214} + %stack.1156 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.797), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/stack" stack_frame_id=214} + %stack.1157 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1155, %stack.1156), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/stack" stack_frame_id=214} + %reshape.537 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1157), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/reshape"} + %cos.145 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.537), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/cos" stack_frame_id=218} + %convert_element_type.1295 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.145), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/convert_element_type" stack_frame_id=222} + %mul.2386 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1295), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=242} + %mul.2387 = bf16[1,1,128]{2,1,0} reshape(%mul.2386), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=242} + %mul.2388 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2387), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=242} + %mul.2389 = bf16[1,1,32,128]{3,2,1,0} multiply(%dot_general.796, %mul.2388), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=242} + %split.291 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.796), slice={[0:1], [0:1], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/split" stack_frame_id=234} + %neg.299 = bf16[1,1,32,64]{3,2,1,0} negate(%split.291), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/neg" stack_frame_id=235} + %stack.1158 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%neg.299), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/stack" stack_frame_id=238} + %split.290 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.796), slice={[0:1], [0:1], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/split" stack_frame_id=234} + %stack.1159 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%split.290), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/stack" stack_frame_id=238} + %stack.1160 = bf16[1,1,32,2,64]{4,3,2,1,0} concatenate(%stack.1158, %stack.1159), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/stack" stack_frame_id=238} + %reshape.538 = bf16[1,1,32,128]{3,2,1,0} reshape(%stack.1160), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/reshape" stack_frame_id=241} + %sin.145 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.537), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/sin" stack_frame_id=226} + %convert_element_type.1296 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.145), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/convert_element_type" stack_frame_id=230} + %mul.2390 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1296), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=243} + %mul.2391 = bf16[1,1,128]{2,1,0} reshape(%mul.2390), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=243} + %mul.2392 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2391), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=243} + %mul.2393 = bf16[1,1,32,128]{3,2,1,0} multiply(%reshape.538, %mul.2392), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=243} + %add.821 = bf16[1,1,32,128]{3,2,1,0} add(%mul.2389, %mul.2393), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/add" stack_frame_id=244} + %reshape.543 = bf16[1,1,32,1,128]{4,3,2,1,0} reshape(%add.821), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/reshape" stack_frame_id=83} + %slice.124 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.127), slice={[0:1], [0:1], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/slice" stack_frame_id=245} + %squeeze.128 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.124), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/squeeze" stack_frame_id=245} + %dot_general.798 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1293, %arg_tuple.5#91), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=256} + %convert_element_type.1297 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/convert_element_type" stack_frame_id=287} + %add.822 = f32[1,1]{1,0} reshape(%convert_element_type.1297), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/add"} + %iota.421 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/iota" stack_frame_id=276} + %mul.2394 = f32[64]{0} multiply(%iota.421, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=276} + %div.736 = f32[64]{0} divide(%mul.2394, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/div" stack_frame_id=277} + %neg.300 = f32[64]{0} negate(%div.736), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/neg" stack_frame_id=278} + %pow.293 = f32[64]{0} power(%broadcast.39, %neg.300), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/pow" stack_frame_id=281} + %div.737 = f32[64]{0} divide(%pow.293, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/div" stack_frame_id=282} + %dot_general.800 = f32[1,1,64]{2,1,0} dot(%add.822, %div.737), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/bi,j->bij/dot_general" stack_frame_id=290} + %stack.1161 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.800), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/stack" stack_frame_id=293} + %stack.1162 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.800), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/stack" stack_frame_id=293} + %stack.1163 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1161, %stack.1162), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/stack" stack_frame_id=293} + %reshape.539 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1163), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/reshape"} + %cos.146 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.539), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/cos" stack_frame_id=297} + %convert_element_type.1298 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.146), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/convert_element_type" stack_frame_id=301} + %mul.2395 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1298), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=321} + %mul.2396 = bf16[1,1,128]{2,1,0} reshape(%mul.2395), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=321} + %mul.2397 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2396), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=321} + %mul.2398 = bf16[1,1,8,128]{3,2,1,0} multiply(%dot_general.798, %mul.2397), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=321} + %split.293 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.798), slice={[0:1], [0:1], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/split" stack_frame_id=313} + %neg.301 = bf16[1,1,8,64]{3,2,1,0} negate(%split.293), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/neg" stack_frame_id=314} + %stack.1164 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%neg.301), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/stack" stack_frame_id=317} + %split.292 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.798), slice={[0:1], [0:1], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/split" stack_frame_id=313} + %stack.1165 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%split.292), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/stack" stack_frame_id=317} + %stack.1166 = bf16[1,1,8,2,64]{4,3,2,1,0} concatenate(%stack.1164, %stack.1165), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/stack" stack_frame_id=317} + %reshape.540 = bf16[1,1,8,128]{3,2,1,0} reshape(%stack.1166), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/reshape" stack_frame_id=320} + %sin.146 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.539), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/sin" stack_frame_id=305} + %convert_element_type.1299 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.146), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/convert_element_type" stack_frame_id=309} + %mul.2399 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1299), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=322} + %mul.2400 = bf16[1,1,128]{2,1,0} reshape(%mul.2399), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=322} + %mul.2401 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2400), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=322} + %mul.2402 = bf16[1,1,8,128]{3,2,1,0} multiply(%reshape.540, %mul.2401), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=322} + %add.823 = bf16[1,1,8,128]{3,2,1,0} add(%mul.2398, %mul.2402), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/rotary_embedding_9/add" stack_frame_id=323} + %lt.484 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/lt" stack_frame_id=326} + %add.824 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/add" stack_frame_id=326} + %select_n.109 = s32[] select(%lt.484, %add.824, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/select_n" stack_frame_id=326} + %dynamic_update_slice.83 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.128, %add.823, %constant.127, %select_n.109, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/dynamic_update_slice" stack_frame_id=326} + %broadcast_in_dim.1169 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.83), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/broadcast_in_dim" stack_frame_id=335} + %reshape.541 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1169), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/reshape" stack_frame_id=335} + %dot_general.801 = f32[1,32,1,1,41]{4,3,2,1,0} dot(%reshape.543, %reshape.541), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=83} + %mul.2403 = f32[1,32,1,1,41]{4,3,2,1,0} multiply(%dot_general.801, %broadcast.36), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/vmap()/mul" stack_frame_id=83} + %vmap_jit__where__.72 = f32[1,1,32,1,41]{4,3,2,1,0} call(%and.226, %mul.2403, %constant.114), to_apply=%_where_3.137, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/vmap(jit(_where))" stack_frame_id=83} + %reduce_max.457 = f32[1,1,32,1]{3,2,1,0} reduce(%vmap_jit__where__.72, %constant.113), dimensions={4}, to_apply=%region_167.175, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/vmap()/reduce_max" stack_frame_id=83} + %max.81 = f32[1,1,32,1]{3,2,1,0} maximum(%reduce_max.457, %broadcast.35), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/vmap()/max" stack_frame_id=83} + %broadcast_in_dim.1173 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%max.81), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %sub.307 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1173), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/vmap()/sub" stack_frame_id=83} + %sub.308 = f32[1,1,32,1]{3,2,1,0} reshape(%sub.307), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/vmap()/sub" stack_frame_id=83} + %sub.309 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%sub.308), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/vmap()/sub" stack_frame_id=83} + %sub.310 = f32[1,1,32,1,41]{4,3,2,1,0} subtract(%vmap_jit__where__.72, %sub.309), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/vmap()/sub" stack_frame_id=83} + %exp.77 = f32[1,1,32,1,41]{4,3,2,1,0} exponential(%sub.310), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/vmap()/exp" stack_frame_id=83} + %reduce_sum.1379 = f32[1,1,32,1]{3,2,1,0} reduce(%exp.77, %constant.123), dimensions={4}, to_apply=%region_168.176, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/vmap()/reduce_sum" stack_frame_id=83} + %broadcast_in_dim.1174 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%reduce_sum.1379), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %div.738 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1174), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/vmap()/div" stack_frame_id=83} + %div.739 = f32[1,1,32,1]{3,2,1,0} reshape(%div.738), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/vmap()/div" stack_frame_id=83} + %div.740 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%div.739), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/vmap()/div" stack_frame_id=83} + %div.741 = f32[1,1,32,1,41]{4,3,2,1,0} divide(%exp.77, %div.740), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/vmap()/div" stack_frame_id=83} + %convert_element_type.1301 = bf16[1,1,32,1,41]{4,3,2,1,0} convert(%div.741), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/vmap()/convert_element_type" stack_frame_id=83} + %dot_general.802 = bf16[1,32,128,1,1]{4,3,2,1,0} dot(%reshape.542, %convert_element_type.1301), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=83} + %transpose.76 = bf16[1,1,32,128]{3,2,1,0} reshape(%dot_general.802), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/reshape;jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/vmap()/transpose;jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %dot_general.803 = bf16[1,1,4096]{2,1,0} dot(%transpose.76, %arg_tuple.5#93), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=354} + %add.826 = bf16[1,1,4096]{2,1,0} add(%dot_general.803, %add.815), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/add" stack_frame_id=355} + %convert_element_type.1302 = f32[1,1,4096]{2,1,0} convert(%add.826), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/feedforward_layernorm/convert_element_type" stack_frame_id=359} + %pow.294 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1302, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/feedforward_layernorm/pow" stack_frame_id=362} + %reduce_sum.1380 = f32[1,1]{1,0} reduce(%pow.294, %constant.123), dimensions={2}, to_apply=%region_169.177, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/feedforward_layernorm/reduce_sum" stack_frame_id=92} + %broadcast_in_dim.1175 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1380), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/feedforward_layernorm/broadcast_in_dim" stack_frame_id=92} + %div.742 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1175, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/feedforward_layernorm/div" stack_frame_id=92} + %add.827 = f32[1,1,1]{2,1,0} add(%div.742, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/feedforward_layernorm/add" stack_frame_id=363} + %rsqrt.147 = f32[1,1,1]{2,1,0} rsqrt(%add.827), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/feedforward_layernorm/rsqrt" stack_frame_id=366} + %mul.2404 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.147), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2405 = f32[1,1]{1,0} reshape(%mul.2404), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2406 = f32[1,1,4096]{2,1,0} broadcast(%mul.2405), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2407 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1302, %mul.2406), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/feedforward_layernorm/mul" stack_frame_id=367} + %convert_element_type.1303 = f32[4096]{0} convert(%arg_tuple.5#94), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/feedforward_layernorm/convert_element_type" stack_frame_id=368} + %broadcast_in_dim.1176 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1303), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/feedforward_layernorm/broadcast_in_dim" stack_frame_id=368} + %mul.2408 = f32[1,1,4096]{2,1,0} multiply(%mul.2407, %broadcast_in_dim.1176), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/feedforward_layernorm/mul" stack_frame_id=368} + %convert_element_type.1304 = bf16[1,1,4096]{2,1,0} convert(%mul.2408), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/feedforward_layernorm/convert_element_type" stack_frame_id=372} + %dot_general.805 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1304, %arg_tuple.5#96), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/feedforward_intermediate_dense/dot_general" stack_frame_id=401} + %dot_general.804 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1304, %arg_tuple.5#95), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/feedforward_gate_dense/dot_general" stack_frame_id=381} + %convert_element_type.1305 = f32[1,1,14336]{2,1,0} convert(%dot_general.804), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/convert_element_type" stack_frame_id=385} + %jit_silu_.72 = f32[1,1,14336]{2,1,0} call(%convert_element_type.1305), to_apply=%silu_4.141, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/jit(silu)" stack_frame_id=388} + %convert_element_type.1306 = bf16[1,1,14336]{2,1,0} convert(%jit_silu_.72), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/convert_element_type" stack_frame_id=392} + %mul.2409 = bf16[1,1,14336]{2,1,0} multiply(%dot_general.805, %convert_element_type.1306), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/mul" stack_frame_id=404} + %dot_general.806 = bf16[1,1,4096]{2,1,0} dot(%mul.2409, %arg_tuple.5#97), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/feedforward_output_dense/dot_general" stack_frame_id=413} + %add.828 = bf16[1,1,4096]{2,1,0} add(%dot_general.806, %add.826), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/add" stack_frame_id=414} + %convert_element_type.1310 = f32[1,1,4096]{2,1,0} convert(%add.828), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention_layernorm/convert_element_type" stack_frame_id=164} + %pow.295 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1310, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention_layernorm/pow" stack_frame_id=167} + %reduce_sum.1381 = f32[1,1]{1,0} reduce(%pow.295, %constant.123), dimensions={2}, to_apply=%region_170.178, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention_layernorm/reduce_sum" stack_frame_id=73} + %broadcast_in_dim.1181 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1381), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention_layernorm/broadcast_in_dim" stack_frame_id=73} + %div.743 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1181, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention_layernorm/div" stack_frame_id=73} + %add.832 = f32[1,1,1]{2,1,0} add(%div.743, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention_layernorm/add" stack_frame_id=168} + %rsqrt.148 = f32[1,1,1]{2,1,0} rsqrt(%add.832), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention_layernorm/rsqrt" stack_frame_id=171} + %mul.2410 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.148), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2411 = f32[1,1]{1,0} reshape(%mul.2410), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2412 = f32[1,1,4096]{2,1,0} broadcast(%mul.2411), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2413 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1310, %mul.2412), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention_layernorm/mul" stack_frame_id=172} + %convert_element_type.1311 = f32[4096]{0} convert(%arg_tuple.5#98), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention_layernorm/convert_element_type" stack_frame_id=173} + %broadcast_in_dim.1182 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1311), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention_layernorm/broadcast_in_dim" stack_frame_id=173} + %mul.2414 = f32[1,1,4096]{2,1,0} multiply(%mul.2413, %broadcast_in_dim.1182), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention_layernorm/mul" stack_frame_id=173} + %convert_element_type.1312 = bf16[1,1,4096]{2,1,0} convert(%mul.2414), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention_layernorm/convert_element_type" stack_frame_id=177} + %dot_general.810 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1312, %arg_tuple.5#101), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=265} + %lt.491 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/lt" stack_frame_id=329} + %add.838 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/add" stack_frame_id=329} + %select_n.112 = s32[] select(%lt.491, %add.838, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/select_n" stack_frame_id=329} + %dynamic_update_slice.86 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.132, %dot_general.810, %constant.127, %select_n.112, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/dynamic_update_slice" stack_frame_id=329} + %broadcast_in_dim.1184 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.86), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/broadcast_in_dim" stack_frame_id=338} + %reshape.549 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1184), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/reshape" stack_frame_id=338} + %convert_element_type.1307 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/convert_element_type" stack_frame_id=139} + %add.829 = f32[1,1]{1,0} reshape(%convert_element_type.1307), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_10/add"} + %ge.393 = f32[1,1]{1,0} broadcast(%add.829), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/ge" stack_frame_id=143} + %ge.394 = f32[1]{0} reshape(%ge.393), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/ge" stack_frame_id=143} + %ge.395 = f32[1,41]{1,0} broadcast(%ge.394), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/ge" stack_frame_id=143} + %iota.422 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/iota" stack_frame_id=142} + %broadcast_in_dim.1177 = f32[1,41]{1,0} reshape(%iota.422), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/broadcast_in_dim" stack_frame_id=143} + %ge.396 = pred[1,41]{1,0} compare(%ge.395, %broadcast_in_dim.1177), direction=GE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/ge" stack_frame_id=143} + %broadcast_in_dim.1178 = pred[1,1,41]{2,1,0} reshape(%ge.396), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/broadcast_in_dim" stack_frame_id=146} + %convert_element_type.1309 = s32[1,1,41]{2,1,0} convert(%broadcast_in_dim.1178), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/convert_element_type" stack_frame_id=160} + %add.830 = s32[1,1]{1,0} reshape(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/add" stack_frame_id=147} + %lt.486 = s32[1,1]{1,0} broadcast(%add.830), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/lt" stack_frame_id=152} + %lt.487 = s32[1]{0} reshape(%lt.486), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/lt" stack_frame_id=152} + %lt.488 = s32[1,41]{1,0} broadcast(%lt.487), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/lt" stack_frame_id=152} + %iota.423 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/iota" stack_frame_id=150} + %broadcast_in_dim.1179 = s32[1,41]{1,0} reshape(%iota.423), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/broadcast_in_dim" stack_frame_id=148} + %add.831 = s32[1,41]{1,0} add(%broadcast_in_dim.1179, %broadcast.43), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/add" stack_frame_id=151} + %lt.489 = pred[1,41]{1,0} compare(%lt.488, %add.831), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/lt" stack_frame_id=152} + %convert_element_type.1308 = s32[1,41]{1,0} convert(%lt.489), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/convert_element_type" stack_frame_id=156} + %broadcast_in_dim.1180 = s32[1,1,41]{2,1,0} reshape(%convert_element_type.1308), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/broadcast_in_dim" stack_frame_id=160} + %min.73 = s32[1,1,41]{2,1,0} minimum(%convert_element_type.1309, %broadcast_in_dim.1180), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/min" stack_frame_id=160} + %broadcast_in_dim.1185 = s32[1,1,1,41]{3,2,1,0} reshape(%min.73), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/broadcast_in_dim" stack_frame_id=341} + %convert_element_type.1319 = pred[1,1,1,41]{3,2,1,0} compare(%broadcast_in_dim.1185, %broadcast.37), direction=NE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/convert_element_type" stack_frame_id=345} + %broadcast_in_dim.1186 = pred[1,1,1,1,41]{4,3,2,1,0} reshape(%convert_element_type.1319), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/broadcast_in_dim" stack_frame_id=83} + %and.227 = pred[1,1,1,1,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1186), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/vmap()/and" stack_frame_id=83} + %and.228 = pred[1,1,1,41]{3,2,1,0} reshape(%and.227), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/vmap()/and" stack_frame_id=83} + %and.229 = pred[1,1,32,1,41]{4,3,2,1,0} broadcast(%and.228), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/vmap()/and" stack_frame_id=83} + %dot_general.807 = bf16[1,1,32,128]{3,2,1,0} dot(%convert_element_type.1312, %arg_tuple.5#99), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=186} + %convert_element_type.1313 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/convert_element_type" stack_frame_id=208} + %add.833 = f32[1,1]{1,0} reshape(%convert_element_type.1313), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/add"} + %iota.424 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/iota" stack_frame_id=197} + %mul.2415 = f32[64]{0} multiply(%iota.424, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=197} + %div.744 = f32[64]{0} divide(%mul.2415, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/div" stack_frame_id=198} + %neg.302 = f32[64]{0} negate(%div.744), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/neg" stack_frame_id=199} + %pow.296 = f32[64]{0} power(%broadcast.39, %neg.302), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/pow" stack_frame_id=202} + %div.745 = f32[64]{0} divide(%pow.296, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/div" stack_frame_id=203} + %dot_general.808 = f32[1,1,64]{2,1,0} dot(%add.833, %div.745), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/bi,j->bij/dot_general" stack_frame_id=211} + %stack.1170 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.808), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/stack" stack_frame_id=214} + %stack.1171 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.808), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/stack" stack_frame_id=214} + %stack.1172 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1170, %stack.1171), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/stack" stack_frame_id=214} + %reshape.544 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1172), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/reshape"} + %cos.147 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.544), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/cos" stack_frame_id=218} + %convert_element_type.1314 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.147), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/convert_element_type" stack_frame_id=222} + %mul.2416 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1314), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=242} + %mul.2417 = bf16[1,1,128]{2,1,0} reshape(%mul.2416), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=242} + %mul.2418 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2417), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=242} + %mul.2419 = bf16[1,1,32,128]{3,2,1,0} multiply(%dot_general.807, %mul.2418), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=242} + %split.295 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.807), slice={[0:1], [0:1], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/split" stack_frame_id=234} + %neg.303 = bf16[1,1,32,64]{3,2,1,0} negate(%split.295), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/neg" stack_frame_id=235} + %stack.1173 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%neg.303), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/stack" stack_frame_id=238} + %split.294 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.807), slice={[0:1], [0:1], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/split" stack_frame_id=234} + %stack.1174 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%split.294), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/stack" stack_frame_id=238} + %stack.1175 = bf16[1,1,32,2,64]{4,3,2,1,0} concatenate(%stack.1173, %stack.1174), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/stack" stack_frame_id=238} + %reshape.545 = bf16[1,1,32,128]{3,2,1,0} reshape(%stack.1175), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/reshape" stack_frame_id=241} + %sin.147 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.544), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/sin" stack_frame_id=226} + %convert_element_type.1315 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.147), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/convert_element_type" stack_frame_id=230} + %mul.2420 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1315), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=243} + %mul.2421 = bf16[1,1,128]{2,1,0} reshape(%mul.2420), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=243} + %mul.2422 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2421), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=243} + %mul.2423 = bf16[1,1,32,128]{3,2,1,0} multiply(%reshape.545, %mul.2422), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=243} + %add.834 = bf16[1,1,32,128]{3,2,1,0} add(%mul.2419, %mul.2423), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/add" stack_frame_id=244} + %reshape.550 = bf16[1,1,32,1,128]{4,3,2,1,0} reshape(%add.834), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/reshape" stack_frame_id=83} + %slice.127 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.130), slice={[0:1], [0:1], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/slice" stack_frame_id=245} + %squeeze.131 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/squeeze" stack_frame_id=245} + %dot_general.809 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1312, %arg_tuple.5#100), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=256} + %convert_element_type.1316 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/convert_element_type" stack_frame_id=287} + %add.835 = f32[1,1]{1,0} reshape(%convert_element_type.1316), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/add"} + %iota.425 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/iota" stack_frame_id=276} + %mul.2424 = f32[64]{0} multiply(%iota.425, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=276} + %div.746 = f32[64]{0} divide(%mul.2424, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/div" stack_frame_id=277} + %neg.304 = f32[64]{0} negate(%div.746), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/neg" stack_frame_id=278} + %pow.297 = f32[64]{0} power(%broadcast.39, %neg.304), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/pow" stack_frame_id=281} + %div.747 = f32[64]{0} divide(%pow.297, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/div" stack_frame_id=282} + %dot_general.811 = f32[1,1,64]{2,1,0} dot(%add.835, %div.747), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/bi,j->bij/dot_general" stack_frame_id=290} + %stack.1176 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.811), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/stack" stack_frame_id=293} + %stack.1177 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.811), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/stack" stack_frame_id=293} + %stack.1178 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1176, %stack.1177), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/stack" stack_frame_id=293} + %reshape.546 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1178), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/reshape"} + %cos.148 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.546), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/cos" stack_frame_id=297} + %convert_element_type.1317 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.148), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/convert_element_type" stack_frame_id=301} + %mul.2425 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1317), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=321} + %mul.2426 = bf16[1,1,128]{2,1,0} reshape(%mul.2425), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=321} + %mul.2427 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2426), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=321} + %mul.2428 = bf16[1,1,8,128]{3,2,1,0} multiply(%dot_general.809, %mul.2427), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=321} + %split.297 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.809), slice={[0:1], [0:1], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/split" stack_frame_id=313} + %neg.305 = bf16[1,1,8,64]{3,2,1,0} negate(%split.297), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/neg" stack_frame_id=314} + %stack.1179 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%neg.305), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/stack" stack_frame_id=317} + %split.296 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.809), slice={[0:1], [0:1], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/split" stack_frame_id=313} + %stack.1180 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%split.296), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/stack" stack_frame_id=317} + %stack.1181 = bf16[1,1,8,2,64]{4,3,2,1,0} concatenate(%stack.1179, %stack.1180), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/stack" stack_frame_id=317} + %reshape.547 = bf16[1,1,8,128]{3,2,1,0} reshape(%stack.1181), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/reshape" stack_frame_id=320} + %sin.148 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.546), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/sin" stack_frame_id=305} + %convert_element_type.1318 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.148), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/convert_element_type" stack_frame_id=309} + %mul.2429 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1318), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=322} + %mul.2430 = bf16[1,1,128]{2,1,0} reshape(%mul.2429), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=322} + %mul.2431 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2430), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=322} + %mul.2432 = bf16[1,1,8,128]{3,2,1,0} multiply(%reshape.547, %mul.2431), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=322} + %add.836 = bf16[1,1,8,128]{3,2,1,0} add(%mul.2428, %mul.2432), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/rotary_embedding_10/add" stack_frame_id=323} + %lt.490 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/lt" stack_frame_id=326} + %add.837 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/add" stack_frame_id=326} + %select_n.111 = s32[] select(%lt.490, %add.837, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/select_n" stack_frame_id=326} + %dynamic_update_slice.85 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.131, %add.836, %constant.127, %select_n.111, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/dynamic_update_slice" stack_frame_id=326} + %broadcast_in_dim.1183 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.85), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/broadcast_in_dim" stack_frame_id=335} + %reshape.548 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1183), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/reshape" stack_frame_id=335} + %dot_general.812 = f32[1,32,1,1,41]{4,3,2,1,0} dot(%reshape.550, %reshape.548), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=83} + %mul.2433 = f32[1,32,1,1,41]{4,3,2,1,0} multiply(%dot_general.812, %broadcast.36), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/vmap()/mul" stack_frame_id=83} + %vmap_jit__where__.73 = f32[1,1,32,1,41]{4,3,2,1,0} call(%and.229, %mul.2433, %constant.114), to_apply=%_where_3.137, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/vmap(jit(_where))" stack_frame_id=83} + %reduce_max.458 = f32[1,1,32,1]{3,2,1,0} reduce(%vmap_jit__where__.73, %constant.113), dimensions={4}, to_apply=%region_171.179, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/vmap()/reduce_max" stack_frame_id=83} + %max.82 = f32[1,1,32,1]{3,2,1,0} maximum(%reduce_max.458, %broadcast.35), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/vmap()/max" stack_frame_id=83} + %broadcast_in_dim.1187 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%max.82), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %sub.311 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1187), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/vmap()/sub" stack_frame_id=83} + %sub.312 = f32[1,1,32,1]{3,2,1,0} reshape(%sub.311), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/vmap()/sub" stack_frame_id=83} + %sub.313 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%sub.312), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/vmap()/sub" stack_frame_id=83} + %sub.314 = f32[1,1,32,1,41]{4,3,2,1,0} subtract(%vmap_jit__where__.73, %sub.313), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/vmap()/sub" stack_frame_id=83} + %exp.78 = f32[1,1,32,1,41]{4,3,2,1,0} exponential(%sub.314), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/vmap()/exp" stack_frame_id=83} + %reduce_sum.1382 = f32[1,1,32,1]{3,2,1,0} reduce(%exp.78, %constant.123), dimensions={4}, to_apply=%region_172.180, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/vmap()/reduce_sum" stack_frame_id=83} + %broadcast_in_dim.1188 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%reduce_sum.1382), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %div.748 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1188), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/vmap()/div" stack_frame_id=83} + %div.749 = f32[1,1,32,1]{3,2,1,0} reshape(%div.748), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/vmap()/div" stack_frame_id=83} + %div.750 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%div.749), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/vmap()/div" stack_frame_id=83} + %div.751 = f32[1,1,32,1,41]{4,3,2,1,0} divide(%exp.78, %div.750), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/vmap()/div" stack_frame_id=83} + %convert_element_type.1320 = bf16[1,1,32,1,41]{4,3,2,1,0} convert(%div.751), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/vmap()/convert_element_type" stack_frame_id=83} + %dot_general.813 = bf16[1,32,128,1,1]{4,3,2,1,0} dot(%reshape.549, %convert_element_type.1320), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=83} + %transpose.77 = bf16[1,1,32,128]{3,2,1,0} reshape(%dot_general.813), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/reshape;jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/vmap()/transpose;jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %dot_general.814 = bf16[1,1,4096]{2,1,0} dot(%transpose.77, %arg_tuple.5#102), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=354} + %add.839 = bf16[1,1,4096]{2,1,0} add(%dot_general.814, %add.828), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/add" stack_frame_id=355} + %convert_element_type.1321 = f32[1,1,4096]{2,1,0} convert(%add.839), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/feedforward_layernorm/convert_element_type" stack_frame_id=359} + %pow.298 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1321, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/feedforward_layernorm/pow" stack_frame_id=362} + %reduce_sum.1383 = f32[1,1]{1,0} reduce(%pow.298, %constant.123), dimensions={2}, to_apply=%region_173.181, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/feedforward_layernorm/reduce_sum" stack_frame_id=92} + %broadcast_in_dim.1189 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1383), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/feedforward_layernorm/broadcast_in_dim" stack_frame_id=92} + %div.752 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1189, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/feedforward_layernorm/div" stack_frame_id=92} + %add.840 = f32[1,1,1]{2,1,0} add(%div.752, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/feedforward_layernorm/add" stack_frame_id=363} + %rsqrt.149 = f32[1,1,1]{2,1,0} rsqrt(%add.840), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/feedforward_layernorm/rsqrt" stack_frame_id=366} + %mul.2434 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.149), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2435 = f32[1,1]{1,0} reshape(%mul.2434), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2436 = f32[1,1,4096]{2,1,0} broadcast(%mul.2435), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2437 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1321, %mul.2436), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/feedforward_layernorm/mul" stack_frame_id=367} + %convert_element_type.1322 = f32[4096]{0} convert(%arg_tuple.5#103), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/feedforward_layernorm/convert_element_type" stack_frame_id=368} + %broadcast_in_dim.1190 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1322), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/feedforward_layernorm/broadcast_in_dim" stack_frame_id=368} + %mul.2438 = f32[1,1,4096]{2,1,0} multiply(%mul.2437, %broadcast_in_dim.1190), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/feedforward_layernorm/mul" stack_frame_id=368} + %convert_element_type.1323 = bf16[1,1,4096]{2,1,0} convert(%mul.2438), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/feedforward_layernorm/convert_element_type" stack_frame_id=372} + %dot_general.816 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1323, %arg_tuple.5#105), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/feedforward_intermediate_dense/dot_general" stack_frame_id=401} + %dot_general.815 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1323, %arg_tuple.5#104), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/feedforward_gate_dense/dot_general" stack_frame_id=381} + %convert_element_type.1324 = f32[1,1,14336]{2,1,0} convert(%dot_general.815), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/convert_element_type" stack_frame_id=385} + %jit_silu_.73 = f32[1,1,14336]{2,1,0} call(%convert_element_type.1324), to_apply=%silu_4.141, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/jit(silu)" stack_frame_id=388} + %convert_element_type.1325 = bf16[1,1,14336]{2,1,0} convert(%jit_silu_.73), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/convert_element_type" stack_frame_id=392} + %mul.2439 = bf16[1,1,14336]{2,1,0} multiply(%dot_general.816, %convert_element_type.1325), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/mul" stack_frame_id=404} + %dot_general.817 = bf16[1,1,4096]{2,1,0} dot(%mul.2439, %arg_tuple.5#106), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/feedforward_output_dense/dot_general" stack_frame_id=413} + %add.841 = bf16[1,1,4096]{2,1,0} add(%dot_general.817, %add.839), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/add" stack_frame_id=414} + %convert_element_type.1329 = f32[1,1,4096]{2,1,0} convert(%add.841), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention_layernorm/convert_element_type" stack_frame_id=164} + %pow.299 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1329, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention_layernorm/pow" stack_frame_id=167} + %reduce_sum.1384 = f32[1,1]{1,0} reduce(%pow.299, %constant.123), dimensions={2}, to_apply=%region_174.182, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention_layernorm/reduce_sum" stack_frame_id=73} + %broadcast_in_dim.1195 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1384), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention_layernorm/broadcast_in_dim" stack_frame_id=73} + %div.753 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1195, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention_layernorm/div" stack_frame_id=73} + %add.845 = f32[1,1,1]{2,1,0} add(%div.753, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention_layernorm/add" stack_frame_id=168} + %rsqrt.150 = f32[1,1,1]{2,1,0} rsqrt(%add.845), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention_layernorm/rsqrt" stack_frame_id=171} + %mul.2440 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.150), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2441 = f32[1,1]{1,0} reshape(%mul.2440), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2442 = f32[1,1,4096]{2,1,0} broadcast(%mul.2441), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2443 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1329, %mul.2442), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention_layernorm/mul" stack_frame_id=172} + %convert_element_type.1330 = f32[4096]{0} convert(%arg_tuple.5#107), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention_layernorm/convert_element_type" stack_frame_id=173} + %broadcast_in_dim.1196 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1330), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention_layernorm/broadcast_in_dim" stack_frame_id=173} + %mul.2444 = f32[1,1,4096]{2,1,0} multiply(%mul.2443, %broadcast_in_dim.1196), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention_layernorm/mul" stack_frame_id=173} + %convert_element_type.1331 = bf16[1,1,4096]{2,1,0} convert(%mul.2444), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention_layernorm/convert_element_type" stack_frame_id=177} + %dot_general.821 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1331, %arg_tuple.5#110), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=265} + %lt.497 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/lt" stack_frame_id=329} + %add.851 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/add" stack_frame_id=329} + %select_n.114 = s32[] select(%lt.497, %add.851, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/select_n" stack_frame_id=329} + %dynamic_update_slice.88 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.135, %dot_general.821, %constant.127, %select_n.114, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/dynamic_update_slice" stack_frame_id=329} + %broadcast_in_dim.1198 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.88), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/broadcast_in_dim" stack_frame_id=338} + %reshape.556 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1198), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/reshape" stack_frame_id=338} + %convert_element_type.1326 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/convert_element_type" stack_frame_id=139} + %add.842 = f32[1,1]{1,0} reshape(%convert_element_type.1326), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_11/add"} + %ge.397 = f32[1,1]{1,0} broadcast(%add.842), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/ge" stack_frame_id=143} + %ge.398 = f32[1]{0} reshape(%ge.397), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/ge" stack_frame_id=143} + %ge.399 = f32[1,41]{1,0} broadcast(%ge.398), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/ge" stack_frame_id=143} + %iota.426 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/iota" stack_frame_id=142} + %broadcast_in_dim.1191 = f32[1,41]{1,0} reshape(%iota.426), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/broadcast_in_dim" stack_frame_id=143} + %ge.400 = pred[1,41]{1,0} compare(%ge.399, %broadcast_in_dim.1191), direction=GE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/ge" stack_frame_id=143} + %broadcast_in_dim.1192 = pred[1,1,41]{2,1,0} reshape(%ge.400), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/broadcast_in_dim" stack_frame_id=146} + %convert_element_type.1328 = s32[1,1,41]{2,1,0} convert(%broadcast_in_dim.1192), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/convert_element_type" stack_frame_id=160} + %add.843 = s32[1,1]{1,0} reshape(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/add" stack_frame_id=147} + %lt.492 = s32[1,1]{1,0} broadcast(%add.843), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/lt" stack_frame_id=152} + %lt.493 = s32[1]{0} reshape(%lt.492), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/lt" stack_frame_id=152} + %lt.494 = s32[1,41]{1,0} broadcast(%lt.493), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/lt" stack_frame_id=152} + %iota.427 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/iota" stack_frame_id=150} + %broadcast_in_dim.1193 = s32[1,41]{1,0} reshape(%iota.427), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/broadcast_in_dim" stack_frame_id=148} + %add.844 = s32[1,41]{1,0} add(%broadcast_in_dim.1193, %broadcast.43), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/add" stack_frame_id=151} + %lt.495 = pred[1,41]{1,0} compare(%lt.494, %add.844), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/lt" stack_frame_id=152} + %convert_element_type.1327 = s32[1,41]{1,0} convert(%lt.495), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/convert_element_type" stack_frame_id=156} + %broadcast_in_dim.1194 = s32[1,1,41]{2,1,0} reshape(%convert_element_type.1327), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/broadcast_in_dim" stack_frame_id=160} + %min.74 = s32[1,1,41]{2,1,0} minimum(%convert_element_type.1328, %broadcast_in_dim.1194), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/min" stack_frame_id=160} + %broadcast_in_dim.1199 = s32[1,1,1,41]{3,2,1,0} reshape(%min.74), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/broadcast_in_dim" stack_frame_id=341} + %convert_element_type.1338 = pred[1,1,1,41]{3,2,1,0} compare(%broadcast_in_dim.1199, %broadcast.37), direction=NE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/convert_element_type" stack_frame_id=345} + %broadcast_in_dim.1200 = pred[1,1,1,1,41]{4,3,2,1,0} reshape(%convert_element_type.1338), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/broadcast_in_dim" stack_frame_id=83} + %and.230 = pred[1,1,1,1,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1200), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/vmap()/and" stack_frame_id=83} + %and.231 = pred[1,1,1,41]{3,2,1,0} reshape(%and.230), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/vmap()/and" stack_frame_id=83} + %and.232 = pred[1,1,32,1,41]{4,3,2,1,0} broadcast(%and.231), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/vmap()/and" stack_frame_id=83} + %dot_general.818 = bf16[1,1,32,128]{3,2,1,0} dot(%convert_element_type.1331, %arg_tuple.5#108), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=186} + %convert_element_type.1332 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/convert_element_type" stack_frame_id=208} + %add.846 = f32[1,1]{1,0} reshape(%convert_element_type.1332), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/add"} + %iota.428 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/iota" stack_frame_id=197} + %mul.2445 = f32[64]{0} multiply(%iota.428, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=197} + %div.754 = f32[64]{0} divide(%mul.2445, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/div" stack_frame_id=198} + %neg.306 = f32[64]{0} negate(%div.754), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/neg" stack_frame_id=199} + %pow.300 = f32[64]{0} power(%broadcast.39, %neg.306), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/pow" stack_frame_id=202} + %div.755 = f32[64]{0} divide(%pow.300, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/div" stack_frame_id=203} + %dot_general.819 = f32[1,1,64]{2,1,0} dot(%add.846, %div.755), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/bi,j->bij/dot_general" stack_frame_id=211} + %stack.1185 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.819), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/stack" stack_frame_id=214} + %stack.1186 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.819), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/stack" stack_frame_id=214} + %stack.1187 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1185, %stack.1186), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/stack" stack_frame_id=214} + %reshape.551 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1187), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/reshape"} + %cos.149 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.551), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/cos" stack_frame_id=218} + %convert_element_type.1333 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.149), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/convert_element_type" stack_frame_id=222} + %mul.2446 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1333), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=242} + %mul.2447 = bf16[1,1,128]{2,1,0} reshape(%mul.2446), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=242} + %mul.2448 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2447), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=242} + %mul.2449 = bf16[1,1,32,128]{3,2,1,0} multiply(%dot_general.818, %mul.2448), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=242} + %split.299 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.818), slice={[0:1], [0:1], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/split" stack_frame_id=234} + %neg.307 = bf16[1,1,32,64]{3,2,1,0} negate(%split.299), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/neg" stack_frame_id=235} + %stack.1188 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%neg.307), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/stack" stack_frame_id=238} + %split.298 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.818), slice={[0:1], [0:1], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/split" stack_frame_id=234} + %stack.1189 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%split.298), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/stack" stack_frame_id=238} + %stack.1190 = bf16[1,1,32,2,64]{4,3,2,1,0} concatenate(%stack.1188, %stack.1189), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/stack" stack_frame_id=238} + %reshape.552 = bf16[1,1,32,128]{3,2,1,0} reshape(%stack.1190), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/reshape" stack_frame_id=241} + %sin.149 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.551), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/sin" stack_frame_id=226} + %convert_element_type.1334 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.149), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/convert_element_type" stack_frame_id=230} + %mul.2450 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1334), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=243} + %mul.2451 = bf16[1,1,128]{2,1,0} reshape(%mul.2450), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=243} + %mul.2452 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2451), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=243} + %mul.2453 = bf16[1,1,32,128]{3,2,1,0} multiply(%reshape.552, %mul.2452), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=243} + %add.847 = bf16[1,1,32,128]{3,2,1,0} add(%mul.2449, %mul.2453), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/add" stack_frame_id=244} + %reshape.557 = bf16[1,1,32,1,128]{4,3,2,1,0} reshape(%add.847), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/reshape" stack_frame_id=83} + %slice.130 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.133), slice={[0:1], [0:1], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/slice" stack_frame_id=245} + %squeeze.134 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.130), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/squeeze" stack_frame_id=245} + %dot_general.820 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1331, %arg_tuple.5#109), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=256} + %convert_element_type.1335 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/convert_element_type" stack_frame_id=287} + %add.848 = f32[1,1]{1,0} reshape(%convert_element_type.1335), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/add"} + %iota.429 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/iota" stack_frame_id=276} + %mul.2454 = f32[64]{0} multiply(%iota.429, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=276} + %div.756 = f32[64]{0} divide(%mul.2454, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/div" stack_frame_id=277} + %neg.308 = f32[64]{0} negate(%div.756), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/neg" stack_frame_id=278} + %pow.301 = f32[64]{0} power(%broadcast.39, %neg.308), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/pow" stack_frame_id=281} + %div.757 = f32[64]{0} divide(%pow.301, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/div" stack_frame_id=282} + %dot_general.822 = f32[1,1,64]{2,1,0} dot(%add.848, %div.757), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/bi,j->bij/dot_general" stack_frame_id=290} + %stack.1191 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.822), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/stack" stack_frame_id=293} + %stack.1192 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.822), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/stack" stack_frame_id=293} + %stack.1193 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1191, %stack.1192), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/stack" stack_frame_id=293} + %reshape.553 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1193), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/reshape"} + %cos.150 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.553), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/cos" stack_frame_id=297} + %convert_element_type.1336 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.150), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/convert_element_type" stack_frame_id=301} + %mul.2455 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1336), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=321} + %mul.2456 = bf16[1,1,128]{2,1,0} reshape(%mul.2455), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=321} + %mul.2457 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2456), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=321} + %mul.2458 = bf16[1,1,8,128]{3,2,1,0} multiply(%dot_general.820, %mul.2457), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=321} + %split.301 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.820), slice={[0:1], [0:1], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/split" stack_frame_id=313} + %neg.309 = bf16[1,1,8,64]{3,2,1,0} negate(%split.301), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/neg" stack_frame_id=314} + %stack.1194 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%neg.309), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/stack" stack_frame_id=317} + %split.300 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.820), slice={[0:1], [0:1], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/split" stack_frame_id=313} + %stack.1195 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%split.300), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/stack" stack_frame_id=317} + %stack.1196 = bf16[1,1,8,2,64]{4,3,2,1,0} concatenate(%stack.1194, %stack.1195), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/stack" stack_frame_id=317} + %reshape.554 = bf16[1,1,8,128]{3,2,1,0} reshape(%stack.1196), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/reshape" stack_frame_id=320} + %sin.150 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.553), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/sin" stack_frame_id=305} + %convert_element_type.1337 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.150), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/convert_element_type" stack_frame_id=309} + %mul.2459 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1337), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=322} + %mul.2460 = bf16[1,1,128]{2,1,0} reshape(%mul.2459), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=322} + %mul.2461 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2460), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=322} + %mul.2462 = bf16[1,1,8,128]{3,2,1,0} multiply(%reshape.554, %mul.2461), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=322} + %add.849 = bf16[1,1,8,128]{3,2,1,0} add(%mul.2458, %mul.2462), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/rotary_embedding_11/add" stack_frame_id=323} + %lt.496 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/lt" stack_frame_id=326} + %add.850 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/add" stack_frame_id=326} + %select_n.113 = s32[] select(%lt.496, %add.850, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/select_n" stack_frame_id=326} + %dynamic_update_slice.87 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.134, %add.849, %constant.127, %select_n.113, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/dynamic_update_slice" stack_frame_id=326} + %broadcast_in_dim.1197 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.87), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/broadcast_in_dim" stack_frame_id=335} + %reshape.555 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1197), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/reshape" stack_frame_id=335} + %dot_general.823 = f32[1,32,1,1,41]{4,3,2,1,0} dot(%reshape.557, %reshape.555), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=83} + %mul.2463 = f32[1,32,1,1,41]{4,3,2,1,0} multiply(%dot_general.823, %broadcast.36), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/vmap()/mul" stack_frame_id=83} + %vmap_jit__where__.74 = f32[1,1,32,1,41]{4,3,2,1,0} call(%and.232, %mul.2463, %constant.114), to_apply=%_where_3.137, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/vmap(jit(_where))" stack_frame_id=83} + %reduce_max.459 = f32[1,1,32,1]{3,2,1,0} reduce(%vmap_jit__where__.74, %constant.113), dimensions={4}, to_apply=%region_175.183, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/vmap()/reduce_max" stack_frame_id=83} + %max.83 = f32[1,1,32,1]{3,2,1,0} maximum(%reduce_max.459, %broadcast.35), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/vmap()/max" stack_frame_id=83} + %broadcast_in_dim.1201 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%max.83), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %sub.315 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1201), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/vmap()/sub" stack_frame_id=83} + %sub.316 = f32[1,1,32,1]{3,2,1,0} reshape(%sub.315), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/vmap()/sub" stack_frame_id=83} + %sub.317 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%sub.316), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/vmap()/sub" stack_frame_id=83} + %sub.318 = f32[1,1,32,1,41]{4,3,2,1,0} subtract(%vmap_jit__where__.74, %sub.317), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/vmap()/sub" stack_frame_id=83} + %exp.79 = f32[1,1,32,1,41]{4,3,2,1,0} exponential(%sub.318), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/vmap()/exp" stack_frame_id=83} + %reduce_sum.1385 = f32[1,1,32,1]{3,2,1,0} reduce(%exp.79, %constant.123), dimensions={4}, to_apply=%region_176.184, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/vmap()/reduce_sum" stack_frame_id=83} + %broadcast_in_dim.1202 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%reduce_sum.1385), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %div.758 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1202), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/vmap()/div" stack_frame_id=83} + %div.759 = f32[1,1,32,1]{3,2,1,0} reshape(%div.758), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/vmap()/div" stack_frame_id=83} + %div.760 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%div.759), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/vmap()/div" stack_frame_id=83} + %div.761 = f32[1,1,32,1,41]{4,3,2,1,0} divide(%exp.79, %div.760), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/vmap()/div" stack_frame_id=83} + %convert_element_type.1339 = bf16[1,1,32,1,41]{4,3,2,1,0} convert(%div.761), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/vmap()/convert_element_type" stack_frame_id=83} + %dot_general.824 = bf16[1,32,128,1,1]{4,3,2,1,0} dot(%reshape.556, %convert_element_type.1339), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=83} + %transpose.78 = bf16[1,1,32,128]{3,2,1,0} reshape(%dot_general.824), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/reshape;jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/vmap()/transpose;jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %dot_general.825 = bf16[1,1,4096]{2,1,0} dot(%transpose.78, %arg_tuple.5#111), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=354} + %add.852 = bf16[1,1,4096]{2,1,0} add(%dot_general.825, %add.841), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/add" stack_frame_id=355} + %convert_element_type.1340 = f32[1,1,4096]{2,1,0} convert(%add.852), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/feedforward_layernorm/convert_element_type" stack_frame_id=359} + %pow.302 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1340, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/feedforward_layernorm/pow" stack_frame_id=362} + %reduce_sum.1386 = f32[1,1]{1,0} reduce(%pow.302, %constant.123), dimensions={2}, to_apply=%region_177.185, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/feedforward_layernorm/reduce_sum" stack_frame_id=92} + %broadcast_in_dim.1203 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1386), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/feedforward_layernorm/broadcast_in_dim" stack_frame_id=92} + %div.762 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1203, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/feedforward_layernorm/div" stack_frame_id=92} + %add.853 = f32[1,1,1]{2,1,0} add(%div.762, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/feedforward_layernorm/add" stack_frame_id=363} + %rsqrt.151 = f32[1,1,1]{2,1,0} rsqrt(%add.853), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/feedforward_layernorm/rsqrt" stack_frame_id=366} + %mul.2464 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.151), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2465 = f32[1,1]{1,0} reshape(%mul.2464), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2466 = f32[1,1,4096]{2,1,0} broadcast(%mul.2465), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2467 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1340, %mul.2466), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/feedforward_layernorm/mul" stack_frame_id=367} + %convert_element_type.1341 = f32[4096]{0} convert(%arg_tuple.5#112), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/feedforward_layernorm/convert_element_type" stack_frame_id=368} + %broadcast_in_dim.1204 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1341), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/feedforward_layernorm/broadcast_in_dim" stack_frame_id=368} + %mul.2468 = f32[1,1,4096]{2,1,0} multiply(%mul.2467, %broadcast_in_dim.1204), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/feedforward_layernorm/mul" stack_frame_id=368} + %convert_element_type.1342 = bf16[1,1,4096]{2,1,0} convert(%mul.2468), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/feedforward_layernorm/convert_element_type" stack_frame_id=372} + %dot_general.827 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1342, %arg_tuple.5#114), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/feedforward_intermediate_dense/dot_general" stack_frame_id=401} + %dot_general.826 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1342, %arg_tuple.5#113), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/feedforward_gate_dense/dot_general" stack_frame_id=381} + %convert_element_type.1343 = f32[1,1,14336]{2,1,0} convert(%dot_general.826), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/convert_element_type" stack_frame_id=385} + %jit_silu_.74 = f32[1,1,14336]{2,1,0} call(%convert_element_type.1343), to_apply=%silu_4.141, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/jit(silu)" stack_frame_id=388} + %convert_element_type.1344 = bf16[1,1,14336]{2,1,0} convert(%jit_silu_.74), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/convert_element_type" stack_frame_id=392} + %mul.2469 = bf16[1,1,14336]{2,1,0} multiply(%dot_general.827, %convert_element_type.1344), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/mul" stack_frame_id=404} + %dot_general.828 = bf16[1,1,4096]{2,1,0} dot(%mul.2469, %arg_tuple.5#115), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/feedforward_output_dense/dot_general" stack_frame_id=413} + %add.854 = bf16[1,1,4096]{2,1,0} add(%dot_general.828, %add.852), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/add" stack_frame_id=414} + %convert_element_type.1348 = f32[1,1,4096]{2,1,0} convert(%add.854), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention_layernorm/convert_element_type" stack_frame_id=164} + %pow.303 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1348, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention_layernorm/pow" stack_frame_id=167} + %reduce_sum.1387 = f32[1,1]{1,0} reduce(%pow.303, %constant.123), dimensions={2}, to_apply=%region_178.186, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention_layernorm/reduce_sum" stack_frame_id=73} + %broadcast_in_dim.1209 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1387), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention_layernorm/broadcast_in_dim" stack_frame_id=73} + %div.763 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1209, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention_layernorm/div" stack_frame_id=73} + %add.858 = f32[1,1,1]{2,1,0} add(%div.763, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention_layernorm/add" stack_frame_id=168} + %rsqrt.152 = f32[1,1,1]{2,1,0} rsqrt(%add.858), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention_layernorm/rsqrt" stack_frame_id=171} + %mul.2470 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.152), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2471 = f32[1,1]{1,0} reshape(%mul.2470), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2472 = f32[1,1,4096]{2,1,0} broadcast(%mul.2471), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2473 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1348, %mul.2472), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention_layernorm/mul" stack_frame_id=172} + %convert_element_type.1349 = f32[4096]{0} convert(%arg_tuple.5#116), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention_layernorm/convert_element_type" stack_frame_id=173} + %broadcast_in_dim.1210 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1349), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention_layernorm/broadcast_in_dim" stack_frame_id=173} + %mul.2474 = f32[1,1,4096]{2,1,0} multiply(%mul.2473, %broadcast_in_dim.1210), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention_layernorm/mul" stack_frame_id=173} + %convert_element_type.1350 = bf16[1,1,4096]{2,1,0} convert(%mul.2474), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention_layernorm/convert_element_type" stack_frame_id=177} + %dot_general.832 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1350, %arg_tuple.5#119), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=265} + %lt.503 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/lt" stack_frame_id=329} + %add.864 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/add" stack_frame_id=329} + %select_n.116 = s32[] select(%lt.503, %add.864, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/select_n" stack_frame_id=329} + %dynamic_update_slice.90 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.138, %dot_general.832, %constant.127, %select_n.116, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/dynamic_update_slice" stack_frame_id=329} + %broadcast_in_dim.1212 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.90), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/broadcast_in_dim" stack_frame_id=338} + %reshape.563 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1212), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/reshape" stack_frame_id=338} + %convert_element_type.1345 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/convert_element_type" stack_frame_id=139} + %add.855 = f32[1,1]{1,0} reshape(%convert_element_type.1345), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_12/add"} + %ge.401 = f32[1,1]{1,0} broadcast(%add.855), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/ge" stack_frame_id=143} + %ge.402 = f32[1]{0} reshape(%ge.401), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/ge" stack_frame_id=143} + %ge.403 = f32[1,41]{1,0} broadcast(%ge.402), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/ge" stack_frame_id=143} + %iota.430 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/iota" stack_frame_id=142} + %broadcast_in_dim.1205 = f32[1,41]{1,0} reshape(%iota.430), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/broadcast_in_dim" stack_frame_id=143} + %ge.404 = pred[1,41]{1,0} compare(%ge.403, %broadcast_in_dim.1205), direction=GE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/ge" stack_frame_id=143} + %broadcast_in_dim.1206 = pred[1,1,41]{2,1,0} reshape(%ge.404), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/broadcast_in_dim" stack_frame_id=146} + %convert_element_type.1347 = s32[1,1,41]{2,1,0} convert(%broadcast_in_dim.1206), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/convert_element_type" stack_frame_id=160} + %add.856 = s32[1,1]{1,0} reshape(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/add" stack_frame_id=147} + %lt.498 = s32[1,1]{1,0} broadcast(%add.856), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/lt" stack_frame_id=152} + %lt.499 = s32[1]{0} reshape(%lt.498), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/lt" stack_frame_id=152} + %lt.500 = s32[1,41]{1,0} broadcast(%lt.499), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/lt" stack_frame_id=152} + %iota.431 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/iota" stack_frame_id=150} + %broadcast_in_dim.1207 = s32[1,41]{1,0} reshape(%iota.431), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/broadcast_in_dim" stack_frame_id=148} + %add.857 = s32[1,41]{1,0} add(%broadcast_in_dim.1207, %broadcast.43), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/add" stack_frame_id=151} + %lt.501 = pred[1,41]{1,0} compare(%lt.500, %add.857), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/lt" stack_frame_id=152} + %convert_element_type.1346 = s32[1,41]{1,0} convert(%lt.501), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/convert_element_type" stack_frame_id=156} + %broadcast_in_dim.1208 = s32[1,1,41]{2,1,0} reshape(%convert_element_type.1346), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/broadcast_in_dim" stack_frame_id=160} + %min.75 = s32[1,1,41]{2,1,0} minimum(%convert_element_type.1347, %broadcast_in_dim.1208), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/min" stack_frame_id=160} + %broadcast_in_dim.1213 = s32[1,1,1,41]{3,2,1,0} reshape(%min.75), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/broadcast_in_dim" stack_frame_id=341} + %convert_element_type.1357 = pred[1,1,1,41]{3,2,1,0} compare(%broadcast_in_dim.1213, %broadcast.37), direction=NE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/convert_element_type" stack_frame_id=345} + %broadcast_in_dim.1214 = pred[1,1,1,1,41]{4,3,2,1,0} reshape(%convert_element_type.1357), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/broadcast_in_dim" stack_frame_id=83} + %and.233 = pred[1,1,1,1,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1214), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/vmap()/and" stack_frame_id=83} + %and.234 = pred[1,1,1,41]{3,2,1,0} reshape(%and.233), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/vmap()/and" stack_frame_id=83} + %and.235 = pred[1,1,32,1,41]{4,3,2,1,0} broadcast(%and.234), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/vmap()/and" stack_frame_id=83} + %dot_general.829 = bf16[1,1,32,128]{3,2,1,0} dot(%convert_element_type.1350, %arg_tuple.5#117), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=186} + %convert_element_type.1351 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/convert_element_type" stack_frame_id=208} + %add.859 = f32[1,1]{1,0} reshape(%convert_element_type.1351), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/add"} + %iota.432 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/iota" stack_frame_id=197} + %mul.2475 = f32[64]{0} multiply(%iota.432, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=197} + %div.764 = f32[64]{0} divide(%mul.2475, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/div" stack_frame_id=198} + %neg.310 = f32[64]{0} negate(%div.764), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/neg" stack_frame_id=199} + %pow.304 = f32[64]{0} power(%broadcast.39, %neg.310), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/pow" stack_frame_id=202} + %div.765 = f32[64]{0} divide(%pow.304, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/div" stack_frame_id=203} + %dot_general.830 = f32[1,1,64]{2,1,0} dot(%add.859, %div.765), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/bi,j->bij/dot_general" stack_frame_id=211} + %stack.1200 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.830), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/stack" stack_frame_id=214} + %stack.1201 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.830), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/stack" stack_frame_id=214} + %stack.1202 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1200, %stack.1201), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/stack" stack_frame_id=214} + %reshape.558 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1202), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/reshape"} + %cos.151 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.558), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/cos" stack_frame_id=218} + %convert_element_type.1352 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.151), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/convert_element_type" stack_frame_id=222} + %mul.2476 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1352), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=242} + %mul.2477 = bf16[1,1,128]{2,1,0} reshape(%mul.2476), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=242} + %mul.2478 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2477), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=242} + %mul.2479 = bf16[1,1,32,128]{3,2,1,0} multiply(%dot_general.829, %mul.2478), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=242} + %split.303 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.829), slice={[0:1], [0:1], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/split" stack_frame_id=234} + %neg.311 = bf16[1,1,32,64]{3,2,1,0} negate(%split.303), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/neg" stack_frame_id=235} + %stack.1203 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%neg.311), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/stack" stack_frame_id=238} + %split.302 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.829), slice={[0:1], [0:1], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/split" stack_frame_id=234} + %stack.1204 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%split.302), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/stack" stack_frame_id=238} + %stack.1205 = bf16[1,1,32,2,64]{4,3,2,1,0} concatenate(%stack.1203, %stack.1204), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/stack" stack_frame_id=238} + %reshape.559 = bf16[1,1,32,128]{3,2,1,0} reshape(%stack.1205), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/reshape" stack_frame_id=241} + %sin.151 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.558), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/sin" stack_frame_id=226} + %convert_element_type.1353 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.151), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/convert_element_type" stack_frame_id=230} + %mul.2480 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1353), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=243} + %mul.2481 = bf16[1,1,128]{2,1,0} reshape(%mul.2480), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=243} + %mul.2482 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2481), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=243} + %mul.2483 = bf16[1,1,32,128]{3,2,1,0} multiply(%reshape.559, %mul.2482), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=243} + %add.860 = bf16[1,1,32,128]{3,2,1,0} add(%mul.2479, %mul.2483), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/add" stack_frame_id=244} + %reshape.564 = bf16[1,1,32,1,128]{4,3,2,1,0} reshape(%add.860), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/reshape" stack_frame_id=83} + %slice.133 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.136), slice={[0:1], [0:1], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/slice" stack_frame_id=245} + %squeeze.137 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.133), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/squeeze" stack_frame_id=245} + %dot_general.831 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1350, %arg_tuple.5#118), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=256} + %convert_element_type.1354 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/convert_element_type" stack_frame_id=287} + %add.861 = f32[1,1]{1,0} reshape(%convert_element_type.1354), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/add"} + %iota.433 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/iota" stack_frame_id=276} + %mul.2484 = f32[64]{0} multiply(%iota.433, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=276} + %div.766 = f32[64]{0} divide(%mul.2484, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/div" stack_frame_id=277} + %neg.312 = f32[64]{0} negate(%div.766), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/neg" stack_frame_id=278} + %pow.305 = f32[64]{0} power(%broadcast.39, %neg.312), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/pow" stack_frame_id=281} + %div.767 = f32[64]{0} divide(%pow.305, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/div" stack_frame_id=282} + %dot_general.833 = f32[1,1,64]{2,1,0} dot(%add.861, %div.767), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/bi,j->bij/dot_general" stack_frame_id=290} + %stack.1206 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.833), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/stack" stack_frame_id=293} + %stack.1207 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.833), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/stack" stack_frame_id=293} + %stack.1208 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1206, %stack.1207), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/stack" stack_frame_id=293} + %reshape.560 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1208), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/reshape"} + %cos.152 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.560), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/cos" stack_frame_id=297} + %convert_element_type.1355 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.152), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/convert_element_type" stack_frame_id=301} + %mul.2485 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1355), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=321} + %mul.2486 = bf16[1,1,128]{2,1,0} reshape(%mul.2485), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=321} + %mul.2487 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2486), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=321} + %mul.2488 = bf16[1,1,8,128]{3,2,1,0} multiply(%dot_general.831, %mul.2487), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=321} + %split.305 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.831), slice={[0:1], [0:1], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/split" stack_frame_id=313} + %neg.313 = bf16[1,1,8,64]{3,2,1,0} negate(%split.305), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/neg" stack_frame_id=314} + %stack.1209 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%neg.313), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/stack" stack_frame_id=317} + %split.304 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.831), slice={[0:1], [0:1], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/split" stack_frame_id=313} + %stack.1210 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%split.304), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/stack" stack_frame_id=317} + %stack.1211 = bf16[1,1,8,2,64]{4,3,2,1,0} concatenate(%stack.1209, %stack.1210), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/stack" stack_frame_id=317} + %reshape.561 = bf16[1,1,8,128]{3,2,1,0} reshape(%stack.1211), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/reshape" stack_frame_id=320} + %sin.152 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.560), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/sin" stack_frame_id=305} + %convert_element_type.1356 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.152), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/convert_element_type" stack_frame_id=309} + %mul.2489 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1356), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=322} + %mul.2490 = bf16[1,1,128]{2,1,0} reshape(%mul.2489), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=322} + %mul.2491 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2490), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=322} + %mul.2492 = bf16[1,1,8,128]{3,2,1,0} multiply(%reshape.561, %mul.2491), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=322} + %add.862 = bf16[1,1,8,128]{3,2,1,0} add(%mul.2488, %mul.2492), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/rotary_embedding_12/add" stack_frame_id=323} + %lt.502 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/lt" stack_frame_id=326} + %add.863 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/add" stack_frame_id=326} + %select_n.115 = s32[] select(%lt.502, %add.863, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/select_n" stack_frame_id=326} + %dynamic_update_slice.89 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.137, %add.862, %constant.127, %select_n.115, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/dynamic_update_slice" stack_frame_id=326} + %broadcast_in_dim.1211 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.89), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/broadcast_in_dim" stack_frame_id=335} + %reshape.562 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1211), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/reshape" stack_frame_id=335} + %dot_general.834 = f32[1,32,1,1,41]{4,3,2,1,0} dot(%reshape.564, %reshape.562), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=83} + %mul.2493 = f32[1,32,1,1,41]{4,3,2,1,0} multiply(%dot_general.834, %broadcast.36), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/vmap()/mul" stack_frame_id=83} + %vmap_jit__where__.75 = f32[1,1,32,1,41]{4,3,2,1,0} call(%and.235, %mul.2493, %constant.114), to_apply=%_where_3.137, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/vmap(jit(_where))" stack_frame_id=83} + %reduce_max.460 = f32[1,1,32,1]{3,2,1,0} reduce(%vmap_jit__where__.75, %constant.113), dimensions={4}, to_apply=%region_179.187, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/vmap()/reduce_max" stack_frame_id=83} + %max.84 = f32[1,1,32,1]{3,2,1,0} maximum(%reduce_max.460, %broadcast.35), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/vmap()/max" stack_frame_id=83} + %broadcast_in_dim.1215 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%max.84), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %sub.319 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1215), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/vmap()/sub" stack_frame_id=83} + %sub.320 = f32[1,1,32,1]{3,2,1,0} reshape(%sub.319), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/vmap()/sub" stack_frame_id=83} + %sub.321 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%sub.320), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/vmap()/sub" stack_frame_id=83} + %sub.322 = f32[1,1,32,1,41]{4,3,2,1,0} subtract(%vmap_jit__where__.75, %sub.321), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/vmap()/sub" stack_frame_id=83} + %exp.80 = f32[1,1,32,1,41]{4,3,2,1,0} exponential(%sub.322), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/vmap()/exp" stack_frame_id=83} + %reduce_sum.1388 = f32[1,1,32,1]{3,2,1,0} reduce(%exp.80, %constant.123), dimensions={4}, to_apply=%region_180.188, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/vmap()/reduce_sum" stack_frame_id=83} + %broadcast_in_dim.1216 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%reduce_sum.1388), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %div.768 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1216), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/vmap()/div" stack_frame_id=83} + %div.769 = f32[1,1,32,1]{3,2,1,0} reshape(%div.768), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/vmap()/div" stack_frame_id=83} + %div.770 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%div.769), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/vmap()/div" stack_frame_id=83} + %div.771 = f32[1,1,32,1,41]{4,3,2,1,0} divide(%exp.80, %div.770), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/vmap()/div" stack_frame_id=83} + %convert_element_type.1358 = bf16[1,1,32,1,41]{4,3,2,1,0} convert(%div.771), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/vmap()/convert_element_type" stack_frame_id=83} + %dot_general.835 = bf16[1,32,128,1,1]{4,3,2,1,0} dot(%reshape.563, %convert_element_type.1358), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=83} + %transpose.79 = bf16[1,1,32,128]{3,2,1,0} reshape(%dot_general.835), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/reshape;jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/vmap()/transpose;jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %dot_general.836 = bf16[1,1,4096]{2,1,0} dot(%transpose.79, %arg_tuple.5#120), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=354} + %add.865 = bf16[1,1,4096]{2,1,0} add(%dot_general.836, %add.854), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/add" stack_frame_id=355} + %convert_element_type.1359 = f32[1,1,4096]{2,1,0} convert(%add.865), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/feedforward_layernorm/convert_element_type" stack_frame_id=359} + %pow.306 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1359, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/feedforward_layernorm/pow" stack_frame_id=362} + %reduce_sum.1389 = f32[1,1]{1,0} reduce(%pow.306, %constant.123), dimensions={2}, to_apply=%region_181.189, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/feedforward_layernorm/reduce_sum" stack_frame_id=92} + %broadcast_in_dim.1217 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1389), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/feedforward_layernorm/broadcast_in_dim" stack_frame_id=92} + %div.772 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1217, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/feedforward_layernorm/div" stack_frame_id=92} + %add.866 = f32[1,1,1]{2,1,0} add(%div.772, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/feedforward_layernorm/add" stack_frame_id=363} + %rsqrt.153 = f32[1,1,1]{2,1,0} rsqrt(%add.866), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/feedforward_layernorm/rsqrt" stack_frame_id=366} + %mul.2494 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.153), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2495 = f32[1,1]{1,0} reshape(%mul.2494), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2496 = f32[1,1,4096]{2,1,0} broadcast(%mul.2495), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2497 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1359, %mul.2496), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/feedforward_layernorm/mul" stack_frame_id=367} + %convert_element_type.1360 = f32[4096]{0} convert(%arg_tuple.5#121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/feedforward_layernorm/convert_element_type" stack_frame_id=368} + %broadcast_in_dim.1218 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1360), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/feedforward_layernorm/broadcast_in_dim" stack_frame_id=368} + %mul.2498 = f32[1,1,4096]{2,1,0} multiply(%mul.2497, %broadcast_in_dim.1218), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/feedforward_layernorm/mul" stack_frame_id=368} + %convert_element_type.1361 = bf16[1,1,4096]{2,1,0} convert(%mul.2498), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/feedforward_layernorm/convert_element_type" stack_frame_id=372} + %dot_general.838 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1361, %arg_tuple.5#123), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/feedforward_intermediate_dense/dot_general" stack_frame_id=401} + %dot_general.837 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1361, %arg_tuple.5#122), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/feedforward_gate_dense/dot_general" stack_frame_id=381} + %convert_element_type.1362 = f32[1,1,14336]{2,1,0} convert(%dot_general.837), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/convert_element_type" stack_frame_id=385} + %jit_silu_.75 = f32[1,1,14336]{2,1,0} call(%convert_element_type.1362), to_apply=%silu_4.141, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/jit(silu)" stack_frame_id=388} + %convert_element_type.1363 = bf16[1,1,14336]{2,1,0} convert(%jit_silu_.75), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/convert_element_type" stack_frame_id=392} + %mul.2499 = bf16[1,1,14336]{2,1,0} multiply(%dot_general.838, %convert_element_type.1363), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/mul" stack_frame_id=404} + %dot_general.839 = bf16[1,1,4096]{2,1,0} dot(%mul.2499, %arg_tuple.5#124), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/feedforward_output_dense/dot_general" stack_frame_id=413} + %add.867 = bf16[1,1,4096]{2,1,0} add(%dot_general.839, %add.865), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/add" stack_frame_id=414} + %convert_element_type.1367 = f32[1,1,4096]{2,1,0} convert(%add.867), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention_layernorm/convert_element_type" stack_frame_id=164} + %pow.307 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1367, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention_layernorm/pow" stack_frame_id=167} + %reduce_sum.1390 = f32[1,1]{1,0} reduce(%pow.307, %constant.123), dimensions={2}, to_apply=%region_182.190, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention_layernorm/reduce_sum" stack_frame_id=73} + %broadcast_in_dim.1223 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1390), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention_layernorm/broadcast_in_dim" stack_frame_id=73} + %div.773 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1223, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention_layernorm/div" stack_frame_id=73} + %add.871 = f32[1,1,1]{2,1,0} add(%div.773, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention_layernorm/add" stack_frame_id=168} + %rsqrt.154 = f32[1,1,1]{2,1,0} rsqrt(%add.871), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention_layernorm/rsqrt" stack_frame_id=171} + %mul.2500 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.154), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2501 = f32[1,1]{1,0} reshape(%mul.2500), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2502 = f32[1,1,4096]{2,1,0} broadcast(%mul.2501), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2503 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1367, %mul.2502), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention_layernorm/mul" stack_frame_id=172} + %convert_element_type.1368 = f32[4096]{0} convert(%arg_tuple.5#125), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention_layernorm/convert_element_type" stack_frame_id=173} + %broadcast_in_dim.1224 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1368), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention_layernorm/broadcast_in_dim" stack_frame_id=173} + %mul.2504 = f32[1,1,4096]{2,1,0} multiply(%mul.2503, %broadcast_in_dim.1224), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention_layernorm/mul" stack_frame_id=173} + %convert_element_type.1369 = bf16[1,1,4096]{2,1,0} convert(%mul.2504), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention_layernorm/convert_element_type" stack_frame_id=177} + %dot_general.843 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1369, %arg_tuple.5#128), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=265} + %lt.509 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/lt" stack_frame_id=329} + %add.877 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/add" stack_frame_id=329} + %select_n.118 = s32[] select(%lt.509, %add.877, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/select_n" stack_frame_id=329} + %dynamic_update_slice.92 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.141, %dot_general.843, %constant.127, %select_n.118, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/dynamic_update_slice" stack_frame_id=329} + %broadcast_in_dim.1226 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.92), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/broadcast_in_dim" stack_frame_id=338} + %reshape.570 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1226), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/reshape" stack_frame_id=338} + %convert_element_type.1364 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/convert_element_type" stack_frame_id=139} + %add.868 = f32[1,1]{1,0} reshape(%convert_element_type.1364), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_13/add"} + %ge.405 = f32[1,1]{1,0} broadcast(%add.868), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/ge" stack_frame_id=143} + %ge.406 = f32[1]{0} reshape(%ge.405), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/ge" stack_frame_id=143} + %ge.407 = f32[1,41]{1,0} broadcast(%ge.406), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/ge" stack_frame_id=143} + %iota.434 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/iota" stack_frame_id=142} + %broadcast_in_dim.1219 = f32[1,41]{1,0} reshape(%iota.434), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/broadcast_in_dim" stack_frame_id=143} + %ge.408 = pred[1,41]{1,0} compare(%ge.407, %broadcast_in_dim.1219), direction=GE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/ge" stack_frame_id=143} + %broadcast_in_dim.1220 = pred[1,1,41]{2,1,0} reshape(%ge.408), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/broadcast_in_dim" stack_frame_id=146} + %convert_element_type.1366 = s32[1,1,41]{2,1,0} convert(%broadcast_in_dim.1220), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/convert_element_type" stack_frame_id=160} + %add.869 = s32[1,1]{1,0} reshape(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/add" stack_frame_id=147} + %lt.504 = s32[1,1]{1,0} broadcast(%add.869), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/lt" stack_frame_id=152} + %lt.505 = s32[1]{0} reshape(%lt.504), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/lt" stack_frame_id=152} + %lt.506 = s32[1,41]{1,0} broadcast(%lt.505), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/lt" stack_frame_id=152} + %iota.435 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/iota" stack_frame_id=150} + %broadcast_in_dim.1221 = s32[1,41]{1,0} reshape(%iota.435), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/broadcast_in_dim" stack_frame_id=148} + %add.870 = s32[1,41]{1,0} add(%broadcast_in_dim.1221, %broadcast.43), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/add" stack_frame_id=151} + %lt.507 = pred[1,41]{1,0} compare(%lt.506, %add.870), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/lt" stack_frame_id=152} + %convert_element_type.1365 = s32[1,41]{1,0} convert(%lt.507), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/convert_element_type" stack_frame_id=156} + %broadcast_in_dim.1222 = s32[1,1,41]{2,1,0} reshape(%convert_element_type.1365), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/broadcast_in_dim" stack_frame_id=160} + %min.76 = s32[1,1,41]{2,1,0} minimum(%convert_element_type.1366, %broadcast_in_dim.1222), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/min" stack_frame_id=160} + %broadcast_in_dim.1227 = s32[1,1,1,41]{3,2,1,0} reshape(%min.76), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/broadcast_in_dim" stack_frame_id=341} + %convert_element_type.1376 = pred[1,1,1,41]{3,2,1,0} compare(%broadcast_in_dim.1227, %broadcast.37), direction=NE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/convert_element_type" stack_frame_id=345} + %broadcast_in_dim.1228 = pred[1,1,1,1,41]{4,3,2,1,0} reshape(%convert_element_type.1376), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/broadcast_in_dim" stack_frame_id=83} + %and.236 = pred[1,1,1,1,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1228), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/vmap()/and" stack_frame_id=83} + %and.237 = pred[1,1,1,41]{3,2,1,0} reshape(%and.236), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/vmap()/and" stack_frame_id=83} + %and.238 = pred[1,1,32,1,41]{4,3,2,1,0} broadcast(%and.237), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/vmap()/and" stack_frame_id=83} + %dot_general.840 = bf16[1,1,32,128]{3,2,1,0} dot(%convert_element_type.1369, %arg_tuple.5#126), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=186} + %convert_element_type.1370 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/convert_element_type" stack_frame_id=208} + %add.872 = f32[1,1]{1,0} reshape(%convert_element_type.1370), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/add"} + %iota.436 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/iota" stack_frame_id=197} + %mul.2505 = f32[64]{0} multiply(%iota.436, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=197} + %div.774 = f32[64]{0} divide(%mul.2505, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/div" stack_frame_id=198} + %neg.314 = f32[64]{0} negate(%div.774), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/neg" stack_frame_id=199} + %pow.308 = f32[64]{0} power(%broadcast.39, %neg.314), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/pow" stack_frame_id=202} + %div.775 = f32[64]{0} divide(%pow.308, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/div" stack_frame_id=203} + %dot_general.841 = f32[1,1,64]{2,1,0} dot(%add.872, %div.775), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/bi,j->bij/dot_general" stack_frame_id=211} + %stack.1215 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.841), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/stack" stack_frame_id=214} + %stack.1216 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.841), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/stack" stack_frame_id=214} + %stack.1217 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1215, %stack.1216), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/stack" stack_frame_id=214} + %reshape.565 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1217), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/reshape"} + %cos.153 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.565), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/cos" stack_frame_id=218} + %convert_element_type.1371 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.153), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/convert_element_type" stack_frame_id=222} + %mul.2506 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1371), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=242} + %mul.2507 = bf16[1,1,128]{2,1,0} reshape(%mul.2506), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=242} + %mul.2508 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2507), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=242} + %mul.2509 = bf16[1,1,32,128]{3,2,1,0} multiply(%dot_general.840, %mul.2508), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=242} + %split.307 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.840), slice={[0:1], [0:1], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/split" stack_frame_id=234} + %neg.315 = bf16[1,1,32,64]{3,2,1,0} negate(%split.307), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/neg" stack_frame_id=235} + %stack.1218 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%neg.315), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/stack" stack_frame_id=238} + %split.306 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.840), slice={[0:1], [0:1], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/split" stack_frame_id=234} + %stack.1219 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%split.306), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/stack" stack_frame_id=238} + %stack.1220 = bf16[1,1,32,2,64]{4,3,2,1,0} concatenate(%stack.1218, %stack.1219), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/stack" stack_frame_id=238} + %reshape.566 = bf16[1,1,32,128]{3,2,1,0} reshape(%stack.1220), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/reshape" stack_frame_id=241} + %sin.153 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.565), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/sin" stack_frame_id=226} + %convert_element_type.1372 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.153), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/convert_element_type" stack_frame_id=230} + %mul.2510 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1372), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=243} + %mul.2511 = bf16[1,1,128]{2,1,0} reshape(%mul.2510), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=243} + %mul.2512 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2511), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=243} + %mul.2513 = bf16[1,1,32,128]{3,2,1,0} multiply(%reshape.566, %mul.2512), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=243} + %add.873 = bf16[1,1,32,128]{3,2,1,0} add(%mul.2509, %mul.2513), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/add" stack_frame_id=244} + %reshape.571 = bf16[1,1,32,1,128]{4,3,2,1,0} reshape(%add.873), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/reshape" stack_frame_id=83} + %slice.136 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.139), slice={[0:1], [0:1], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/slice" stack_frame_id=245} + %squeeze.140 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.136), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/squeeze" stack_frame_id=245} + %dot_general.842 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1369, %arg_tuple.5#127), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=256} + %convert_element_type.1373 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/convert_element_type" stack_frame_id=287} + %add.874 = f32[1,1]{1,0} reshape(%convert_element_type.1373), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/add"} + %iota.437 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/iota" stack_frame_id=276} + %mul.2514 = f32[64]{0} multiply(%iota.437, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=276} + %div.776 = f32[64]{0} divide(%mul.2514, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/div" stack_frame_id=277} + %neg.316 = f32[64]{0} negate(%div.776), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/neg" stack_frame_id=278} + %pow.309 = f32[64]{0} power(%broadcast.39, %neg.316), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/pow" stack_frame_id=281} + %div.777 = f32[64]{0} divide(%pow.309, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/div" stack_frame_id=282} + %dot_general.844 = f32[1,1,64]{2,1,0} dot(%add.874, %div.777), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/bi,j->bij/dot_general" stack_frame_id=290} + %stack.1221 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.844), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/stack" stack_frame_id=293} + %stack.1222 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.844), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/stack" stack_frame_id=293} + %stack.1223 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1221, %stack.1222), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/stack" stack_frame_id=293} + %reshape.567 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1223), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/reshape"} + %cos.154 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.567), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/cos" stack_frame_id=297} + %convert_element_type.1374 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.154), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/convert_element_type" stack_frame_id=301} + %mul.2515 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1374), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=321} + %mul.2516 = bf16[1,1,128]{2,1,0} reshape(%mul.2515), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=321} + %mul.2517 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2516), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=321} + %mul.2518 = bf16[1,1,8,128]{3,2,1,0} multiply(%dot_general.842, %mul.2517), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=321} + %split.309 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.842), slice={[0:1], [0:1], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/split" stack_frame_id=313} + %neg.317 = bf16[1,1,8,64]{3,2,1,0} negate(%split.309), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/neg" stack_frame_id=314} + %stack.1224 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%neg.317), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/stack" stack_frame_id=317} + %split.308 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.842), slice={[0:1], [0:1], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/split" stack_frame_id=313} + %stack.1225 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%split.308), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/stack" stack_frame_id=317} + %stack.1226 = bf16[1,1,8,2,64]{4,3,2,1,0} concatenate(%stack.1224, %stack.1225), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/stack" stack_frame_id=317} + %reshape.568 = bf16[1,1,8,128]{3,2,1,0} reshape(%stack.1226), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/reshape" stack_frame_id=320} + %sin.154 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.567), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/sin" stack_frame_id=305} + %convert_element_type.1375 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.154), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/convert_element_type" stack_frame_id=309} + %mul.2519 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1375), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=322} + %mul.2520 = bf16[1,1,128]{2,1,0} reshape(%mul.2519), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=322} + %mul.2521 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2520), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=322} + %mul.2522 = bf16[1,1,8,128]{3,2,1,0} multiply(%reshape.568, %mul.2521), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=322} + %add.875 = bf16[1,1,8,128]{3,2,1,0} add(%mul.2518, %mul.2522), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/rotary_embedding_13/add" stack_frame_id=323} + %lt.508 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/lt" stack_frame_id=326} + %add.876 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/add" stack_frame_id=326} + %select_n.117 = s32[] select(%lt.508, %add.876, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/select_n" stack_frame_id=326} + %dynamic_update_slice.91 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.140, %add.875, %constant.127, %select_n.117, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/dynamic_update_slice" stack_frame_id=326} + %broadcast_in_dim.1225 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.91), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/broadcast_in_dim" stack_frame_id=335} + %reshape.569 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1225), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/reshape" stack_frame_id=335} + %dot_general.845 = f32[1,32,1,1,41]{4,3,2,1,0} dot(%reshape.571, %reshape.569), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=83} + %mul.2523 = f32[1,32,1,1,41]{4,3,2,1,0} multiply(%dot_general.845, %broadcast.36), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/vmap()/mul" stack_frame_id=83} + %vmap_jit__where__.76 = f32[1,1,32,1,41]{4,3,2,1,0} call(%and.238, %mul.2523, %constant.114), to_apply=%_where_3.137, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/vmap(jit(_where))" stack_frame_id=83} + %reduce_max.461 = f32[1,1,32,1]{3,2,1,0} reduce(%vmap_jit__where__.76, %constant.113), dimensions={4}, to_apply=%region_183.191, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/vmap()/reduce_max" stack_frame_id=83} + %max.85 = f32[1,1,32,1]{3,2,1,0} maximum(%reduce_max.461, %broadcast.35), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/vmap()/max" stack_frame_id=83} + %broadcast_in_dim.1229 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%max.85), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %sub.323 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1229), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/vmap()/sub" stack_frame_id=83} + %sub.324 = f32[1,1,32,1]{3,2,1,0} reshape(%sub.323), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/vmap()/sub" stack_frame_id=83} + %sub.325 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%sub.324), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/vmap()/sub" stack_frame_id=83} + %sub.326 = f32[1,1,32,1,41]{4,3,2,1,0} subtract(%vmap_jit__where__.76, %sub.325), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/vmap()/sub" stack_frame_id=83} + %exp.81 = f32[1,1,32,1,41]{4,3,2,1,0} exponential(%sub.326), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/vmap()/exp" stack_frame_id=83} + %reduce_sum.1391 = f32[1,1,32,1]{3,2,1,0} reduce(%exp.81, %constant.123), dimensions={4}, to_apply=%region_184.192, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/vmap()/reduce_sum" stack_frame_id=83} + %broadcast_in_dim.1230 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%reduce_sum.1391), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %div.778 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1230), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/vmap()/div" stack_frame_id=83} + %div.779 = f32[1,1,32,1]{3,2,1,0} reshape(%div.778), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/vmap()/div" stack_frame_id=83} + %div.780 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%div.779), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/vmap()/div" stack_frame_id=83} + %div.781 = f32[1,1,32,1,41]{4,3,2,1,0} divide(%exp.81, %div.780), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/vmap()/div" stack_frame_id=83} + %convert_element_type.1377 = bf16[1,1,32,1,41]{4,3,2,1,0} convert(%div.781), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/vmap()/convert_element_type" stack_frame_id=83} + %dot_general.846 = bf16[1,32,128,1,1]{4,3,2,1,0} dot(%reshape.570, %convert_element_type.1377), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=83} + %transpose.80 = bf16[1,1,32,128]{3,2,1,0} reshape(%dot_general.846), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/reshape;jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/vmap()/transpose;jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %dot_general.847 = bf16[1,1,4096]{2,1,0} dot(%transpose.80, %arg_tuple.5#129), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=354} + %add.878 = bf16[1,1,4096]{2,1,0} add(%dot_general.847, %add.867), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/add" stack_frame_id=355} + %convert_element_type.1378 = f32[1,1,4096]{2,1,0} convert(%add.878), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/feedforward_layernorm/convert_element_type" stack_frame_id=359} + %pow.310 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1378, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/feedforward_layernorm/pow" stack_frame_id=362} + %reduce_sum.1392 = f32[1,1]{1,0} reduce(%pow.310, %constant.123), dimensions={2}, to_apply=%region_185.193, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/feedforward_layernorm/reduce_sum" stack_frame_id=92} + %broadcast_in_dim.1231 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1392), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/feedforward_layernorm/broadcast_in_dim" stack_frame_id=92} + %div.782 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1231, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/feedforward_layernorm/div" stack_frame_id=92} + %add.879 = f32[1,1,1]{2,1,0} add(%div.782, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/feedforward_layernorm/add" stack_frame_id=363} + %rsqrt.155 = f32[1,1,1]{2,1,0} rsqrt(%add.879), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/feedforward_layernorm/rsqrt" stack_frame_id=366} + %mul.2524 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.155), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2525 = f32[1,1]{1,0} reshape(%mul.2524), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2526 = f32[1,1,4096]{2,1,0} broadcast(%mul.2525), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2527 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1378, %mul.2526), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/feedforward_layernorm/mul" stack_frame_id=367} + %convert_element_type.1379 = f32[4096]{0} convert(%arg_tuple.5#130), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/feedforward_layernorm/convert_element_type" stack_frame_id=368} + %broadcast_in_dim.1232 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1379), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/feedforward_layernorm/broadcast_in_dim" stack_frame_id=368} + %mul.2528 = f32[1,1,4096]{2,1,0} multiply(%mul.2527, %broadcast_in_dim.1232), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/feedforward_layernorm/mul" stack_frame_id=368} + %convert_element_type.1380 = bf16[1,1,4096]{2,1,0} convert(%mul.2528), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/feedforward_layernorm/convert_element_type" stack_frame_id=372} + %dot_general.849 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1380, %arg_tuple.5#132), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/feedforward_intermediate_dense/dot_general" stack_frame_id=401} + %dot_general.848 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1380, %arg_tuple.5#131), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/feedforward_gate_dense/dot_general" stack_frame_id=381} + %convert_element_type.1381 = f32[1,1,14336]{2,1,0} convert(%dot_general.848), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/convert_element_type" stack_frame_id=385} + %jit_silu_.76 = f32[1,1,14336]{2,1,0} call(%convert_element_type.1381), to_apply=%silu_4.141, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/jit(silu)" stack_frame_id=388} + %convert_element_type.1382 = bf16[1,1,14336]{2,1,0} convert(%jit_silu_.76), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/convert_element_type" stack_frame_id=392} + %mul.2529 = bf16[1,1,14336]{2,1,0} multiply(%dot_general.849, %convert_element_type.1382), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/mul" stack_frame_id=404} + %dot_general.850 = bf16[1,1,4096]{2,1,0} dot(%mul.2529, %arg_tuple.5#133), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/feedforward_output_dense/dot_general" stack_frame_id=413} + %add.880 = bf16[1,1,4096]{2,1,0} add(%dot_general.850, %add.878), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/add" stack_frame_id=414} + %convert_element_type.1386 = f32[1,1,4096]{2,1,0} convert(%add.880), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention_layernorm/convert_element_type" stack_frame_id=164} + %pow.311 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1386, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention_layernorm/pow" stack_frame_id=167} + %reduce_sum.1393 = f32[1,1]{1,0} reduce(%pow.311, %constant.123), dimensions={2}, to_apply=%region_186.194, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention_layernorm/reduce_sum" stack_frame_id=73} + %broadcast_in_dim.1237 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1393), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention_layernorm/broadcast_in_dim" stack_frame_id=73} + %div.783 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1237, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention_layernorm/div" stack_frame_id=73} + %add.884 = f32[1,1,1]{2,1,0} add(%div.783, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention_layernorm/add" stack_frame_id=168} + %rsqrt.156 = f32[1,1,1]{2,1,0} rsqrt(%add.884), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention_layernorm/rsqrt" stack_frame_id=171} + %mul.2530 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.156), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2531 = f32[1,1]{1,0} reshape(%mul.2530), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2532 = f32[1,1,4096]{2,1,0} broadcast(%mul.2531), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2533 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1386, %mul.2532), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention_layernorm/mul" stack_frame_id=172} + %convert_element_type.1387 = f32[4096]{0} convert(%arg_tuple.5#134), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention_layernorm/convert_element_type" stack_frame_id=173} + %broadcast_in_dim.1238 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1387), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention_layernorm/broadcast_in_dim" stack_frame_id=173} + %mul.2534 = f32[1,1,4096]{2,1,0} multiply(%mul.2533, %broadcast_in_dim.1238), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention_layernorm/mul" stack_frame_id=173} + %convert_element_type.1388 = bf16[1,1,4096]{2,1,0} convert(%mul.2534), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention_layernorm/convert_element_type" stack_frame_id=177} + %dot_general.854 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1388, %arg_tuple.5#137), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=265} + %lt.515 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/lt" stack_frame_id=329} + %add.890 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/add" stack_frame_id=329} + %select_n.120 = s32[] select(%lt.515, %add.890, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/select_n" stack_frame_id=329} + %dynamic_update_slice.94 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.144, %dot_general.854, %constant.127, %select_n.120, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/dynamic_update_slice" stack_frame_id=329} + %broadcast_in_dim.1240 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.94), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/broadcast_in_dim" stack_frame_id=338} + %reshape.577 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1240), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/reshape" stack_frame_id=338} + %convert_element_type.1383 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/convert_element_type" stack_frame_id=139} + %add.881 = f32[1,1]{1,0} reshape(%convert_element_type.1383), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_14/add"} + %ge.409 = f32[1,1]{1,0} broadcast(%add.881), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/ge" stack_frame_id=143} + %ge.410 = f32[1]{0} reshape(%ge.409), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/ge" stack_frame_id=143} + %ge.411 = f32[1,41]{1,0} broadcast(%ge.410), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/ge" stack_frame_id=143} + %iota.438 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/iota" stack_frame_id=142} + %broadcast_in_dim.1233 = f32[1,41]{1,0} reshape(%iota.438), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/broadcast_in_dim" stack_frame_id=143} + %ge.412 = pred[1,41]{1,0} compare(%ge.411, %broadcast_in_dim.1233), direction=GE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/ge" stack_frame_id=143} + %broadcast_in_dim.1234 = pred[1,1,41]{2,1,0} reshape(%ge.412), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/broadcast_in_dim" stack_frame_id=146} + %convert_element_type.1385 = s32[1,1,41]{2,1,0} convert(%broadcast_in_dim.1234), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/convert_element_type" stack_frame_id=160} + %add.882 = s32[1,1]{1,0} reshape(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/add" stack_frame_id=147} + %lt.510 = s32[1,1]{1,0} broadcast(%add.882), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/lt" stack_frame_id=152} + %lt.511 = s32[1]{0} reshape(%lt.510), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/lt" stack_frame_id=152} + %lt.512 = s32[1,41]{1,0} broadcast(%lt.511), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/lt" stack_frame_id=152} + %iota.439 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/iota" stack_frame_id=150} + %broadcast_in_dim.1235 = s32[1,41]{1,0} reshape(%iota.439), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/broadcast_in_dim" stack_frame_id=148} + %add.883 = s32[1,41]{1,0} add(%broadcast_in_dim.1235, %broadcast.43), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/add" stack_frame_id=151} + %lt.513 = pred[1,41]{1,0} compare(%lt.512, %add.883), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/lt" stack_frame_id=152} + %convert_element_type.1384 = s32[1,41]{1,0} convert(%lt.513), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/convert_element_type" stack_frame_id=156} + %broadcast_in_dim.1236 = s32[1,1,41]{2,1,0} reshape(%convert_element_type.1384), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/broadcast_in_dim" stack_frame_id=160} + %min.77 = s32[1,1,41]{2,1,0} minimum(%convert_element_type.1385, %broadcast_in_dim.1236), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/min" stack_frame_id=160} + %broadcast_in_dim.1241 = s32[1,1,1,41]{3,2,1,0} reshape(%min.77), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/broadcast_in_dim" stack_frame_id=341} + %convert_element_type.1395 = pred[1,1,1,41]{3,2,1,0} compare(%broadcast_in_dim.1241, %broadcast.37), direction=NE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/convert_element_type" stack_frame_id=345} + %broadcast_in_dim.1242 = pred[1,1,1,1,41]{4,3,2,1,0} reshape(%convert_element_type.1395), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/broadcast_in_dim" stack_frame_id=83} + %and.239 = pred[1,1,1,1,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1242), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/vmap()/and" stack_frame_id=83} + %and.240 = pred[1,1,1,41]{3,2,1,0} reshape(%and.239), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/vmap()/and" stack_frame_id=83} + %and.241 = pred[1,1,32,1,41]{4,3,2,1,0} broadcast(%and.240), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/vmap()/and" stack_frame_id=83} + %dot_general.851 = bf16[1,1,32,128]{3,2,1,0} dot(%convert_element_type.1388, %arg_tuple.5#135), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=186} + %convert_element_type.1389 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/convert_element_type" stack_frame_id=208} + %add.885 = f32[1,1]{1,0} reshape(%convert_element_type.1389), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/add"} + %iota.440 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/iota" stack_frame_id=197} + %mul.2535 = f32[64]{0} multiply(%iota.440, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=197} + %div.784 = f32[64]{0} divide(%mul.2535, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/div" stack_frame_id=198} + %neg.318 = f32[64]{0} negate(%div.784), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/neg" stack_frame_id=199} + %pow.312 = f32[64]{0} power(%broadcast.39, %neg.318), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/pow" stack_frame_id=202} + %div.785 = f32[64]{0} divide(%pow.312, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/div" stack_frame_id=203} + %dot_general.852 = f32[1,1,64]{2,1,0} dot(%add.885, %div.785), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/bi,j->bij/dot_general" stack_frame_id=211} + %stack.1230 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.852), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/stack" stack_frame_id=214} + %stack.1231 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.852), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/stack" stack_frame_id=214} + %stack.1232 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1230, %stack.1231), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/stack" stack_frame_id=214} + %reshape.572 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1232), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/reshape"} + %cos.155 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.572), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/cos" stack_frame_id=218} + %convert_element_type.1390 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.155), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/convert_element_type" stack_frame_id=222} + %mul.2536 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1390), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=242} + %mul.2537 = bf16[1,1,128]{2,1,0} reshape(%mul.2536), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=242} + %mul.2538 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2537), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=242} + %mul.2539 = bf16[1,1,32,128]{3,2,1,0} multiply(%dot_general.851, %mul.2538), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=242} + %split.311 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.851), slice={[0:1], [0:1], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/split" stack_frame_id=234} + %neg.319 = bf16[1,1,32,64]{3,2,1,0} negate(%split.311), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/neg" stack_frame_id=235} + %stack.1233 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%neg.319), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/stack" stack_frame_id=238} + %split.310 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.851), slice={[0:1], [0:1], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/split" stack_frame_id=234} + %stack.1234 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%split.310), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/stack" stack_frame_id=238} + %stack.1235 = bf16[1,1,32,2,64]{4,3,2,1,0} concatenate(%stack.1233, %stack.1234), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/stack" stack_frame_id=238} + %reshape.573 = bf16[1,1,32,128]{3,2,1,0} reshape(%stack.1235), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/reshape" stack_frame_id=241} + %sin.155 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.572), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/sin" stack_frame_id=226} + %convert_element_type.1391 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.155), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/convert_element_type" stack_frame_id=230} + %mul.2540 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1391), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=243} + %mul.2541 = bf16[1,1,128]{2,1,0} reshape(%mul.2540), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=243} + %mul.2542 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2541), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=243} + %mul.2543 = bf16[1,1,32,128]{3,2,1,0} multiply(%reshape.573, %mul.2542), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=243} + %add.886 = bf16[1,1,32,128]{3,2,1,0} add(%mul.2539, %mul.2543), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/add" stack_frame_id=244} + %reshape.578 = bf16[1,1,32,1,128]{4,3,2,1,0} reshape(%add.886), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/reshape" stack_frame_id=83} + %slice.139 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.142), slice={[0:1], [0:1], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/slice" stack_frame_id=245} + %squeeze.143 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.139), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/squeeze" stack_frame_id=245} + %dot_general.853 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1388, %arg_tuple.5#136), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=256} + %convert_element_type.1392 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/convert_element_type" stack_frame_id=287} + %add.887 = f32[1,1]{1,0} reshape(%convert_element_type.1392), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/add"} + %iota.441 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/iota" stack_frame_id=276} + %mul.2544 = f32[64]{0} multiply(%iota.441, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=276} + %div.786 = f32[64]{0} divide(%mul.2544, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/div" stack_frame_id=277} + %neg.320 = f32[64]{0} negate(%div.786), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/neg" stack_frame_id=278} + %pow.313 = f32[64]{0} power(%broadcast.39, %neg.320), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/pow" stack_frame_id=281} + %div.787 = f32[64]{0} divide(%pow.313, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/div" stack_frame_id=282} + %dot_general.855 = f32[1,1,64]{2,1,0} dot(%add.887, %div.787), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/bi,j->bij/dot_general" stack_frame_id=290} + %stack.1236 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.855), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/stack" stack_frame_id=293} + %stack.1237 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.855), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/stack" stack_frame_id=293} + %stack.1238 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1236, %stack.1237), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/stack" stack_frame_id=293} + %reshape.574 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1238), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/reshape"} + %cos.156 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.574), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/cos" stack_frame_id=297} + %convert_element_type.1393 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.156), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/convert_element_type" stack_frame_id=301} + %mul.2545 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1393), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=321} + %mul.2546 = bf16[1,1,128]{2,1,0} reshape(%mul.2545), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=321} + %mul.2547 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2546), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=321} + %mul.2548 = bf16[1,1,8,128]{3,2,1,0} multiply(%dot_general.853, %mul.2547), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=321} + %split.313 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.853), slice={[0:1], [0:1], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/split" stack_frame_id=313} + %neg.321 = bf16[1,1,8,64]{3,2,1,0} negate(%split.313), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/neg" stack_frame_id=314} + %stack.1239 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%neg.321), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/stack" stack_frame_id=317} + %split.312 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.853), slice={[0:1], [0:1], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/split" stack_frame_id=313} + %stack.1240 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%split.312), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/stack" stack_frame_id=317} + %stack.1241 = bf16[1,1,8,2,64]{4,3,2,1,0} concatenate(%stack.1239, %stack.1240), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/stack" stack_frame_id=317} + %reshape.575 = bf16[1,1,8,128]{3,2,1,0} reshape(%stack.1241), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/reshape" stack_frame_id=320} + %sin.156 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.574), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/sin" stack_frame_id=305} + %convert_element_type.1394 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.156), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/convert_element_type" stack_frame_id=309} + %mul.2549 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1394), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=322} + %mul.2550 = bf16[1,1,128]{2,1,0} reshape(%mul.2549), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=322} + %mul.2551 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2550), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=322} + %mul.2552 = bf16[1,1,8,128]{3,2,1,0} multiply(%reshape.575, %mul.2551), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=322} + %add.888 = bf16[1,1,8,128]{3,2,1,0} add(%mul.2548, %mul.2552), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/rotary_embedding_14/add" stack_frame_id=323} + %lt.514 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/lt" stack_frame_id=326} + %add.889 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/add" stack_frame_id=326} + %select_n.119 = s32[] select(%lt.514, %add.889, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/select_n" stack_frame_id=326} + %dynamic_update_slice.93 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.143, %add.888, %constant.127, %select_n.119, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/dynamic_update_slice" stack_frame_id=326} + %broadcast_in_dim.1239 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.93), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/broadcast_in_dim" stack_frame_id=335} + %reshape.576 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1239), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/reshape" stack_frame_id=335} + %dot_general.856 = f32[1,32,1,1,41]{4,3,2,1,0} dot(%reshape.578, %reshape.576), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=83} + %mul.2553 = f32[1,32,1,1,41]{4,3,2,1,0} multiply(%dot_general.856, %broadcast.36), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/vmap()/mul" stack_frame_id=83} + %vmap_jit__where__.77 = f32[1,1,32,1,41]{4,3,2,1,0} call(%and.241, %mul.2553, %constant.114), to_apply=%_where_3.137, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/vmap(jit(_where))" stack_frame_id=83} + %reduce_max.462 = f32[1,1,32,1]{3,2,1,0} reduce(%vmap_jit__where__.77, %constant.113), dimensions={4}, to_apply=%region_187.195, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/vmap()/reduce_max" stack_frame_id=83} + %max.86 = f32[1,1,32,1]{3,2,1,0} maximum(%reduce_max.462, %broadcast.35), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/vmap()/max" stack_frame_id=83} + %broadcast_in_dim.1243 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%max.86), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %sub.327 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1243), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/vmap()/sub" stack_frame_id=83} + %sub.328 = f32[1,1,32,1]{3,2,1,0} reshape(%sub.327), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/vmap()/sub" stack_frame_id=83} + %sub.329 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%sub.328), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/vmap()/sub" stack_frame_id=83} + %sub.330 = f32[1,1,32,1,41]{4,3,2,1,0} subtract(%vmap_jit__where__.77, %sub.329), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/vmap()/sub" stack_frame_id=83} + %exp.82 = f32[1,1,32,1,41]{4,3,2,1,0} exponential(%sub.330), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/vmap()/exp" stack_frame_id=83} + %reduce_sum.1394 = f32[1,1,32,1]{3,2,1,0} reduce(%exp.82, %constant.123), dimensions={4}, to_apply=%region_188.196, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/vmap()/reduce_sum" stack_frame_id=83} + %broadcast_in_dim.1244 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%reduce_sum.1394), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %div.788 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1244), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/vmap()/div" stack_frame_id=83} + %div.789 = f32[1,1,32,1]{3,2,1,0} reshape(%div.788), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/vmap()/div" stack_frame_id=83} + %div.790 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%div.789), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/vmap()/div" stack_frame_id=83} + %div.791 = f32[1,1,32,1,41]{4,3,2,1,0} divide(%exp.82, %div.790), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/vmap()/div" stack_frame_id=83} + %convert_element_type.1396 = bf16[1,1,32,1,41]{4,3,2,1,0} convert(%div.791), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/vmap()/convert_element_type" stack_frame_id=83} + %dot_general.857 = bf16[1,32,128,1,1]{4,3,2,1,0} dot(%reshape.577, %convert_element_type.1396), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=83} + %transpose.81 = bf16[1,1,32,128]{3,2,1,0} reshape(%dot_general.857), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/reshape;jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/vmap()/transpose;jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %dot_general.858 = bf16[1,1,4096]{2,1,0} dot(%transpose.81, %arg_tuple.5#138), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=354} + %add.891 = bf16[1,1,4096]{2,1,0} add(%dot_general.858, %add.880), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/add" stack_frame_id=355} + %convert_element_type.1397 = f32[1,1,4096]{2,1,0} convert(%add.891), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/feedforward_layernorm/convert_element_type" stack_frame_id=359} + %pow.314 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1397, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/feedforward_layernorm/pow" stack_frame_id=362} + %reduce_sum.1395 = f32[1,1]{1,0} reduce(%pow.314, %constant.123), dimensions={2}, to_apply=%region_189.197, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/feedforward_layernorm/reduce_sum" stack_frame_id=92} + %broadcast_in_dim.1245 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1395), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/feedforward_layernorm/broadcast_in_dim" stack_frame_id=92} + %div.792 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1245, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/feedforward_layernorm/div" stack_frame_id=92} + %add.892 = f32[1,1,1]{2,1,0} add(%div.792, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/feedforward_layernorm/add" stack_frame_id=363} + %rsqrt.157 = f32[1,1,1]{2,1,0} rsqrt(%add.892), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/feedforward_layernorm/rsqrt" stack_frame_id=366} + %mul.2554 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.157), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2555 = f32[1,1]{1,0} reshape(%mul.2554), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2556 = f32[1,1,4096]{2,1,0} broadcast(%mul.2555), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2557 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1397, %mul.2556), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/feedforward_layernorm/mul" stack_frame_id=367} + %convert_element_type.1398 = f32[4096]{0} convert(%arg_tuple.5#139), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/feedforward_layernorm/convert_element_type" stack_frame_id=368} + %broadcast_in_dim.1246 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1398), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/feedforward_layernorm/broadcast_in_dim" stack_frame_id=368} + %mul.2558 = f32[1,1,4096]{2,1,0} multiply(%mul.2557, %broadcast_in_dim.1246), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/feedforward_layernorm/mul" stack_frame_id=368} + %convert_element_type.1399 = bf16[1,1,4096]{2,1,0} convert(%mul.2558), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/feedforward_layernorm/convert_element_type" stack_frame_id=372} + %dot_general.860 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1399, %arg_tuple.5#141), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/feedforward_intermediate_dense/dot_general" stack_frame_id=401} + %dot_general.859 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1399, %arg_tuple.5#140), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/feedforward_gate_dense/dot_general" stack_frame_id=381} + %convert_element_type.1400 = f32[1,1,14336]{2,1,0} convert(%dot_general.859), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/convert_element_type" stack_frame_id=385} + %jit_silu_.77 = f32[1,1,14336]{2,1,0} call(%convert_element_type.1400), to_apply=%silu_4.141, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/jit(silu)" stack_frame_id=388} + %convert_element_type.1401 = bf16[1,1,14336]{2,1,0} convert(%jit_silu_.77), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/convert_element_type" stack_frame_id=392} + %mul.2559 = bf16[1,1,14336]{2,1,0} multiply(%dot_general.860, %convert_element_type.1401), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/mul" stack_frame_id=404} + %dot_general.861 = bf16[1,1,4096]{2,1,0} dot(%mul.2559, %arg_tuple.5#142), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/feedforward_output_dense/dot_general" stack_frame_id=413} + %add.893 = bf16[1,1,4096]{2,1,0} add(%dot_general.861, %add.891), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/add" stack_frame_id=414} + %convert_element_type.1405 = f32[1,1,4096]{2,1,0} convert(%add.893), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention_layernorm/convert_element_type" stack_frame_id=164} + %pow.315 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1405, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention_layernorm/pow" stack_frame_id=167} + %reduce_sum.1396 = f32[1,1]{1,0} reduce(%pow.315, %constant.123), dimensions={2}, to_apply=%region_190.198, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention_layernorm/reduce_sum" stack_frame_id=73} + %broadcast_in_dim.1251 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1396), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention_layernorm/broadcast_in_dim" stack_frame_id=73} + %div.793 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1251, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention_layernorm/div" stack_frame_id=73} + %add.897 = f32[1,1,1]{2,1,0} add(%div.793, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention_layernorm/add" stack_frame_id=168} + %rsqrt.158 = f32[1,1,1]{2,1,0} rsqrt(%add.897), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention_layernorm/rsqrt" stack_frame_id=171} + %mul.2560 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.158), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2561 = f32[1,1]{1,0} reshape(%mul.2560), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2562 = f32[1,1,4096]{2,1,0} broadcast(%mul.2561), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2563 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1405, %mul.2562), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention_layernorm/mul" stack_frame_id=172} + %convert_element_type.1406 = f32[4096]{0} convert(%arg_tuple.5#143), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention_layernorm/convert_element_type" stack_frame_id=173} + %broadcast_in_dim.1252 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1406), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention_layernorm/broadcast_in_dim" stack_frame_id=173} + %mul.2564 = f32[1,1,4096]{2,1,0} multiply(%mul.2563, %broadcast_in_dim.1252), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention_layernorm/mul" stack_frame_id=173} + %convert_element_type.1407 = bf16[1,1,4096]{2,1,0} convert(%mul.2564), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention_layernorm/convert_element_type" stack_frame_id=177} + %dot_general.865 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1407, %arg_tuple.5#146), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=265} + %lt.521 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/lt" stack_frame_id=329} + %add.903 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/add" stack_frame_id=329} + %select_n.122 = s32[] select(%lt.521, %add.903, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/select_n" stack_frame_id=329} + %dynamic_update_slice.96 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.147, %dot_general.865, %constant.127, %select_n.122, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/dynamic_update_slice" stack_frame_id=329} + %broadcast_in_dim.1254 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.96), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/broadcast_in_dim" stack_frame_id=338} + %reshape.584 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1254), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/reshape" stack_frame_id=338} + %convert_element_type.1402 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/convert_element_type" stack_frame_id=139} + %add.894 = f32[1,1]{1,0} reshape(%convert_element_type.1402), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_15/add"} + %ge.413 = f32[1,1]{1,0} broadcast(%add.894), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/ge" stack_frame_id=143} + %ge.414 = f32[1]{0} reshape(%ge.413), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/ge" stack_frame_id=143} + %ge.415 = f32[1,41]{1,0} broadcast(%ge.414), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/ge" stack_frame_id=143} + %iota.442 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/iota" stack_frame_id=142} + %broadcast_in_dim.1247 = f32[1,41]{1,0} reshape(%iota.442), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/broadcast_in_dim" stack_frame_id=143} + %ge.416 = pred[1,41]{1,0} compare(%ge.415, %broadcast_in_dim.1247), direction=GE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/ge" stack_frame_id=143} + %broadcast_in_dim.1248 = pred[1,1,41]{2,1,0} reshape(%ge.416), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/broadcast_in_dim" stack_frame_id=146} + %convert_element_type.1404 = s32[1,1,41]{2,1,0} convert(%broadcast_in_dim.1248), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/convert_element_type" stack_frame_id=160} + %add.895 = s32[1,1]{1,0} reshape(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/add" stack_frame_id=147} + %lt.516 = s32[1,1]{1,0} broadcast(%add.895), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/lt" stack_frame_id=152} + %lt.517 = s32[1]{0} reshape(%lt.516), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/lt" stack_frame_id=152} + %lt.518 = s32[1,41]{1,0} broadcast(%lt.517), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/lt" stack_frame_id=152} + %iota.443 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/iota" stack_frame_id=150} + %broadcast_in_dim.1249 = s32[1,41]{1,0} reshape(%iota.443), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/broadcast_in_dim" stack_frame_id=148} + %add.896 = s32[1,41]{1,0} add(%broadcast_in_dim.1249, %broadcast.43), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/add" stack_frame_id=151} + %lt.519 = pred[1,41]{1,0} compare(%lt.518, %add.896), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/lt" stack_frame_id=152} + %convert_element_type.1403 = s32[1,41]{1,0} convert(%lt.519), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/convert_element_type" stack_frame_id=156} + %broadcast_in_dim.1250 = s32[1,1,41]{2,1,0} reshape(%convert_element_type.1403), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/broadcast_in_dim" stack_frame_id=160} + %min.78 = s32[1,1,41]{2,1,0} minimum(%convert_element_type.1404, %broadcast_in_dim.1250), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/min" stack_frame_id=160} + %broadcast_in_dim.1255 = s32[1,1,1,41]{3,2,1,0} reshape(%min.78), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/broadcast_in_dim" stack_frame_id=341} + %convert_element_type.1414 = pred[1,1,1,41]{3,2,1,0} compare(%broadcast_in_dim.1255, %broadcast.37), direction=NE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/convert_element_type" stack_frame_id=345} + %broadcast_in_dim.1256 = pred[1,1,1,1,41]{4,3,2,1,0} reshape(%convert_element_type.1414), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/broadcast_in_dim" stack_frame_id=83} + %and.242 = pred[1,1,1,1,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1256), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/vmap()/and" stack_frame_id=83} + %and.243 = pred[1,1,1,41]{3,2,1,0} reshape(%and.242), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/vmap()/and" stack_frame_id=83} + %and.244 = pred[1,1,32,1,41]{4,3,2,1,0} broadcast(%and.243), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/vmap()/and" stack_frame_id=83} + %dot_general.862 = bf16[1,1,32,128]{3,2,1,0} dot(%convert_element_type.1407, %arg_tuple.5#144), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=186} + %convert_element_type.1408 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/convert_element_type" stack_frame_id=208} + %add.898 = f32[1,1]{1,0} reshape(%convert_element_type.1408), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/add"} + %iota.444 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/iota" stack_frame_id=197} + %mul.2565 = f32[64]{0} multiply(%iota.444, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=197} + %div.794 = f32[64]{0} divide(%mul.2565, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/div" stack_frame_id=198} + %neg.322 = f32[64]{0} negate(%div.794), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/neg" stack_frame_id=199} + %pow.316 = f32[64]{0} power(%broadcast.39, %neg.322), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/pow" stack_frame_id=202} + %div.795 = f32[64]{0} divide(%pow.316, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/div" stack_frame_id=203} + %dot_general.863 = f32[1,1,64]{2,1,0} dot(%add.898, %div.795), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/bi,j->bij/dot_general" stack_frame_id=211} + %stack.1245 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.863), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/stack" stack_frame_id=214} + %stack.1246 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.863), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/stack" stack_frame_id=214} + %stack.1247 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1245, %stack.1246), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/stack" stack_frame_id=214} + %reshape.579 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1247), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/reshape"} + %cos.157 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.579), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/cos" stack_frame_id=218} + %convert_element_type.1409 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.157), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/convert_element_type" stack_frame_id=222} + %mul.2566 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1409), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=242} + %mul.2567 = bf16[1,1,128]{2,1,0} reshape(%mul.2566), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=242} + %mul.2568 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2567), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=242} + %mul.2569 = bf16[1,1,32,128]{3,2,1,0} multiply(%dot_general.862, %mul.2568), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=242} + %split.315 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.862), slice={[0:1], [0:1], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/split" stack_frame_id=234} + %neg.323 = bf16[1,1,32,64]{3,2,1,0} negate(%split.315), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/neg" stack_frame_id=235} + %stack.1248 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%neg.323), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/stack" stack_frame_id=238} + %split.314 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.862), slice={[0:1], [0:1], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/split" stack_frame_id=234} + %stack.1249 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%split.314), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/stack" stack_frame_id=238} + %stack.1250 = bf16[1,1,32,2,64]{4,3,2,1,0} concatenate(%stack.1248, %stack.1249), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/stack" stack_frame_id=238} + %reshape.580 = bf16[1,1,32,128]{3,2,1,0} reshape(%stack.1250), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/reshape" stack_frame_id=241} + %sin.157 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.579), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/sin" stack_frame_id=226} + %convert_element_type.1410 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.157), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/convert_element_type" stack_frame_id=230} + %mul.2570 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1410), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=243} + %mul.2571 = bf16[1,1,128]{2,1,0} reshape(%mul.2570), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=243} + %mul.2572 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2571), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=243} + %mul.2573 = bf16[1,1,32,128]{3,2,1,0} multiply(%reshape.580, %mul.2572), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=243} + %add.899 = bf16[1,1,32,128]{3,2,1,0} add(%mul.2569, %mul.2573), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/add" stack_frame_id=244} + %reshape.585 = bf16[1,1,32,1,128]{4,3,2,1,0} reshape(%add.899), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/reshape" stack_frame_id=83} + %slice.142 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.145), slice={[0:1], [0:1], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/slice" stack_frame_id=245} + %squeeze.146 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.142), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/squeeze" stack_frame_id=245} + %dot_general.864 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1407, %arg_tuple.5#145), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=256} + %convert_element_type.1411 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/convert_element_type" stack_frame_id=287} + %add.900 = f32[1,1]{1,0} reshape(%convert_element_type.1411), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/add"} + %iota.445 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/iota" stack_frame_id=276} + %mul.2574 = f32[64]{0} multiply(%iota.445, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=276} + %div.796 = f32[64]{0} divide(%mul.2574, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/div" stack_frame_id=277} + %neg.324 = f32[64]{0} negate(%div.796), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/neg" stack_frame_id=278} + %pow.317 = f32[64]{0} power(%broadcast.39, %neg.324), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/pow" stack_frame_id=281} + %div.797 = f32[64]{0} divide(%pow.317, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/div" stack_frame_id=282} + %dot_general.866 = f32[1,1,64]{2,1,0} dot(%add.900, %div.797), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/bi,j->bij/dot_general" stack_frame_id=290} + %stack.1251 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.866), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/stack" stack_frame_id=293} + %stack.1252 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.866), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/stack" stack_frame_id=293} + %stack.1253 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1251, %stack.1252), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/stack" stack_frame_id=293} + %reshape.581 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1253), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/reshape"} + %cos.158 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.581), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/cos" stack_frame_id=297} + %convert_element_type.1412 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.158), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/convert_element_type" stack_frame_id=301} + %mul.2575 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1412), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=321} + %mul.2576 = bf16[1,1,128]{2,1,0} reshape(%mul.2575), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=321} + %mul.2577 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2576), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=321} + %mul.2578 = bf16[1,1,8,128]{3,2,1,0} multiply(%dot_general.864, %mul.2577), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=321} + %split.317 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.864), slice={[0:1], [0:1], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/split" stack_frame_id=313} + %neg.325 = bf16[1,1,8,64]{3,2,1,0} negate(%split.317), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/neg" stack_frame_id=314} + %stack.1254 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%neg.325), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/stack" stack_frame_id=317} + %split.316 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.864), slice={[0:1], [0:1], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/split" stack_frame_id=313} + %stack.1255 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%split.316), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/stack" stack_frame_id=317} + %stack.1256 = bf16[1,1,8,2,64]{4,3,2,1,0} concatenate(%stack.1254, %stack.1255), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/stack" stack_frame_id=317} + %reshape.582 = bf16[1,1,8,128]{3,2,1,0} reshape(%stack.1256), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/reshape" stack_frame_id=320} + %sin.158 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.581), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/sin" stack_frame_id=305} + %convert_element_type.1413 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.158), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/convert_element_type" stack_frame_id=309} + %mul.2579 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1413), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=322} + %mul.2580 = bf16[1,1,128]{2,1,0} reshape(%mul.2579), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=322} + %mul.2581 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2580), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=322} + %mul.2582 = bf16[1,1,8,128]{3,2,1,0} multiply(%reshape.582, %mul.2581), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=322} + %add.901 = bf16[1,1,8,128]{3,2,1,0} add(%mul.2578, %mul.2582), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/rotary_embedding_15/add" stack_frame_id=323} + %lt.520 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/lt" stack_frame_id=326} + %add.902 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/add" stack_frame_id=326} + %select_n.121 = s32[] select(%lt.520, %add.902, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/select_n" stack_frame_id=326} + %dynamic_update_slice.95 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.146, %add.901, %constant.127, %select_n.121, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/dynamic_update_slice" stack_frame_id=326} + %broadcast_in_dim.1253 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.95), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/broadcast_in_dim" stack_frame_id=335} + %reshape.583 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1253), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/reshape" stack_frame_id=335} + %dot_general.867 = f32[1,32,1,1,41]{4,3,2,1,0} dot(%reshape.585, %reshape.583), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=83} + %mul.2583 = f32[1,32,1,1,41]{4,3,2,1,0} multiply(%dot_general.867, %broadcast.36), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/vmap()/mul" stack_frame_id=83} + %vmap_jit__where__.78 = f32[1,1,32,1,41]{4,3,2,1,0} call(%and.244, %mul.2583, %constant.114), to_apply=%_where_3.137, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/vmap(jit(_where))" stack_frame_id=83} + %reduce_max.463 = f32[1,1,32,1]{3,2,1,0} reduce(%vmap_jit__where__.78, %constant.113), dimensions={4}, to_apply=%region_191.199, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/vmap()/reduce_max" stack_frame_id=83} + %max.87 = f32[1,1,32,1]{3,2,1,0} maximum(%reduce_max.463, %broadcast.35), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/vmap()/max" stack_frame_id=83} + %broadcast_in_dim.1257 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%max.87), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %sub.331 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1257), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/vmap()/sub" stack_frame_id=83} + %sub.332 = f32[1,1,32,1]{3,2,1,0} reshape(%sub.331), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/vmap()/sub" stack_frame_id=83} + %sub.333 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%sub.332), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/vmap()/sub" stack_frame_id=83} + %sub.334 = f32[1,1,32,1,41]{4,3,2,1,0} subtract(%vmap_jit__where__.78, %sub.333), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/vmap()/sub" stack_frame_id=83} + %exp.83 = f32[1,1,32,1,41]{4,3,2,1,0} exponential(%sub.334), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/vmap()/exp" stack_frame_id=83} + %reduce_sum.1397 = f32[1,1,32,1]{3,2,1,0} reduce(%exp.83, %constant.123), dimensions={4}, to_apply=%region_192.200, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/vmap()/reduce_sum" stack_frame_id=83} + %broadcast_in_dim.1258 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%reduce_sum.1397), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %div.798 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1258), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/vmap()/div" stack_frame_id=83} + %div.799 = f32[1,1,32,1]{3,2,1,0} reshape(%div.798), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/vmap()/div" stack_frame_id=83} + %div.800 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%div.799), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/vmap()/div" stack_frame_id=83} + %div.801 = f32[1,1,32,1,41]{4,3,2,1,0} divide(%exp.83, %div.800), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/vmap()/div" stack_frame_id=83} + %convert_element_type.1415 = bf16[1,1,32,1,41]{4,3,2,1,0} convert(%div.801), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/vmap()/convert_element_type" stack_frame_id=83} + %dot_general.868 = bf16[1,32,128,1,1]{4,3,2,1,0} dot(%reshape.584, %convert_element_type.1415), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=83} + %transpose.82 = bf16[1,1,32,128]{3,2,1,0} reshape(%dot_general.868), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/reshape;jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/vmap()/transpose;jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %dot_general.869 = bf16[1,1,4096]{2,1,0} dot(%transpose.82, %arg_tuple.5#147), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=354} + %add.904 = bf16[1,1,4096]{2,1,0} add(%dot_general.869, %add.893), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/add" stack_frame_id=355} + %convert_element_type.1416 = f32[1,1,4096]{2,1,0} convert(%add.904), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/feedforward_layernorm/convert_element_type" stack_frame_id=359} + %pow.318 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1416, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/feedforward_layernorm/pow" stack_frame_id=362} + %reduce_sum.1398 = f32[1,1]{1,0} reduce(%pow.318, %constant.123), dimensions={2}, to_apply=%region_193.201, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/feedforward_layernorm/reduce_sum" stack_frame_id=92} + %broadcast_in_dim.1259 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1398), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/feedforward_layernorm/broadcast_in_dim" stack_frame_id=92} + %div.802 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1259, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/feedforward_layernorm/div" stack_frame_id=92} + %add.905 = f32[1,1,1]{2,1,0} add(%div.802, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/feedforward_layernorm/add" stack_frame_id=363} + %rsqrt.159 = f32[1,1,1]{2,1,0} rsqrt(%add.905), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/feedforward_layernorm/rsqrt" stack_frame_id=366} + %mul.2584 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.159), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2585 = f32[1,1]{1,0} reshape(%mul.2584), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2586 = f32[1,1,4096]{2,1,0} broadcast(%mul.2585), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2587 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1416, %mul.2586), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/feedforward_layernorm/mul" stack_frame_id=367} + %convert_element_type.1417 = f32[4096]{0} convert(%arg_tuple.5#148), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/feedforward_layernorm/convert_element_type" stack_frame_id=368} + %broadcast_in_dim.1260 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1417), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/feedforward_layernorm/broadcast_in_dim" stack_frame_id=368} + %mul.2588 = f32[1,1,4096]{2,1,0} multiply(%mul.2587, %broadcast_in_dim.1260), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/feedforward_layernorm/mul" stack_frame_id=368} + %convert_element_type.1418 = bf16[1,1,4096]{2,1,0} convert(%mul.2588), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/feedforward_layernorm/convert_element_type" stack_frame_id=372} + %dot_general.871 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1418, %arg_tuple.5#150), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/feedforward_intermediate_dense/dot_general" stack_frame_id=401} + %dot_general.870 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1418, %arg_tuple.5#149), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/feedforward_gate_dense/dot_general" stack_frame_id=381} + %convert_element_type.1419 = f32[1,1,14336]{2,1,0} convert(%dot_general.870), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/convert_element_type" stack_frame_id=385} + %jit_silu_.78 = f32[1,1,14336]{2,1,0} call(%convert_element_type.1419), to_apply=%silu_4.141, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/jit(silu)" stack_frame_id=388} + %convert_element_type.1420 = bf16[1,1,14336]{2,1,0} convert(%jit_silu_.78), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/convert_element_type" stack_frame_id=392} + %mul.2589 = bf16[1,1,14336]{2,1,0} multiply(%dot_general.871, %convert_element_type.1420), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/mul" stack_frame_id=404} + %dot_general.872 = bf16[1,1,4096]{2,1,0} dot(%mul.2589, %arg_tuple.5#151), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/feedforward_output_dense/dot_general" stack_frame_id=413} + %add.906 = bf16[1,1,4096]{2,1,0} add(%dot_general.872, %add.904), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/add" stack_frame_id=414} + %convert_element_type.1424 = f32[1,1,4096]{2,1,0} convert(%add.906), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention_layernorm/convert_element_type" stack_frame_id=164} + %pow.319 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1424, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention_layernorm/pow" stack_frame_id=167} + %reduce_sum.1399 = f32[1,1]{1,0} reduce(%pow.319, %constant.123), dimensions={2}, to_apply=%region_194.202, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention_layernorm/reduce_sum" stack_frame_id=73} + %broadcast_in_dim.1265 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1399), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention_layernorm/broadcast_in_dim" stack_frame_id=73} + %div.803 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1265, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention_layernorm/div" stack_frame_id=73} + %add.910 = f32[1,1,1]{2,1,0} add(%div.803, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention_layernorm/add" stack_frame_id=168} + %rsqrt.160 = f32[1,1,1]{2,1,0} rsqrt(%add.910), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention_layernorm/rsqrt" stack_frame_id=171} + %mul.2590 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.160), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2591 = f32[1,1]{1,0} reshape(%mul.2590), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2592 = f32[1,1,4096]{2,1,0} broadcast(%mul.2591), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2593 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1424, %mul.2592), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention_layernorm/mul" stack_frame_id=172} + %convert_element_type.1425 = f32[4096]{0} convert(%arg_tuple.5#152), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention_layernorm/convert_element_type" stack_frame_id=173} + %broadcast_in_dim.1266 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1425), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention_layernorm/broadcast_in_dim" stack_frame_id=173} + %mul.2594 = f32[1,1,4096]{2,1,0} multiply(%mul.2593, %broadcast_in_dim.1266), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention_layernorm/mul" stack_frame_id=173} + %convert_element_type.1426 = bf16[1,1,4096]{2,1,0} convert(%mul.2594), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention_layernorm/convert_element_type" stack_frame_id=177} + %dot_general.876 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1426, %arg_tuple.5#155), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=265} + %lt.527 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/lt" stack_frame_id=329} + %add.916 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/add" stack_frame_id=329} + %select_n.124 = s32[] select(%lt.527, %add.916, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/select_n" stack_frame_id=329} + %dynamic_update_slice.98 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.150, %dot_general.876, %constant.127, %select_n.124, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/dynamic_update_slice" stack_frame_id=329} + %broadcast_in_dim.1268 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.98), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/broadcast_in_dim" stack_frame_id=338} + %reshape.591 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1268), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/reshape" stack_frame_id=338} + %convert_element_type.1421 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/convert_element_type" stack_frame_id=139} + %add.907 = f32[1,1]{1,0} reshape(%convert_element_type.1421), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_16/add"} + %ge.417 = f32[1,1]{1,0} broadcast(%add.907), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/ge" stack_frame_id=143} + %ge.418 = f32[1]{0} reshape(%ge.417), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/ge" stack_frame_id=143} + %ge.419 = f32[1,41]{1,0} broadcast(%ge.418), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/ge" stack_frame_id=143} + %iota.446 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/iota" stack_frame_id=142} + %broadcast_in_dim.1261 = f32[1,41]{1,0} reshape(%iota.446), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/broadcast_in_dim" stack_frame_id=143} + %ge.420 = pred[1,41]{1,0} compare(%ge.419, %broadcast_in_dim.1261), direction=GE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/ge" stack_frame_id=143} + %broadcast_in_dim.1262 = pred[1,1,41]{2,1,0} reshape(%ge.420), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/broadcast_in_dim" stack_frame_id=146} + %convert_element_type.1423 = s32[1,1,41]{2,1,0} convert(%broadcast_in_dim.1262), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/convert_element_type" stack_frame_id=160} + %add.908 = s32[1,1]{1,0} reshape(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/add" stack_frame_id=147} + %lt.522 = s32[1,1]{1,0} broadcast(%add.908), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/lt" stack_frame_id=152} + %lt.523 = s32[1]{0} reshape(%lt.522), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/lt" stack_frame_id=152} + %lt.524 = s32[1,41]{1,0} broadcast(%lt.523), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/lt" stack_frame_id=152} + %iota.447 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/iota" stack_frame_id=150} + %broadcast_in_dim.1263 = s32[1,41]{1,0} reshape(%iota.447), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/broadcast_in_dim" stack_frame_id=148} + %add.909 = s32[1,41]{1,0} add(%broadcast_in_dim.1263, %broadcast.43), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/add" stack_frame_id=151} + %lt.525 = pred[1,41]{1,0} compare(%lt.524, %add.909), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/lt" stack_frame_id=152} + %convert_element_type.1422 = s32[1,41]{1,0} convert(%lt.525), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/convert_element_type" stack_frame_id=156} + %broadcast_in_dim.1264 = s32[1,1,41]{2,1,0} reshape(%convert_element_type.1422), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/broadcast_in_dim" stack_frame_id=160} + %min.79 = s32[1,1,41]{2,1,0} minimum(%convert_element_type.1423, %broadcast_in_dim.1264), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/min" stack_frame_id=160} + %broadcast_in_dim.1269 = s32[1,1,1,41]{3,2,1,0} reshape(%min.79), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/broadcast_in_dim" stack_frame_id=341} + %convert_element_type.1433 = pred[1,1,1,41]{3,2,1,0} compare(%broadcast_in_dim.1269, %broadcast.37), direction=NE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/convert_element_type" stack_frame_id=345} + %broadcast_in_dim.1270 = pred[1,1,1,1,41]{4,3,2,1,0} reshape(%convert_element_type.1433), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/broadcast_in_dim" stack_frame_id=83} + %and.245 = pred[1,1,1,1,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1270), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/vmap()/and" stack_frame_id=83} + %and.246 = pred[1,1,1,41]{3,2,1,0} reshape(%and.245), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/vmap()/and" stack_frame_id=83} + %and.247 = pred[1,1,32,1,41]{4,3,2,1,0} broadcast(%and.246), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/vmap()/and" stack_frame_id=83} + %dot_general.873 = bf16[1,1,32,128]{3,2,1,0} dot(%convert_element_type.1426, %arg_tuple.5#153), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=186} + %convert_element_type.1427 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/convert_element_type" stack_frame_id=208} + %add.911 = f32[1,1]{1,0} reshape(%convert_element_type.1427), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/add"} + %iota.448 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/iota" stack_frame_id=197} + %mul.2595 = f32[64]{0} multiply(%iota.448, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=197} + %div.804 = f32[64]{0} divide(%mul.2595, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/div" stack_frame_id=198} + %neg.326 = f32[64]{0} negate(%div.804), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/neg" stack_frame_id=199} + %pow.320 = f32[64]{0} power(%broadcast.39, %neg.326), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/pow" stack_frame_id=202} + %div.805 = f32[64]{0} divide(%pow.320, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/div" stack_frame_id=203} + %dot_general.874 = f32[1,1,64]{2,1,0} dot(%add.911, %div.805), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/bi,j->bij/dot_general" stack_frame_id=211} + %stack.1260 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.874), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/stack" stack_frame_id=214} + %stack.1261 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.874), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/stack" stack_frame_id=214} + %stack.1262 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1260, %stack.1261), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/stack" stack_frame_id=214} + %reshape.586 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1262), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/reshape"} + %cos.159 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.586), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/cos" stack_frame_id=218} + %convert_element_type.1428 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.159), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/convert_element_type" stack_frame_id=222} + %mul.2596 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1428), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=242} + %mul.2597 = bf16[1,1,128]{2,1,0} reshape(%mul.2596), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=242} + %mul.2598 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2597), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=242} + %mul.2599 = bf16[1,1,32,128]{3,2,1,0} multiply(%dot_general.873, %mul.2598), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=242} + %split.319 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.873), slice={[0:1], [0:1], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/split" stack_frame_id=234} + %neg.327 = bf16[1,1,32,64]{3,2,1,0} negate(%split.319), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/neg" stack_frame_id=235} + %stack.1263 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%neg.327), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/stack" stack_frame_id=238} + %split.318 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.873), slice={[0:1], [0:1], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/split" stack_frame_id=234} + %stack.1264 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%split.318), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/stack" stack_frame_id=238} + %stack.1265 = bf16[1,1,32,2,64]{4,3,2,1,0} concatenate(%stack.1263, %stack.1264), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/stack" stack_frame_id=238} + %reshape.587 = bf16[1,1,32,128]{3,2,1,0} reshape(%stack.1265), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/reshape" stack_frame_id=241} + %sin.159 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.586), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/sin" stack_frame_id=226} + %convert_element_type.1429 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.159), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/convert_element_type" stack_frame_id=230} + %mul.2600 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1429), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=243} + %mul.2601 = bf16[1,1,128]{2,1,0} reshape(%mul.2600), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=243} + %mul.2602 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2601), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=243} + %mul.2603 = bf16[1,1,32,128]{3,2,1,0} multiply(%reshape.587, %mul.2602), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=243} + %add.912 = bf16[1,1,32,128]{3,2,1,0} add(%mul.2599, %mul.2603), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/add" stack_frame_id=244} + %reshape.592 = bf16[1,1,32,1,128]{4,3,2,1,0} reshape(%add.912), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/reshape" stack_frame_id=83} + %slice.145 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.148), slice={[0:1], [0:1], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/slice" stack_frame_id=245} + %squeeze.149 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.145), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/squeeze" stack_frame_id=245} + %dot_general.875 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1426, %arg_tuple.5#154), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=256} + %convert_element_type.1430 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/convert_element_type" stack_frame_id=287} + %add.913 = f32[1,1]{1,0} reshape(%convert_element_type.1430), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/add"} + %iota.449 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/iota" stack_frame_id=276} + %mul.2604 = f32[64]{0} multiply(%iota.449, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=276} + %div.806 = f32[64]{0} divide(%mul.2604, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/div" stack_frame_id=277} + %neg.328 = f32[64]{0} negate(%div.806), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/neg" stack_frame_id=278} + %pow.321 = f32[64]{0} power(%broadcast.39, %neg.328), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/pow" stack_frame_id=281} + %div.807 = f32[64]{0} divide(%pow.321, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/div" stack_frame_id=282} + %dot_general.877 = f32[1,1,64]{2,1,0} dot(%add.913, %div.807), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/bi,j->bij/dot_general" stack_frame_id=290} + %stack.1266 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.877), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/stack" stack_frame_id=293} + %stack.1267 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.877), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/stack" stack_frame_id=293} + %stack.1268 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1266, %stack.1267), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/stack" stack_frame_id=293} + %reshape.588 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1268), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/reshape"} + %cos.160 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.588), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/cos" stack_frame_id=297} + %convert_element_type.1431 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.160), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/convert_element_type" stack_frame_id=301} + %mul.2605 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1431), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=321} + %mul.2606 = bf16[1,1,128]{2,1,0} reshape(%mul.2605), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=321} + %mul.2607 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2606), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=321} + %mul.2608 = bf16[1,1,8,128]{3,2,1,0} multiply(%dot_general.875, %mul.2607), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=321} + %split.321 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.875), slice={[0:1], [0:1], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/split" stack_frame_id=313} + %neg.329 = bf16[1,1,8,64]{3,2,1,0} negate(%split.321), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/neg" stack_frame_id=314} + %stack.1269 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%neg.329), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/stack" stack_frame_id=317} + %split.320 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.875), slice={[0:1], [0:1], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/split" stack_frame_id=313} + %stack.1270 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%split.320), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/stack" stack_frame_id=317} + %stack.1271 = bf16[1,1,8,2,64]{4,3,2,1,0} concatenate(%stack.1269, %stack.1270), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/stack" stack_frame_id=317} + %reshape.589 = bf16[1,1,8,128]{3,2,1,0} reshape(%stack.1271), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/reshape" stack_frame_id=320} + %sin.160 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.588), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/sin" stack_frame_id=305} + %convert_element_type.1432 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.160), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/convert_element_type" stack_frame_id=309} + %mul.2609 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1432), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=322} + %mul.2610 = bf16[1,1,128]{2,1,0} reshape(%mul.2609), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=322} + %mul.2611 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2610), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=322} + %mul.2612 = bf16[1,1,8,128]{3,2,1,0} multiply(%reshape.589, %mul.2611), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=322} + %add.914 = bf16[1,1,8,128]{3,2,1,0} add(%mul.2608, %mul.2612), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/rotary_embedding_16/add" stack_frame_id=323} + %lt.526 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/lt" stack_frame_id=326} + %add.915 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/add" stack_frame_id=326} + %select_n.123 = s32[] select(%lt.526, %add.915, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/select_n" stack_frame_id=326} + %dynamic_update_slice.97 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.149, %add.914, %constant.127, %select_n.123, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/dynamic_update_slice" stack_frame_id=326} + %broadcast_in_dim.1267 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.97), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/broadcast_in_dim" stack_frame_id=335} + %reshape.590 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1267), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/reshape" stack_frame_id=335} + %dot_general.878 = f32[1,32,1,1,41]{4,3,2,1,0} dot(%reshape.592, %reshape.590), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=83} + %mul.2613 = f32[1,32,1,1,41]{4,3,2,1,0} multiply(%dot_general.878, %broadcast.36), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/vmap()/mul" stack_frame_id=83} + %vmap_jit__where__.79 = f32[1,1,32,1,41]{4,3,2,1,0} call(%and.247, %mul.2613, %constant.114), to_apply=%_where_3.137, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/vmap(jit(_where))" stack_frame_id=83} + %reduce_max.464 = f32[1,1,32,1]{3,2,1,0} reduce(%vmap_jit__where__.79, %constant.113), dimensions={4}, to_apply=%region_195.203, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/vmap()/reduce_max" stack_frame_id=83} + %max.88 = f32[1,1,32,1]{3,2,1,0} maximum(%reduce_max.464, %broadcast.35), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/vmap()/max" stack_frame_id=83} + %broadcast_in_dim.1271 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%max.88), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %sub.335 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1271), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/vmap()/sub" stack_frame_id=83} + %sub.336 = f32[1,1,32,1]{3,2,1,0} reshape(%sub.335), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/vmap()/sub" stack_frame_id=83} + %sub.337 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%sub.336), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/vmap()/sub" stack_frame_id=83} + %sub.338 = f32[1,1,32,1,41]{4,3,2,1,0} subtract(%vmap_jit__where__.79, %sub.337), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/vmap()/sub" stack_frame_id=83} + %exp.84 = f32[1,1,32,1,41]{4,3,2,1,0} exponential(%sub.338), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/vmap()/exp" stack_frame_id=83} + %reduce_sum.1400 = f32[1,1,32,1]{3,2,1,0} reduce(%exp.84, %constant.123), dimensions={4}, to_apply=%region_196.204, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/vmap()/reduce_sum" stack_frame_id=83} + %broadcast_in_dim.1272 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%reduce_sum.1400), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %div.808 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1272), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/vmap()/div" stack_frame_id=83} + %div.809 = f32[1,1,32,1]{3,2,1,0} reshape(%div.808), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/vmap()/div" stack_frame_id=83} + %div.810 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%div.809), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/vmap()/div" stack_frame_id=83} + %div.811 = f32[1,1,32,1,41]{4,3,2,1,0} divide(%exp.84, %div.810), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/vmap()/div" stack_frame_id=83} + %convert_element_type.1434 = bf16[1,1,32,1,41]{4,3,2,1,0} convert(%div.811), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/vmap()/convert_element_type" stack_frame_id=83} + %dot_general.879 = bf16[1,32,128,1,1]{4,3,2,1,0} dot(%reshape.591, %convert_element_type.1434), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=83} + %transpose.83 = bf16[1,1,32,128]{3,2,1,0} reshape(%dot_general.879), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/reshape;jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/vmap()/transpose;jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %dot_general.880 = bf16[1,1,4096]{2,1,0} dot(%transpose.83, %arg_tuple.5#156), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=354} + %add.917 = bf16[1,1,4096]{2,1,0} add(%dot_general.880, %add.906), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/add" stack_frame_id=355} + %convert_element_type.1435 = f32[1,1,4096]{2,1,0} convert(%add.917), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/feedforward_layernorm/convert_element_type" stack_frame_id=359} + %pow.322 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1435, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/feedforward_layernorm/pow" stack_frame_id=362} + %reduce_sum.1401 = f32[1,1]{1,0} reduce(%pow.322, %constant.123), dimensions={2}, to_apply=%region_197.205, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/feedforward_layernorm/reduce_sum" stack_frame_id=92} + %broadcast_in_dim.1273 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1401), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/feedforward_layernorm/broadcast_in_dim" stack_frame_id=92} + %div.812 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1273, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/feedforward_layernorm/div" stack_frame_id=92} + %add.918 = f32[1,1,1]{2,1,0} add(%div.812, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/feedforward_layernorm/add" stack_frame_id=363} + %rsqrt.161 = f32[1,1,1]{2,1,0} rsqrt(%add.918), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/feedforward_layernorm/rsqrt" stack_frame_id=366} + %mul.2614 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.161), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2615 = f32[1,1]{1,0} reshape(%mul.2614), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2616 = f32[1,1,4096]{2,1,0} broadcast(%mul.2615), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2617 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1435, %mul.2616), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/feedforward_layernorm/mul" stack_frame_id=367} + %convert_element_type.1436 = f32[4096]{0} convert(%arg_tuple.5#157), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/feedforward_layernorm/convert_element_type" stack_frame_id=368} + %broadcast_in_dim.1274 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1436), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/feedforward_layernorm/broadcast_in_dim" stack_frame_id=368} + %mul.2618 = f32[1,1,4096]{2,1,0} multiply(%mul.2617, %broadcast_in_dim.1274), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/feedforward_layernorm/mul" stack_frame_id=368} + %convert_element_type.1437 = bf16[1,1,4096]{2,1,0} convert(%mul.2618), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/feedforward_layernorm/convert_element_type" stack_frame_id=372} + %dot_general.882 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1437, %arg_tuple.5#159), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/feedforward_intermediate_dense/dot_general" stack_frame_id=401} + %dot_general.881 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1437, %arg_tuple.5#158), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/feedforward_gate_dense/dot_general" stack_frame_id=381} + %convert_element_type.1438 = f32[1,1,14336]{2,1,0} convert(%dot_general.881), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/convert_element_type" stack_frame_id=385} + %jit_silu_.79 = f32[1,1,14336]{2,1,0} call(%convert_element_type.1438), to_apply=%silu_4.141, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/jit(silu)" stack_frame_id=388} + %convert_element_type.1439 = bf16[1,1,14336]{2,1,0} convert(%jit_silu_.79), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/convert_element_type" stack_frame_id=392} + %mul.2619 = bf16[1,1,14336]{2,1,0} multiply(%dot_general.882, %convert_element_type.1439), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/mul" stack_frame_id=404} + %dot_general.883 = bf16[1,1,4096]{2,1,0} dot(%mul.2619, %arg_tuple.5#160), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/feedforward_output_dense/dot_general" stack_frame_id=413} + %add.919 = bf16[1,1,4096]{2,1,0} add(%dot_general.883, %add.917), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/add" stack_frame_id=414} + %convert_element_type.1443 = f32[1,1,4096]{2,1,0} convert(%add.919), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention_layernorm/convert_element_type" stack_frame_id=164} + %pow.323 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1443, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention_layernorm/pow" stack_frame_id=167} + %reduce_sum.1402 = f32[1,1]{1,0} reduce(%pow.323, %constant.123), dimensions={2}, to_apply=%region_198.206, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention_layernorm/reduce_sum" stack_frame_id=73} + %broadcast_in_dim.1279 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1402), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention_layernorm/broadcast_in_dim" stack_frame_id=73} + %div.813 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1279, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention_layernorm/div" stack_frame_id=73} + %add.923 = f32[1,1,1]{2,1,0} add(%div.813, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention_layernorm/add" stack_frame_id=168} + %rsqrt.162 = f32[1,1,1]{2,1,0} rsqrt(%add.923), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention_layernorm/rsqrt" stack_frame_id=171} + %mul.2620 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.162), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2621 = f32[1,1]{1,0} reshape(%mul.2620), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2622 = f32[1,1,4096]{2,1,0} broadcast(%mul.2621), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2623 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1443, %mul.2622), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention_layernorm/mul" stack_frame_id=172} + %convert_element_type.1444 = f32[4096]{0} convert(%arg_tuple.5#161), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention_layernorm/convert_element_type" stack_frame_id=173} + %broadcast_in_dim.1280 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1444), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention_layernorm/broadcast_in_dim" stack_frame_id=173} + %mul.2624 = f32[1,1,4096]{2,1,0} multiply(%mul.2623, %broadcast_in_dim.1280), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention_layernorm/mul" stack_frame_id=173} + %convert_element_type.1445 = bf16[1,1,4096]{2,1,0} convert(%mul.2624), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention_layernorm/convert_element_type" stack_frame_id=177} + %dot_general.887 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1445, %arg_tuple.5#164), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=265} + %lt.533 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/lt" stack_frame_id=329} + %add.929 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/add" stack_frame_id=329} + %select_n.126 = s32[] select(%lt.533, %add.929, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/select_n" stack_frame_id=329} + %dynamic_update_slice.100 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.153, %dot_general.887, %constant.127, %select_n.126, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/dynamic_update_slice" stack_frame_id=329} + %broadcast_in_dim.1282 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.100), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/broadcast_in_dim" stack_frame_id=338} + %reshape.598 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1282), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/reshape" stack_frame_id=338} + %convert_element_type.1440 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/convert_element_type" stack_frame_id=139} + %add.920 = f32[1,1]{1,0} reshape(%convert_element_type.1440), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_17/add"} + %ge.421 = f32[1,1]{1,0} broadcast(%add.920), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/ge" stack_frame_id=143} + %ge.422 = f32[1]{0} reshape(%ge.421), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/ge" stack_frame_id=143} + %ge.423 = f32[1,41]{1,0} broadcast(%ge.422), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/ge" stack_frame_id=143} + %iota.450 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/iota" stack_frame_id=142} + %broadcast_in_dim.1275 = f32[1,41]{1,0} reshape(%iota.450), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/broadcast_in_dim" stack_frame_id=143} + %ge.424 = pred[1,41]{1,0} compare(%ge.423, %broadcast_in_dim.1275), direction=GE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/ge" stack_frame_id=143} + %broadcast_in_dim.1276 = pred[1,1,41]{2,1,0} reshape(%ge.424), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/broadcast_in_dim" stack_frame_id=146} + %convert_element_type.1442 = s32[1,1,41]{2,1,0} convert(%broadcast_in_dim.1276), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/convert_element_type" stack_frame_id=160} + %add.921 = s32[1,1]{1,0} reshape(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/add" stack_frame_id=147} + %lt.528 = s32[1,1]{1,0} broadcast(%add.921), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/lt" stack_frame_id=152} + %lt.529 = s32[1]{0} reshape(%lt.528), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/lt" stack_frame_id=152} + %lt.530 = s32[1,41]{1,0} broadcast(%lt.529), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/lt" stack_frame_id=152} + %iota.451 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/iota" stack_frame_id=150} + %broadcast_in_dim.1277 = s32[1,41]{1,0} reshape(%iota.451), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/broadcast_in_dim" stack_frame_id=148} + %add.922 = s32[1,41]{1,0} add(%broadcast_in_dim.1277, %broadcast.43), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/add" stack_frame_id=151} + %lt.531 = pred[1,41]{1,0} compare(%lt.530, %add.922), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/lt" stack_frame_id=152} + %convert_element_type.1441 = s32[1,41]{1,0} convert(%lt.531), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/convert_element_type" stack_frame_id=156} + %broadcast_in_dim.1278 = s32[1,1,41]{2,1,0} reshape(%convert_element_type.1441), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/broadcast_in_dim" stack_frame_id=160} + %min.80 = s32[1,1,41]{2,1,0} minimum(%convert_element_type.1442, %broadcast_in_dim.1278), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/min" stack_frame_id=160} + %broadcast_in_dim.1283 = s32[1,1,1,41]{3,2,1,0} reshape(%min.80), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/broadcast_in_dim" stack_frame_id=341} + %convert_element_type.1452 = pred[1,1,1,41]{3,2,1,0} compare(%broadcast_in_dim.1283, %broadcast.37), direction=NE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/convert_element_type" stack_frame_id=345} + %broadcast_in_dim.1284 = pred[1,1,1,1,41]{4,3,2,1,0} reshape(%convert_element_type.1452), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/broadcast_in_dim" stack_frame_id=83} + %and.248 = pred[1,1,1,1,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1284), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/vmap()/and" stack_frame_id=83} + %and.249 = pred[1,1,1,41]{3,2,1,0} reshape(%and.248), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/vmap()/and" stack_frame_id=83} + %and.250 = pred[1,1,32,1,41]{4,3,2,1,0} broadcast(%and.249), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/vmap()/and" stack_frame_id=83} + %dot_general.884 = bf16[1,1,32,128]{3,2,1,0} dot(%convert_element_type.1445, %arg_tuple.5#162), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=186} + %convert_element_type.1446 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/convert_element_type" stack_frame_id=208} + %add.924 = f32[1,1]{1,0} reshape(%convert_element_type.1446), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/add"} + %iota.452 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/iota" stack_frame_id=197} + %mul.2625 = f32[64]{0} multiply(%iota.452, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=197} + %div.814 = f32[64]{0} divide(%mul.2625, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/div" stack_frame_id=198} + %neg.330 = f32[64]{0} negate(%div.814), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/neg" stack_frame_id=199} + %pow.324 = f32[64]{0} power(%broadcast.39, %neg.330), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/pow" stack_frame_id=202} + %div.815 = f32[64]{0} divide(%pow.324, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/div" stack_frame_id=203} + %dot_general.885 = f32[1,1,64]{2,1,0} dot(%add.924, %div.815), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/bi,j->bij/dot_general" stack_frame_id=211} + %stack.1275 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.885), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/stack" stack_frame_id=214} + %stack.1276 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.885), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/stack" stack_frame_id=214} + %stack.1277 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1275, %stack.1276), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/stack" stack_frame_id=214} + %reshape.593 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1277), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/reshape"} + %cos.161 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.593), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/cos" stack_frame_id=218} + %convert_element_type.1447 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.161), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/convert_element_type" stack_frame_id=222} + %mul.2626 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1447), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=242} + %mul.2627 = bf16[1,1,128]{2,1,0} reshape(%mul.2626), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=242} + %mul.2628 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2627), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=242} + %mul.2629 = bf16[1,1,32,128]{3,2,1,0} multiply(%dot_general.884, %mul.2628), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=242} + %split.323 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.884), slice={[0:1], [0:1], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/split" stack_frame_id=234} + %neg.331 = bf16[1,1,32,64]{3,2,1,0} negate(%split.323), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/neg" stack_frame_id=235} + %stack.1278 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%neg.331), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/stack" stack_frame_id=238} + %split.322 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.884), slice={[0:1], [0:1], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/split" stack_frame_id=234} + %stack.1279 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%split.322), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/stack" stack_frame_id=238} + %stack.1280 = bf16[1,1,32,2,64]{4,3,2,1,0} concatenate(%stack.1278, %stack.1279), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/stack" stack_frame_id=238} + %reshape.594 = bf16[1,1,32,128]{3,2,1,0} reshape(%stack.1280), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/reshape" stack_frame_id=241} + %sin.161 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.593), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/sin" stack_frame_id=226} + %convert_element_type.1448 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.161), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/convert_element_type" stack_frame_id=230} + %mul.2630 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1448), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=243} + %mul.2631 = bf16[1,1,128]{2,1,0} reshape(%mul.2630), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=243} + %mul.2632 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2631), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=243} + %mul.2633 = bf16[1,1,32,128]{3,2,1,0} multiply(%reshape.594, %mul.2632), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=243} + %add.925 = bf16[1,1,32,128]{3,2,1,0} add(%mul.2629, %mul.2633), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/add" stack_frame_id=244} + %reshape.599 = bf16[1,1,32,1,128]{4,3,2,1,0} reshape(%add.925), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/reshape" stack_frame_id=83} + %slice.148 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.151), slice={[0:1], [0:1], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/slice" stack_frame_id=245} + %squeeze.152 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.148), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/squeeze" stack_frame_id=245} + %dot_general.886 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1445, %arg_tuple.5#163), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=256} + %convert_element_type.1449 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/convert_element_type" stack_frame_id=287} + %add.926 = f32[1,1]{1,0} reshape(%convert_element_type.1449), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/add"} + %iota.453 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/iota" stack_frame_id=276} + %mul.2634 = f32[64]{0} multiply(%iota.453, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=276} + %div.816 = f32[64]{0} divide(%mul.2634, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/div" stack_frame_id=277} + %neg.332 = f32[64]{0} negate(%div.816), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/neg" stack_frame_id=278} + %pow.325 = f32[64]{0} power(%broadcast.39, %neg.332), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/pow" stack_frame_id=281} + %div.817 = f32[64]{0} divide(%pow.325, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/div" stack_frame_id=282} + %dot_general.888 = f32[1,1,64]{2,1,0} dot(%add.926, %div.817), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/bi,j->bij/dot_general" stack_frame_id=290} + %stack.1281 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.888), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/stack" stack_frame_id=293} + %stack.1282 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.888), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/stack" stack_frame_id=293} + %stack.1283 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1281, %stack.1282), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/stack" stack_frame_id=293} + %reshape.595 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1283), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/reshape"} + %cos.162 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.595), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/cos" stack_frame_id=297} + %convert_element_type.1450 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.162), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/convert_element_type" stack_frame_id=301} + %mul.2635 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1450), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=321} + %mul.2636 = bf16[1,1,128]{2,1,0} reshape(%mul.2635), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=321} + %mul.2637 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2636), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=321} + %mul.2638 = bf16[1,1,8,128]{3,2,1,0} multiply(%dot_general.886, %mul.2637), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=321} + %split.325 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.886), slice={[0:1], [0:1], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/split" stack_frame_id=313} + %neg.333 = bf16[1,1,8,64]{3,2,1,0} negate(%split.325), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/neg" stack_frame_id=314} + %stack.1284 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%neg.333), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/stack" stack_frame_id=317} + %split.324 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.886), slice={[0:1], [0:1], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/split" stack_frame_id=313} + %stack.1285 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%split.324), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/stack" stack_frame_id=317} + %stack.1286 = bf16[1,1,8,2,64]{4,3,2,1,0} concatenate(%stack.1284, %stack.1285), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/stack" stack_frame_id=317} + %reshape.596 = bf16[1,1,8,128]{3,2,1,0} reshape(%stack.1286), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/reshape" stack_frame_id=320} + %sin.162 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.595), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/sin" stack_frame_id=305} + %convert_element_type.1451 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.162), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/convert_element_type" stack_frame_id=309} + %mul.2639 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1451), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=322} + %mul.2640 = bf16[1,1,128]{2,1,0} reshape(%mul.2639), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=322} + %mul.2641 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2640), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=322} + %mul.2642 = bf16[1,1,8,128]{3,2,1,0} multiply(%reshape.596, %mul.2641), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=322} + %add.927 = bf16[1,1,8,128]{3,2,1,0} add(%mul.2638, %mul.2642), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/rotary_embedding_17/add" stack_frame_id=323} + %lt.532 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/lt" stack_frame_id=326} + %add.928 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/add" stack_frame_id=326} + %select_n.125 = s32[] select(%lt.532, %add.928, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/select_n" stack_frame_id=326} + %dynamic_update_slice.99 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.152, %add.927, %constant.127, %select_n.125, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/dynamic_update_slice" stack_frame_id=326} + %broadcast_in_dim.1281 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.99), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/broadcast_in_dim" stack_frame_id=335} + %reshape.597 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1281), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/reshape" stack_frame_id=335} + %dot_general.889 = f32[1,32,1,1,41]{4,3,2,1,0} dot(%reshape.599, %reshape.597), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=83} + %mul.2643 = f32[1,32,1,1,41]{4,3,2,1,0} multiply(%dot_general.889, %broadcast.36), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/vmap()/mul" stack_frame_id=83} + %vmap_jit__where__.80 = f32[1,1,32,1,41]{4,3,2,1,0} call(%and.250, %mul.2643, %constant.114), to_apply=%_where_3.137, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/vmap(jit(_where))" stack_frame_id=83} + %reduce_max.465 = f32[1,1,32,1]{3,2,1,0} reduce(%vmap_jit__where__.80, %constant.113), dimensions={4}, to_apply=%region_199.207, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/vmap()/reduce_max" stack_frame_id=83} + %max.89 = f32[1,1,32,1]{3,2,1,0} maximum(%reduce_max.465, %broadcast.35), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/vmap()/max" stack_frame_id=83} + %broadcast_in_dim.1285 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%max.89), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %sub.339 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1285), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/vmap()/sub" stack_frame_id=83} + %sub.340 = f32[1,1,32,1]{3,2,1,0} reshape(%sub.339), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/vmap()/sub" stack_frame_id=83} + %sub.341 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%sub.340), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/vmap()/sub" stack_frame_id=83} + %sub.342 = f32[1,1,32,1,41]{4,3,2,1,0} subtract(%vmap_jit__where__.80, %sub.341), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/vmap()/sub" stack_frame_id=83} + %exp.85 = f32[1,1,32,1,41]{4,3,2,1,0} exponential(%sub.342), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/vmap()/exp" stack_frame_id=83} + %reduce_sum.1403 = f32[1,1,32,1]{3,2,1,0} reduce(%exp.85, %constant.123), dimensions={4}, to_apply=%region_200.208, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/vmap()/reduce_sum" stack_frame_id=83} + %broadcast_in_dim.1286 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%reduce_sum.1403), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %div.818 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1286), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/vmap()/div" stack_frame_id=83} + %div.819 = f32[1,1,32,1]{3,2,1,0} reshape(%div.818), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/vmap()/div" stack_frame_id=83} + %div.820 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%div.819), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/vmap()/div" stack_frame_id=83} + %div.821 = f32[1,1,32,1,41]{4,3,2,1,0} divide(%exp.85, %div.820), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/vmap()/div" stack_frame_id=83} + %convert_element_type.1453 = bf16[1,1,32,1,41]{4,3,2,1,0} convert(%div.821), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/vmap()/convert_element_type" stack_frame_id=83} + %dot_general.890 = bf16[1,32,128,1,1]{4,3,2,1,0} dot(%reshape.598, %convert_element_type.1453), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=83} + %transpose.84 = bf16[1,1,32,128]{3,2,1,0} reshape(%dot_general.890), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/reshape;jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/vmap()/transpose;jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %dot_general.891 = bf16[1,1,4096]{2,1,0} dot(%transpose.84, %arg_tuple.5#165), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=354} + %add.930 = bf16[1,1,4096]{2,1,0} add(%dot_general.891, %add.919), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/add" stack_frame_id=355} + %convert_element_type.1454 = f32[1,1,4096]{2,1,0} convert(%add.930), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/feedforward_layernorm/convert_element_type" stack_frame_id=359} + %pow.326 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1454, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/feedforward_layernorm/pow" stack_frame_id=362} + %reduce_sum.1404 = f32[1,1]{1,0} reduce(%pow.326, %constant.123), dimensions={2}, to_apply=%region_201.209, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/feedforward_layernorm/reduce_sum" stack_frame_id=92} + %broadcast_in_dim.1287 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1404), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/feedforward_layernorm/broadcast_in_dim" stack_frame_id=92} + %div.822 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1287, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/feedforward_layernorm/div" stack_frame_id=92} + %add.931 = f32[1,1,1]{2,1,0} add(%div.822, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/feedforward_layernorm/add" stack_frame_id=363} + %rsqrt.163 = f32[1,1,1]{2,1,0} rsqrt(%add.931), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/feedforward_layernorm/rsqrt" stack_frame_id=366} + %mul.2644 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.163), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2645 = f32[1,1]{1,0} reshape(%mul.2644), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2646 = f32[1,1,4096]{2,1,0} broadcast(%mul.2645), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2647 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1454, %mul.2646), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/feedforward_layernorm/mul" stack_frame_id=367} + %convert_element_type.1455 = f32[4096]{0} convert(%arg_tuple.5#166), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/feedforward_layernorm/convert_element_type" stack_frame_id=368} + %broadcast_in_dim.1288 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1455), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/feedforward_layernorm/broadcast_in_dim" stack_frame_id=368} + %mul.2648 = f32[1,1,4096]{2,1,0} multiply(%mul.2647, %broadcast_in_dim.1288), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/feedforward_layernorm/mul" stack_frame_id=368} + %convert_element_type.1456 = bf16[1,1,4096]{2,1,0} convert(%mul.2648), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/feedforward_layernorm/convert_element_type" stack_frame_id=372} + %dot_general.893 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1456, %arg_tuple.5#168), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/feedforward_intermediate_dense/dot_general" stack_frame_id=401} + %dot_general.892 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1456, %arg_tuple.5#167), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/feedforward_gate_dense/dot_general" stack_frame_id=381} + %convert_element_type.1457 = f32[1,1,14336]{2,1,0} convert(%dot_general.892), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/convert_element_type" stack_frame_id=385} + %jit_silu_.80 = f32[1,1,14336]{2,1,0} call(%convert_element_type.1457), to_apply=%silu_4.141, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/jit(silu)" stack_frame_id=388} + %convert_element_type.1458 = bf16[1,1,14336]{2,1,0} convert(%jit_silu_.80), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/convert_element_type" stack_frame_id=392} + %mul.2649 = bf16[1,1,14336]{2,1,0} multiply(%dot_general.893, %convert_element_type.1458), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/mul" stack_frame_id=404} + %dot_general.894 = bf16[1,1,4096]{2,1,0} dot(%mul.2649, %arg_tuple.5#169), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/feedforward_output_dense/dot_general" stack_frame_id=413} + %add.932 = bf16[1,1,4096]{2,1,0} add(%dot_general.894, %add.930), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/add" stack_frame_id=414} + %convert_element_type.1462 = f32[1,1,4096]{2,1,0} convert(%add.932), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention_layernorm/convert_element_type" stack_frame_id=164} + %pow.327 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1462, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention_layernorm/pow" stack_frame_id=167} + %reduce_sum.1405 = f32[1,1]{1,0} reduce(%pow.327, %constant.123), dimensions={2}, to_apply=%region_202.210, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention_layernorm/reduce_sum" stack_frame_id=73} + %broadcast_in_dim.1293 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1405), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention_layernorm/broadcast_in_dim" stack_frame_id=73} + %div.823 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1293, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention_layernorm/div" stack_frame_id=73} + %add.936 = f32[1,1,1]{2,1,0} add(%div.823, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention_layernorm/add" stack_frame_id=168} + %rsqrt.164 = f32[1,1,1]{2,1,0} rsqrt(%add.936), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention_layernorm/rsqrt" stack_frame_id=171} + %mul.2650 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.164), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2651 = f32[1,1]{1,0} reshape(%mul.2650), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2652 = f32[1,1,4096]{2,1,0} broadcast(%mul.2651), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2653 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1462, %mul.2652), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention_layernorm/mul" stack_frame_id=172} + %convert_element_type.1463 = f32[4096]{0} convert(%arg_tuple.5#170), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention_layernorm/convert_element_type" stack_frame_id=173} + %broadcast_in_dim.1294 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1463), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention_layernorm/broadcast_in_dim" stack_frame_id=173} + %mul.2654 = f32[1,1,4096]{2,1,0} multiply(%mul.2653, %broadcast_in_dim.1294), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention_layernorm/mul" stack_frame_id=173} + %convert_element_type.1464 = bf16[1,1,4096]{2,1,0} convert(%mul.2654), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention_layernorm/convert_element_type" stack_frame_id=177} + %dot_general.898 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1464, %arg_tuple.5#173), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=265} + %lt.539 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/lt" stack_frame_id=329} + %add.942 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/add" stack_frame_id=329} + %select_n.128 = s32[] select(%lt.539, %add.942, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/select_n" stack_frame_id=329} + %dynamic_update_slice.102 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.156, %dot_general.898, %constant.127, %select_n.128, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/dynamic_update_slice" stack_frame_id=329} + %broadcast_in_dim.1296 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.102), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/broadcast_in_dim" stack_frame_id=338} + %reshape.605 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1296), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/reshape" stack_frame_id=338} + %convert_element_type.1459 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/convert_element_type" stack_frame_id=139} + %add.933 = f32[1,1]{1,0} reshape(%convert_element_type.1459), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_18/add"} + %ge.425 = f32[1,1]{1,0} broadcast(%add.933), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/ge" stack_frame_id=143} + %ge.426 = f32[1]{0} reshape(%ge.425), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/ge" stack_frame_id=143} + %ge.427 = f32[1,41]{1,0} broadcast(%ge.426), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/ge" stack_frame_id=143} + %iota.454 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/iota" stack_frame_id=142} + %broadcast_in_dim.1289 = f32[1,41]{1,0} reshape(%iota.454), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/broadcast_in_dim" stack_frame_id=143} + %ge.428 = pred[1,41]{1,0} compare(%ge.427, %broadcast_in_dim.1289), direction=GE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/ge" stack_frame_id=143} + %broadcast_in_dim.1290 = pred[1,1,41]{2,1,0} reshape(%ge.428), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/broadcast_in_dim" stack_frame_id=146} + %convert_element_type.1461 = s32[1,1,41]{2,1,0} convert(%broadcast_in_dim.1290), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/convert_element_type" stack_frame_id=160} + %add.934 = s32[1,1]{1,0} reshape(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/add" stack_frame_id=147} + %lt.534 = s32[1,1]{1,0} broadcast(%add.934), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/lt" stack_frame_id=152} + %lt.535 = s32[1]{0} reshape(%lt.534), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/lt" stack_frame_id=152} + %lt.536 = s32[1,41]{1,0} broadcast(%lt.535), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/lt" stack_frame_id=152} + %iota.455 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/iota" stack_frame_id=150} + %broadcast_in_dim.1291 = s32[1,41]{1,0} reshape(%iota.455), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/broadcast_in_dim" stack_frame_id=148} + %add.935 = s32[1,41]{1,0} add(%broadcast_in_dim.1291, %broadcast.43), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/add" stack_frame_id=151} + %lt.537 = pred[1,41]{1,0} compare(%lt.536, %add.935), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/lt" stack_frame_id=152} + %convert_element_type.1460 = s32[1,41]{1,0} convert(%lt.537), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/convert_element_type" stack_frame_id=156} + %broadcast_in_dim.1292 = s32[1,1,41]{2,1,0} reshape(%convert_element_type.1460), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/broadcast_in_dim" stack_frame_id=160} + %min.81 = s32[1,1,41]{2,1,0} minimum(%convert_element_type.1461, %broadcast_in_dim.1292), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/min" stack_frame_id=160} + %broadcast_in_dim.1297 = s32[1,1,1,41]{3,2,1,0} reshape(%min.81), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/broadcast_in_dim" stack_frame_id=341} + %convert_element_type.1471 = pred[1,1,1,41]{3,2,1,0} compare(%broadcast_in_dim.1297, %broadcast.37), direction=NE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/convert_element_type" stack_frame_id=345} + %broadcast_in_dim.1298 = pred[1,1,1,1,41]{4,3,2,1,0} reshape(%convert_element_type.1471), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/broadcast_in_dim" stack_frame_id=83} + %and.251 = pred[1,1,1,1,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1298), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/vmap()/and" stack_frame_id=83} + %and.252 = pred[1,1,1,41]{3,2,1,0} reshape(%and.251), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/vmap()/and" stack_frame_id=83} + %and.253 = pred[1,1,32,1,41]{4,3,2,1,0} broadcast(%and.252), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/vmap()/and" stack_frame_id=83} + %dot_general.895 = bf16[1,1,32,128]{3,2,1,0} dot(%convert_element_type.1464, %arg_tuple.5#171), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=186} + %convert_element_type.1465 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/convert_element_type" stack_frame_id=208} + %add.937 = f32[1,1]{1,0} reshape(%convert_element_type.1465), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/add"} + %iota.456 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/iota" stack_frame_id=197} + %mul.2655 = f32[64]{0} multiply(%iota.456, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=197} + %div.824 = f32[64]{0} divide(%mul.2655, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/div" stack_frame_id=198} + %neg.334 = f32[64]{0} negate(%div.824), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/neg" stack_frame_id=199} + %pow.328 = f32[64]{0} power(%broadcast.39, %neg.334), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/pow" stack_frame_id=202} + %div.825 = f32[64]{0} divide(%pow.328, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/div" stack_frame_id=203} + %dot_general.896 = f32[1,1,64]{2,1,0} dot(%add.937, %div.825), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/bi,j->bij/dot_general" stack_frame_id=211} + %stack.1290 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.896), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/stack" stack_frame_id=214} + %stack.1291 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.896), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/stack" stack_frame_id=214} + %stack.1292 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1290, %stack.1291), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/stack" stack_frame_id=214} + %reshape.600 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1292), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/reshape"} + %cos.163 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.600), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/cos" stack_frame_id=218} + %convert_element_type.1466 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.163), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/convert_element_type" stack_frame_id=222} + %mul.2656 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1466), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=242} + %mul.2657 = bf16[1,1,128]{2,1,0} reshape(%mul.2656), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=242} + %mul.2658 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2657), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=242} + %mul.2659 = bf16[1,1,32,128]{3,2,1,0} multiply(%dot_general.895, %mul.2658), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=242} + %split.327 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.895), slice={[0:1], [0:1], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/split" stack_frame_id=234} + %neg.335 = bf16[1,1,32,64]{3,2,1,0} negate(%split.327), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/neg" stack_frame_id=235} + %stack.1293 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%neg.335), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/stack" stack_frame_id=238} + %split.326 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.895), slice={[0:1], [0:1], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/split" stack_frame_id=234} + %stack.1294 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%split.326), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/stack" stack_frame_id=238} + %stack.1295 = bf16[1,1,32,2,64]{4,3,2,1,0} concatenate(%stack.1293, %stack.1294), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/stack" stack_frame_id=238} + %reshape.601 = bf16[1,1,32,128]{3,2,1,0} reshape(%stack.1295), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/reshape" stack_frame_id=241} + %sin.163 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.600), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/sin" stack_frame_id=226} + %convert_element_type.1467 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.163), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/convert_element_type" stack_frame_id=230} + %mul.2660 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1467), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=243} + %mul.2661 = bf16[1,1,128]{2,1,0} reshape(%mul.2660), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=243} + %mul.2662 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2661), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=243} + %mul.2663 = bf16[1,1,32,128]{3,2,1,0} multiply(%reshape.601, %mul.2662), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=243} + %add.938 = bf16[1,1,32,128]{3,2,1,0} add(%mul.2659, %mul.2663), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/add" stack_frame_id=244} + %reshape.606 = bf16[1,1,32,1,128]{4,3,2,1,0} reshape(%add.938), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/reshape" stack_frame_id=83} + %slice.151 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.154), slice={[0:1], [0:1], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/slice" stack_frame_id=245} + %squeeze.155 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.151), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/squeeze" stack_frame_id=245} + %dot_general.897 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1464, %arg_tuple.5#172), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=256} + %convert_element_type.1468 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/convert_element_type" stack_frame_id=287} + %add.939 = f32[1,1]{1,0} reshape(%convert_element_type.1468), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/add"} + %iota.457 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/iota" stack_frame_id=276} + %mul.2664 = f32[64]{0} multiply(%iota.457, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=276} + %div.826 = f32[64]{0} divide(%mul.2664, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/div" stack_frame_id=277} + %neg.336 = f32[64]{0} negate(%div.826), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/neg" stack_frame_id=278} + %pow.329 = f32[64]{0} power(%broadcast.39, %neg.336), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/pow" stack_frame_id=281} + %div.827 = f32[64]{0} divide(%pow.329, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/div" stack_frame_id=282} + %dot_general.899 = f32[1,1,64]{2,1,0} dot(%add.939, %div.827), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/bi,j->bij/dot_general" stack_frame_id=290} + %stack.1296 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.899), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/stack" stack_frame_id=293} + %stack.1297 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.899), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/stack" stack_frame_id=293} + %stack.1298 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1296, %stack.1297), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/stack" stack_frame_id=293} + %reshape.602 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1298), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/reshape"} + %cos.164 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.602), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/cos" stack_frame_id=297} + %convert_element_type.1469 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.164), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/convert_element_type" stack_frame_id=301} + %mul.2665 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1469), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=321} + %mul.2666 = bf16[1,1,128]{2,1,0} reshape(%mul.2665), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=321} + %mul.2667 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2666), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=321} + %mul.2668 = bf16[1,1,8,128]{3,2,1,0} multiply(%dot_general.897, %mul.2667), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=321} + %split.329 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.897), slice={[0:1], [0:1], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/split" stack_frame_id=313} + %neg.337 = bf16[1,1,8,64]{3,2,1,0} negate(%split.329), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/neg" stack_frame_id=314} + %stack.1299 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%neg.337), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/stack" stack_frame_id=317} + %split.328 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.897), slice={[0:1], [0:1], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/split" stack_frame_id=313} + %stack.1300 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%split.328), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/stack" stack_frame_id=317} + %stack.1301 = bf16[1,1,8,2,64]{4,3,2,1,0} concatenate(%stack.1299, %stack.1300), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/stack" stack_frame_id=317} + %reshape.603 = bf16[1,1,8,128]{3,2,1,0} reshape(%stack.1301), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/reshape" stack_frame_id=320} + %sin.164 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.602), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/sin" stack_frame_id=305} + %convert_element_type.1470 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.164), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/convert_element_type" stack_frame_id=309} + %mul.2669 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1470), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=322} + %mul.2670 = bf16[1,1,128]{2,1,0} reshape(%mul.2669), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=322} + %mul.2671 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2670), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=322} + %mul.2672 = bf16[1,1,8,128]{3,2,1,0} multiply(%reshape.603, %mul.2671), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=322} + %add.940 = bf16[1,1,8,128]{3,2,1,0} add(%mul.2668, %mul.2672), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/rotary_embedding_18/add" stack_frame_id=323} + %lt.538 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/lt" stack_frame_id=326} + %add.941 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/add" stack_frame_id=326} + %select_n.127 = s32[] select(%lt.538, %add.941, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/select_n" stack_frame_id=326} + %dynamic_update_slice.101 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.155, %add.940, %constant.127, %select_n.127, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/dynamic_update_slice" stack_frame_id=326} + %broadcast_in_dim.1295 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.101), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/broadcast_in_dim" stack_frame_id=335} + %reshape.604 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1295), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/reshape" stack_frame_id=335} + %dot_general.900 = f32[1,32,1,1,41]{4,3,2,1,0} dot(%reshape.606, %reshape.604), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=83} + %mul.2673 = f32[1,32,1,1,41]{4,3,2,1,0} multiply(%dot_general.900, %broadcast.36), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/vmap()/mul" stack_frame_id=83} + %vmap_jit__where__.81 = f32[1,1,32,1,41]{4,3,2,1,0} call(%and.253, %mul.2673, %constant.114), to_apply=%_where_3.137, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/vmap(jit(_where))" stack_frame_id=83} + %reduce_max.466 = f32[1,1,32,1]{3,2,1,0} reduce(%vmap_jit__where__.81, %constant.113), dimensions={4}, to_apply=%region_203.211, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/vmap()/reduce_max" stack_frame_id=83} + %max.90 = f32[1,1,32,1]{3,2,1,0} maximum(%reduce_max.466, %broadcast.35), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/vmap()/max" stack_frame_id=83} + %broadcast_in_dim.1299 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%max.90), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %sub.343 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1299), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/vmap()/sub" stack_frame_id=83} + %sub.344 = f32[1,1,32,1]{3,2,1,0} reshape(%sub.343), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/vmap()/sub" stack_frame_id=83} + %sub.345 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%sub.344), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/vmap()/sub" stack_frame_id=83} + %sub.346 = f32[1,1,32,1,41]{4,3,2,1,0} subtract(%vmap_jit__where__.81, %sub.345), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/vmap()/sub" stack_frame_id=83} + %exp.86 = f32[1,1,32,1,41]{4,3,2,1,0} exponential(%sub.346), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/vmap()/exp" stack_frame_id=83} + %reduce_sum.1406 = f32[1,1,32,1]{3,2,1,0} reduce(%exp.86, %constant.123), dimensions={4}, to_apply=%region_204.212, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/vmap()/reduce_sum" stack_frame_id=83} + %broadcast_in_dim.1300 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%reduce_sum.1406), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %div.828 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1300), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/vmap()/div" stack_frame_id=83} + %div.829 = f32[1,1,32,1]{3,2,1,0} reshape(%div.828), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/vmap()/div" stack_frame_id=83} + %div.830 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%div.829), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/vmap()/div" stack_frame_id=83} + %div.831 = f32[1,1,32,1,41]{4,3,2,1,0} divide(%exp.86, %div.830), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/vmap()/div" stack_frame_id=83} + %convert_element_type.1472 = bf16[1,1,32,1,41]{4,3,2,1,0} convert(%div.831), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/vmap()/convert_element_type" stack_frame_id=83} + %dot_general.901 = bf16[1,32,128,1,1]{4,3,2,1,0} dot(%reshape.605, %convert_element_type.1472), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=83} + %transpose.85 = bf16[1,1,32,128]{3,2,1,0} reshape(%dot_general.901), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/reshape;jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/vmap()/transpose;jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %dot_general.902 = bf16[1,1,4096]{2,1,0} dot(%transpose.85, %arg_tuple.5#174), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=354} + %add.943 = bf16[1,1,4096]{2,1,0} add(%dot_general.902, %add.932), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/add" stack_frame_id=355} + %convert_element_type.1473 = f32[1,1,4096]{2,1,0} convert(%add.943), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/feedforward_layernorm/convert_element_type" stack_frame_id=359} + %pow.330 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1473, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/feedforward_layernorm/pow" stack_frame_id=362} + %reduce_sum.1407 = f32[1,1]{1,0} reduce(%pow.330, %constant.123), dimensions={2}, to_apply=%region_205.213, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/feedforward_layernorm/reduce_sum" stack_frame_id=92} + %broadcast_in_dim.1301 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1407), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/feedforward_layernorm/broadcast_in_dim" stack_frame_id=92} + %div.832 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1301, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/feedforward_layernorm/div" stack_frame_id=92} + %add.944 = f32[1,1,1]{2,1,0} add(%div.832, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/feedforward_layernorm/add" stack_frame_id=363} + %rsqrt.165 = f32[1,1,1]{2,1,0} rsqrt(%add.944), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/feedforward_layernorm/rsqrt" stack_frame_id=366} + %mul.2674 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.165), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2675 = f32[1,1]{1,0} reshape(%mul.2674), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2676 = f32[1,1,4096]{2,1,0} broadcast(%mul.2675), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2677 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1473, %mul.2676), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/feedforward_layernorm/mul" stack_frame_id=367} + %convert_element_type.1474 = f32[4096]{0} convert(%arg_tuple.5#175), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/feedforward_layernorm/convert_element_type" stack_frame_id=368} + %broadcast_in_dim.1302 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1474), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/feedforward_layernorm/broadcast_in_dim" stack_frame_id=368} + %mul.2678 = f32[1,1,4096]{2,1,0} multiply(%mul.2677, %broadcast_in_dim.1302), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/feedforward_layernorm/mul" stack_frame_id=368} + %convert_element_type.1475 = bf16[1,1,4096]{2,1,0} convert(%mul.2678), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/feedforward_layernorm/convert_element_type" stack_frame_id=372} + %dot_general.904 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1475, %arg_tuple.5#177), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/feedforward_intermediate_dense/dot_general" stack_frame_id=401} + %dot_general.903 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1475, %arg_tuple.5#176), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/feedforward_gate_dense/dot_general" stack_frame_id=381} + %convert_element_type.1476 = f32[1,1,14336]{2,1,0} convert(%dot_general.903), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/convert_element_type" stack_frame_id=385} + %jit_silu_.81 = f32[1,1,14336]{2,1,0} call(%convert_element_type.1476), to_apply=%silu_4.141, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/jit(silu)" stack_frame_id=388} + %convert_element_type.1477 = bf16[1,1,14336]{2,1,0} convert(%jit_silu_.81), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/convert_element_type" stack_frame_id=392} + %mul.2679 = bf16[1,1,14336]{2,1,0} multiply(%dot_general.904, %convert_element_type.1477), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/mul" stack_frame_id=404} + %dot_general.905 = bf16[1,1,4096]{2,1,0} dot(%mul.2679, %arg_tuple.5#178), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/feedforward_output_dense/dot_general" stack_frame_id=413} + %add.945 = bf16[1,1,4096]{2,1,0} add(%dot_general.905, %add.943), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/add" stack_frame_id=414} + %convert_element_type.1481 = f32[1,1,4096]{2,1,0} convert(%add.945), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention_layernorm/convert_element_type" stack_frame_id=164} + %pow.331 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1481, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention_layernorm/pow" stack_frame_id=167} + %reduce_sum.1408 = f32[1,1]{1,0} reduce(%pow.331, %constant.123), dimensions={2}, to_apply=%region_206.214, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention_layernorm/reduce_sum" stack_frame_id=73} + %broadcast_in_dim.1307 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1408), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention_layernorm/broadcast_in_dim" stack_frame_id=73} + %div.833 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1307, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention_layernorm/div" stack_frame_id=73} + %add.949 = f32[1,1,1]{2,1,0} add(%div.833, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention_layernorm/add" stack_frame_id=168} + %rsqrt.166 = f32[1,1,1]{2,1,0} rsqrt(%add.949), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention_layernorm/rsqrt" stack_frame_id=171} + %mul.2680 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.166), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2681 = f32[1,1]{1,0} reshape(%mul.2680), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2682 = f32[1,1,4096]{2,1,0} broadcast(%mul.2681), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2683 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1481, %mul.2682), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention_layernorm/mul" stack_frame_id=172} + %convert_element_type.1482 = f32[4096]{0} convert(%arg_tuple.5#179), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention_layernorm/convert_element_type" stack_frame_id=173} + %broadcast_in_dim.1308 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1482), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention_layernorm/broadcast_in_dim" stack_frame_id=173} + %mul.2684 = f32[1,1,4096]{2,1,0} multiply(%mul.2683, %broadcast_in_dim.1308), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention_layernorm/mul" stack_frame_id=173} + %convert_element_type.1483 = bf16[1,1,4096]{2,1,0} convert(%mul.2684), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention_layernorm/convert_element_type" stack_frame_id=177} + %dot_general.909 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1483, %arg_tuple.5#182), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=265} + %lt.545 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/lt" stack_frame_id=329} + %add.955 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/add" stack_frame_id=329} + %select_n.130 = s32[] select(%lt.545, %add.955, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/select_n" stack_frame_id=329} + %dynamic_update_slice.104 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.159, %dot_general.909, %constant.127, %select_n.130, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/dynamic_update_slice" stack_frame_id=329} + %broadcast_in_dim.1310 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.104), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/broadcast_in_dim" stack_frame_id=338} + %reshape.612 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1310), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/reshape" stack_frame_id=338} + %convert_element_type.1478 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/convert_element_type" stack_frame_id=139} + %add.946 = f32[1,1]{1,0} reshape(%convert_element_type.1478), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_19/add"} + %ge.429 = f32[1,1]{1,0} broadcast(%add.946), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/ge" stack_frame_id=143} + %ge.430 = f32[1]{0} reshape(%ge.429), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/ge" stack_frame_id=143} + %ge.431 = f32[1,41]{1,0} broadcast(%ge.430), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/ge" stack_frame_id=143} + %iota.458 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/iota" stack_frame_id=142} + %broadcast_in_dim.1303 = f32[1,41]{1,0} reshape(%iota.458), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/broadcast_in_dim" stack_frame_id=143} + %ge.432 = pred[1,41]{1,0} compare(%ge.431, %broadcast_in_dim.1303), direction=GE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/ge" stack_frame_id=143} + %broadcast_in_dim.1304 = pred[1,1,41]{2,1,0} reshape(%ge.432), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/broadcast_in_dim" stack_frame_id=146} + %convert_element_type.1480 = s32[1,1,41]{2,1,0} convert(%broadcast_in_dim.1304), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/convert_element_type" stack_frame_id=160} + %add.947 = s32[1,1]{1,0} reshape(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/add" stack_frame_id=147} + %lt.540 = s32[1,1]{1,0} broadcast(%add.947), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/lt" stack_frame_id=152} + %lt.541 = s32[1]{0} reshape(%lt.540), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/lt" stack_frame_id=152} + %lt.542 = s32[1,41]{1,0} broadcast(%lt.541), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/lt" stack_frame_id=152} + %iota.459 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/iota" stack_frame_id=150} + %broadcast_in_dim.1305 = s32[1,41]{1,0} reshape(%iota.459), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/broadcast_in_dim" stack_frame_id=148} + %add.948 = s32[1,41]{1,0} add(%broadcast_in_dim.1305, %broadcast.43), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/add" stack_frame_id=151} + %lt.543 = pred[1,41]{1,0} compare(%lt.542, %add.948), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/lt" stack_frame_id=152} + %convert_element_type.1479 = s32[1,41]{1,0} convert(%lt.543), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/convert_element_type" stack_frame_id=156} + %broadcast_in_dim.1306 = s32[1,1,41]{2,1,0} reshape(%convert_element_type.1479), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/broadcast_in_dim" stack_frame_id=160} + %min.82 = s32[1,1,41]{2,1,0} minimum(%convert_element_type.1480, %broadcast_in_dim.1306), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/min" stack_frame_id=160} + %broadcast_in_dim.1311 = s32[1,1,1,41]{3,2,1,0} reshape(%min.82), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/broadcast_in_dim" stack_frame_id=341} + %convert_element_type.1490 = pred[1,1,1,41]{3,2,1,0} compare(%broadcast_in_dim.1311, %broadcast.37), direction=NE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/convert_element_type" stack_frame_id=345} + %broadcast_in_dim.1312 = pred[1,1,1,1,41]{4,3,2,1,0} reshape(%convert_element_type.1490), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/broadcast_in_dim" stack_frame_id=83} + %and.254 = pred[1,1,1,1,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1312), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/vmap()/and" stack_frame_id=83} + %and.255 = pred[1,1,1,41]{3,2,1,0} reshape(%and.254), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/vmap()/and" stack_frame_id=83} + %and.256 = pred[1,1,32,1,41]{4,3,2,1,0} broadcast(%and.255), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/vmap()/and" stack_frame_id=83} + %dot_general.906 = bf16[1,1,32,128]{3,2,1,0} dot(%convert_element_type.1483, %arg_tuple.5#180), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=186} + %convert_element_type.1484 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/convert_element_type" stack_frame_id=208} + %add.950 = f32[1,1]{1,0} reshape(%convert_element_type.1484), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/add"} + %iota.460 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/iota" stack_frame_id=197} + %mul.2685 = f32[64]{0} multiply(%iota.460, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=197} + %div.834 = f32[64]{0} divide(%mul.2685, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/div" stack_frame_id=198} + %neg.338 = f32[64]{0} negate(%div.834), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/neg" stack_frame_id=199} + %pow.332 = f32[64]{0} power(%broadcast.39, %neg.338), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/pow" stack_frame_id=202} + %div.835 = f32[64]{0} divide(%pow.332, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/div" stack_frame_id=203} + %dot_general.907 = f32[1,1,64]{2,1,0} dot(%add.950, %div.835), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/bi,j->bij/dot_general" stack_frame_id=211} + %stack.1305 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.907), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/stack" stack_frame_id=214} + %stack.1306 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.907), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/stack" stack_frame_id=214} + %stack.1307 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1305, %stack.1306), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/stack" stack_frame_id=214} + %reshape.607 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1307), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/reshape"} + %cos.165 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.607), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/cos" stack_frame_id=218} + %convert_element_type.1485 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.165), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/convert_element_type" stack_frame_id=222} + %mul.2686 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1485), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=242} + %mul.2687 = bf16[1,1,128]{2,1,0} reshape(%mul.2686), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=242} + %mul.2688 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2687), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=242} + %mul.2689 = bf16[1,1,32,128]{3,2,1,0} multiply(%dot_general.906, %mul.2688), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=242} + %split.331 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.906), slice={[0:1], [0:1], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/split" stack_frame_id=234} + %neg.339 = bf16[1,1,32,64]{3,2,1,0} negate(%split.331), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/neg" stack_frame_id=235} + %stack.1308 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%neg.339), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/stack" stack_frame_id=238} + %split.330 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.906), slice={[0:1], [0:1], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/split" stack_frame_id=234} + %stack.1309 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%split.330), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/stack" stack_frame_id=238} + %stack.1310 = bf16[1,1,32,2,64]{4,3,2,1,0} concatenate(%stack.1308, %stack.1309), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/stack" stack_frame_id=238} + %reshape.608 = bf16[1,1,32,128]{3,2,1,0} reshape(%stack.1310), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/reshape" stack_frame_id=241} + %sin.165 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.607), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/sin" stack_frame_id=226} + %convert_element_type.1486 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.165), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/convert_element_type" stack_frame_id=230} + %mul.2690 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1486), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=243} + %mul.2691 = bf16[1,1,128]{2,1,0} reshape(%mul.2690), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=243} + %mul.2692 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2691), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=243} + %mul.2693 = bf16[1,1,32,128]{3,2,1,0} multiply(%reshape.608, %mul.2692), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=243} + %add.951 = bf16[1,1,32,128]{3,2,1,0} add(%mul.2689, %mul.2693), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/add" stack_frame_id=244} + %reshape.613 = bf16[1,1,32,1,128]{4,3,2,1,0} reshape(%add.951), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/reshape" stack_frame_id=83} + %slice.154 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.157), slice={[0:1], [0:1], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/slice" stack_frame_id=245} + %squeeze.158 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.154), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/squeeze" stack_frame_id=245} + %dot_general.908 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1483, %arg_tuple.5#181), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=256} + %convert_element_type.1487 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/convert_element_type" stack_frame_id=287} + %add.952 = f32[1,1]{1,0} reshape(%convert_element_type.1487), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/add"} + %iota.461 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/iota" stack_frame_id=276} + %mul.2694 = f32[64]{0} multiply(%iota.461, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=276} + %div.836 = f32[64]{0} divide(%mul.2694, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/div" stack_frame_id=277} + %neg.340 = f32[64]{0} negate(%div.836), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/neg" stack_frame_id=278} + %pow.333 = f32[64]{0} power(%broadcast.39, %neg.340), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/pow" stack_frame_id=281} + %div.837 = f32[64]{0} divide(%pow.333, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/div" stack_frame_id=282} + %dot_general.910 = f32[1,1,64]{2,1,0} dot(%add.952, %div.837), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/bi,j->bij/dot_general" stack_frame_id=290} + %stack.1311 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.910), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/stack" stack_frame_id=293} + %stack.1312 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.910), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/stack" stack_frame_id=293} + %stack.1313 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1311, %stack.1312), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/stack" stack_frame_id=293} + %reshape.609 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1313), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/reshape"} + %cos.166 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.609), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/cos" stack_frame_id=297} + %convert_element_type.1488 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.166), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/convert_element_type" stack_frame_id=301} + %mul.2695 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1488), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=321} + %mul.2696 = bf16[1,1,128]{2,1,0} reshape(%mul.2695), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=321} + %mul.2697 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2696), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=321} + %mul.2698 = bf16[1,1,8,128]{3,2,1,0} multiply(%dot_general.908, %mul.2697), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=321} + %split.333 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.908), slice={[0:1], [0:1], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/split" stack_frame_id=313} + %neg.341 = bf16[1,1,8,64]{3,2,1,0} negate(%split.333), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/neg" stack_frame_id=314} + %stack.1314 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%neg.341), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/stack" stack_frame_id=317} + %split.332 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.908), slice={[0:1], [0:1], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/split" stack_frame_id=313} + %stack.1315 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%split.332), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/stack" stack_frame_id=317} + %stack.1316 = bf16[1,1,8,2,64]{4,3,2,1,0} concatenate(%stack.1314, %stack.1315), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/stack" stack_frame_id=317} + %reshape.610 = bf16[1,1,8,128]{3,2,1,0} reshape(%stack.1316), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/reshape" stack_frame_id=320} + %sin.166 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.609), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/sin" stack_frame_id=305} + %convert_element_type.1489 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.166), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/convert_element_type" stack_frame_id=309} + %mul.2699 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1489), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=322} + %mul.2700 = bf16[1,1,128]{2,1,0} reshape(%mul.2699), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=322} + %mul.2701 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2700), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=322} + %mul.2702 = bf16[1,1,8,128]{3,2,1,0} multiply(%reshape.610, %mul.2701), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=322} + %add.953 = bf16[1,1,8,128]{3,2,1,0} add(%mul.2698, %mul.2702), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/rotary_embedding_19/add" stack_frame_id=323} + %lt.544 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/lt" stack_frame_id=326} + %add.954 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/add" stack_frame_id=326} + %select_n.129 = s32[] select(%lt.544, %add.954, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/select_n" stack_frame_id=326} + %dynamic_update_slice.103 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.158, %add.953, %constant.127, %select_n.129, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/dynamic_update_slice" stack_frame_id=326} + %broadcast_in_dim.1309 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.103), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/broadcast_in_dim" stack_frame_id=335} + %reshape.611 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1309), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/reshape" stack_frame_id=335} + %dot_general.911 = f32[1,32,1,1,41]{4,3,2,1,0} dot(%reshape.613, %reshape.611), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=83} + %mul.2703 = f32[1,32,1,1,41]{4,3,2,1,0} multiply(%dot_general.911, %broadcast.36), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/vmap()/mul" stack_frame_id=83} + %vmap_jit__where__.82 = f32[1,1,32,1,41]{4,3,2,1,0} call(%and.256, %mul.2703, %constant.114), to_apply=%_where_3.137, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/vmap(jit(_where))" stack_frame_id=83} + %reduce_max.467 = f32[1,1,32,1]{3,2,1,0} reduce(%vmap_jit__where__.82, %constant.113), dimensions={4}, to_apply=%region_207.215, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/vmap()/reduce_max" stack_frame_id=83} + %max.91 = f32[1,1,32,1]{3,2,1,0} maximum(%reduce_max.467, %broadcast.35), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/vmap()/max" stack_frame_id=83} + %broadcast_in_dim.1313 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%max.91), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %sub.347 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1313), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/vmap()/sub" stack_frame_id=83} + %sub.348 = f32[1,1,32,1]{3,2,1,0} reshape(%sub.347), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/vmap()/sub" stack_frame_id=83} + %sub.349 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%sub.348), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/vmap()/sub" stack_frame_id=83} + %sub.350 = f32[1,1,32,1,41]{4,3,2,1,0} subtract(%vmap_jit__where__.82, %sub.349), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/vmap()/sub" stack_frame_id=83} + %exp.87 = f32[1,1,32,1,41]{4,3,2,1,0} exponential(%sub.350), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/vmap()/exp" stack_frame_id=83} + %reduce_sum.1409 = f32[1,1,32,1]{3,2,1,0} reduce(%exp.87, %constant.123), dimensions={4}, to_apply=%region_208.216, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/vmap()/reduce_sum" stack_frame_id=83} + %broadcast_in_dim.1314 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%reduce_sum.1409), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %div.838 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1314), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/vmap()/div" stack_frame_id=83} + %div.839 = f32[1,1,32,1]{3,2,1,0} reshape(%div.838), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/vmap()/div" stack_frame_id=83} + %div.840 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%div.839), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/vmap()/div" stack_frame_id=83} + %div.841 = f32[1,1,32,1,41]{4,3,2,1,0} divide(%exp.87, %div.840), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/vmap()/div" stack_frame_id=83} + %convert_element_type.1491 = bf16[1,1,32,1,41]{4,3,2,1,0} convert(%div.841), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/vmap()/convert_element_type" stack_frame_id=83} + %dot_general.912 = bf16[1,32,128,1,1]{4,3,2,1,0} dot(%reshape.612, %convert_element_type.1491), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=83} + %transpose.86 = bf16[1,1,32,128]{3,2,1,0} reshape(%dot_general.912), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/reshape;jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/vmap()/transpose;jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %dot_general.913 = bf16[1,1,4096]{2,1,0} dot(%transpose.86, %arg_tuple.5#183), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=354} + %add.956 = bf16[1,1,4096]{2,1,0} add(%dot_general.913, %add.945), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/add" stack_frame_id=355} + %convert_element_type.1492 = f32[1,1,4096]{2,1,0} convert(%add.956), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/feedforward_layernorm/convert_element_type" stack_frame_id=359} + %pow.334 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1492, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/feedforward_layernorm/pow" stack_frame_id=362} + %reduce_sum.1410 = f32[1,1]{1,0} reduce(%pow.334, %constant.123), dimensions={2}, to_apply=%region_209.217, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/feedforward_layernorm/reduce_sum" stack_frame_id=92} + %broadcast_in_dim.1315 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1410), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/feedforward_layernorm/broadcast_in_dim" stack_frame_id=92} + %div.842 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1315, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/feedforward_layernorm/div" stack_frame_id=92} + %add.957 = f32[1,1,1]{2,1,0} add(%div.842, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/feedforward_layernorm/add" stack_frame_id=363} + %rsqrt.167 = f32[1,1,1]{2,1,0} rsqrt(%add.957), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/feedforward_layernorm/rsqrt" stack_frame_id=366} + %mul.2704 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.167), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2705 = f32[1,1]{1,0} reshape(%mul.2704), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2706 = f32[1,1,4096]{2,1,0} broadcast(%mul.2705), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2707 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1492, %mul.2706), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/feedforward_layernorm/mul" stack_frame_id=367} + %convert_element_type.1493 = f32[4096]{0} convert(%arg_tuple.5#184), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/feedforward_layernorm/convert_element_type" stack_frame_id=368} + %broadcast_in_dim.1316 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1493), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/feedforward_layernorm/broadcast_in_dim" stack_frame_id=368} + %mul.2708 = f32[1,1,4096]{2,1,0} multiply(%mul.2707, %broadcast_in_dim.1316), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/feedforward_layernorm/mul" stack_frame_id=368} + %convert_element_type.1494 = bf16[1,1,4096]{2,1,0} convert(%mul.2708), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/feedforward_layernorm/convert_element_type" stack_frame_id=372} + %dot_general.915 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1494, %arg_tuple.5#186), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/feedforward_intermediate_dense/dot_general" stack_frame_id=401} + %dot_general.914 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1494, %arg_tuple.5#185), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/feedforward_gate_dense/dot_general" stack_frame_id=381} + %convert_element_type.1495 = f32[1,1,14336]{2,1,0} convert(%dot_general.914), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/convert_element_type" stack_frame_id=385} + %jit_silu_.82 = f32[1,1,14336]{2,1,0} call(%convert_element_type.1495), to_apply=%silu_4.141, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/jit(silu)" stack_frame_id=388} + %convert_element_type.1496 = bf16[1,1,14336]{2,1,0} convert(%jit_silu_.82), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/convert_element_type" stack_frame_id=392} + %mul.2709 = bf16[1,1,14336]{2,1,0} multiply(%dot_general.915, %convert_element_type.1496), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/mul" stack_frame_id=404} + %dot_general.916 = bf16[1,1,4096]{2,1,0} dot(%mul.2709, %arg_tuple.5#187), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/feedforward_output_dense/dot_general" stack_frame_id=413} + %add.958 = bf16[1,1,4096]{2,1,0} add(%dot_general.916, %add.956), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/add" stack_frame_id=414} + %convert_element_type.1500 = f32[1,1,4096]{2,1,0} convert(%add.958), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention_layernorm/convert_element_type" stack_frame_id=164} + %pow.335 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1500, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention_layernorm/pow" stack_frame_id=167} + %reduce_sum.1411 = f32[1,1]{1,0} reduce(%pow.335, %constant.123), dimensions={2}, to_apply=%region_210.218, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention_layernorm/reduce_sum" stack_frame_id=73} + %broadcast_in_dim.1321 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1411), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention_layernorm/broadcast_in_dim" stack_frame_id=73} + %div.843 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1321, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention_layernorm/div" stack_frame_id=73} + %add.962 = f32[1,1,1]{2,1,0} add(%div.843, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention_layernorm/add" stack_frame_id=168} + %rsqrt.168 = f32[1,1,1]{2,1,0} rsqrt(%add.962), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention_layernorm/rsqrt" stack_frame_id=171} + %mul.2710 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.168), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2711 = f32[1,1]{1,0} reshape(%mul.2710), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2712 = f32[1,1,4096]{2,1,0} broadcast(%mul.2711), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2713 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1500, %mul.2712), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention_layernorm/mul" stack_frame_id=172} + %convert_element_type.1501 = f32[4096]{0} convert(%arg_tuple.5#188), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention_layernorm/convert_element_type" stack_frame_id=173} + %broadcast_in_dim.1322 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1501), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention_layernorm/broadcast_in_dim" stack_frame_id=173} + %mul.2714 = f32[1,1,4096]{2,1,0} multiply(%mul.2713, %broadcast_in_dim.1322), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention_layernorm/mul" stack_frame_id=173} + %convert_element_type.1502 = bf16[1,1,4096]{2,1,0} convert(%mul.2714), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention_layernorm/convert_element_type" stack_frame_id=177} + %dot_general.920 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1502, %arg_tuple.5#191), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=265} + %lt.551 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/lt" stack_frame_id=329} + %add.968 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/add" stack_frame_id=329} + %select_n.132 = s32[] select(%lt.551, %add.968, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/select_n" stack_frame_id=329} + %dynamic_update_slice.106 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.162, %dot_general.920, %constant.127, %select_n.132, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/dynamic_update_slice" stack_frame_id=329} + %broadcast_in_dim.1324 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.106), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/broadcast_in_dim" stack_frame_id=338} + %reshape.619 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1324), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/reshape" stack_frame_id=338} + %convert_element_type.1497 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/convert_element_type" stack_frame_id=139} + %add.959 = f32[1,1]{1,0} reshape(%convert_element_type.1497), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_20/add"} + %ge.433 = f32[1,1]{1,0} broadcast(%add.959), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/ge" stack_frame_id=143} + %ge.434 = f32[1]{0} reshape(%ge.433), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/ge" stack_frame_id=143} + %ge.435 = f32[1,41]{1,0} broadcast(%ge.434), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/ge" stack_frame_id=143} + %iota.462 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/iota" stack_frame_id=142} + %broadcast_in_dim.1317 = f32[1,41]{1,0} reshape(%iota.462), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/broadcast_in_dim" stack_frame_id=143} + %ge.436 = pred[1,41]{1,0} compare(%ge.435, %broadcast_in_dim.1317), direction=GE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/ge" stack_frame_id=143} + %broadcast_in_dim.1318 = pred[1,1,41]{2,1,0} reshape(%ge.436), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/broadcast_in_dim" stack_frame_id=146} + %convert_element_type.1499 = s32[1,1,41]{2,1,0} convert(%broadcast_in_dim.1318), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/convert_element_type" stack_frame_id=160} + %add.960 = s32[1,1]{1,0} reshape(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/add" stack_frame_id=147} + %lt.546 = s32[1,1]{1,0} broadcast(%add.960), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/lt" stack_frame_id=152} + %lt.547 = s32[1]{0} reshape(%lt.546), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/lt" stack_frame_id=152} + %lt.548 = s32[1,41]{1,0} broadcast(%lt.547), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/lt" stack_frame_id=152} + %iota.463 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/iota" stack_frame_id=150} + %broadcast_in_dim.1319 = s32[1,41]{1,0} reshape(%iota.463), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/broadcast_in_dim" stack_frame_id=148} + %add.961 = s32[1,41]{1,0} add(%broadcast_in_dim.1319, %broadcast.43), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/add" stack_frame_id=151} + %lt.549 = pred[1,41]{1,0} compare(%lt.548, %add.961), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/lt" stack_frame_id=152} + %convert_element_type.1498 = s32[1,41]{1,0} convert(%lt.549), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/convert_element_type" stack_frame_id=156} + %broadcast_in_dim.1320 = s32[1,1,41]{2,1,0} reshape(%convert_element_type.1498), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/broadcast_in_dim" stack_frame_id=160} + %min.83 = s32[1,1,41]{2,1,0} minimum(%convert_element_type.1499, %broadcast_in_dim.1320), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/min" stack_frame_id=160} + %broadcast_in_dim.1325 = s32[1,1,1,41]{3,2,1,0} reshape(%min.83), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/broadcast_in_dim" stack_frame_id=341} + %convert_element_type.1509 = pred[1,1,1,41]{3,2,1,0} compare(%broadcast_in_dim.1325, %broadcast.37), direction=NE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/convert_element_type" stack_frame_id=345} + %broadcast_in_dim.1326 = pred[1,1,1,1,41]{4,3,2,1,0} reshape(%convert_element_type.1509), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/broadcast_in_dim" stack_frame_id=83} + %and.257 = pred[1,1,1,1,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1326), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/vmap()/and" stack_frame_id=83} + %and.258 = pred[1,1,1,41]{3,2,1,0} reshape(%and.257), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/vmap()/and" stack_frame_id=83} + %and.259 = pred[1,1,32,1,41]{4,3,2,1,0} broadcast(%and.258), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/vmap()/and" stack_frame_id=83} + %dot_general.917 = bf16[1,1,32,128]{3,2,1,0} dot(%convert_element_type.1502, %arg_tuple.5#189), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=186} + %convert_element_type.1503 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/convert_element_type" stack_frame_id=208} + %add.963 = f32[1,1]{1,0} reshape(%convert_element_type.1503), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/add"} + %iota.464 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/iota" stack_frame_id=197} + %mul.2715 = f32[64]{0} multiply(%iota.464, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=197} + %div.844 = f32[64]{0} divide(%mul.2715, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/div" stack_frame_id=198} + %neg.342 = f32[64]{0} negate(%div.844), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/neg" stack_frame_id=199} + %pow.336 = f32[64]{0} power(%broadcast.39, %neg.342), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/pow" stack_frame_id=202} + %div.845 = f32[64]{0} divide(%pow.336, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/div" stack_frame_id=203} + %dot_general.918 = f32[1,1,64]{2,1,0} dot(%add.963, %div.845), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/bi,j->bij/dot_general" stack_frame_id=211} + %stack.1320 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.918), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/stack" stack_frame_id=214} + %stack.1321 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.918), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/stack" stack_frame_id=214} + %stack.1322 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1320, %stack.1321), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/stack" stack_frame_id=214} + %reshape.614 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1322), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/reshape"} + %cos.167 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.614), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/cos" stack_frame_id=218} + %convert_element_type.1504 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.167), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/convert_element_type" stack_frame_id=222} + %mul.2716 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1504), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=242} + %mul.2717 = bf16[1,1,128]{2,1,0} reshape(%mul.2716), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=242} + %mul.2718 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2717), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=242} + %mul.2719 = bf16[1,1,32,128]{3,2,1,0} multiply(%dot_general.917, %mul.2718), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=242} + %split.335 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.917), slice={[0:1], [0:1], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/split" stack_frame_id=234} + %neg.343 = bf16[1,1,32,64]{3,2,1,0} negate(%split.335), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/neg" stack_frame_id=235} + %stack.1323 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%neg.343), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/stack" stack_frame_id=238} + %split.334 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.917), slice={[0:1], [0:1], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/split" stack_frame_id=234} + %stack.1324 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%split.334), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/stack" stack_frame_id=238} + %stack.1325 = bf16[1,1,32,2,64]{4,3,2,1,0} concatenate(%stack.1323, %stack.1324), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/stack" stack_frame_id=238} + %reshape.615 = bf16[1,1,32,128]{3,2,1,0} reshape(%stack.1325), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/reshape" stack_frame_id=241} + %sin.167 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.614), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/sin" stack_frame_id=226} + %convert_element_type.1505 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.167), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/convert_element_type" stack_frame_id=230} + %mul.2720 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1505), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=243} + %mul.2721 = bf16[1,1,128]{2,1,0} reshape(%mul.2720), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=243} + %mul.2722 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2721), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=243} + %mul.2723 = bf16[1,1,32,128]{3,2,1,0} multiply(%reshape.615, %mul.2722), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=243} + %add.964 = bf16[1,1,32,128]{3,2,1,0} add(%mul.2719, %mul.2723), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/add" stack_frame_id=244} + %reshape.620 = bf16[1,1,32,1,128]{4,3,2,1,0} reshape(%add.964), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/reshape" stack_frame_id=83} + %slice.157 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.160), slice={[0:1], [0:1], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/slice" stack_frame_id=245} + %squeeze.161 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.157), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/squeeze" stack_frame_id=245} + %dot_general.919 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1502, %arg_tuple.5#190), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=256} + %convert_element_type.1506 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/convert_element_type" stack_frame_id=287} + %add.965 = f32[1,1]{1,0} reshape(%convert_element_type.1506), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/add"} + %iota.465 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/iota" stack_frame_id=276} + %mul.2724 = f32[64]{0} multiply(%iota.465, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=276} + %div.846 = f32[64]{0} divide(%mul.2724, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/div" stack_frame_id=277} + %neg.344 = f32[64]{0} negate(%div.846), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/neg" stack_frame_id=278} + %pow.337 = f32[64]{0} power(%broadcast.39, %neg.344), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/pow" stack_frame_id=281} + %div.847 = f32[64]{0} divide(%pow.337, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/div" stack_frame_id=282} + %dot_general.921 = f32[1,1,64]{2,1,0} dot(%add.965, %div.847), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/bi,j->bij/dot_general" stack_frame_id=290} + %stack.1326 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.921), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/stack" stack_frame_id=293} + %stack.1327 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.921), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/stack" stack_frame_id=293} + %stack.1328 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1326, %stack.1327), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/stack" stack_frame_id=293} + %reshape.616 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1328), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/reshape"} + %cos.168 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.616), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/cos" stack_frame_id=297} + %convert_element_type.1507 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.168), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/convert_element_type" stack_frame_id=301} + %mul.2725 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1507), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=321} + %mul.2726 = bf16[1,1,128]{2,1,0} reshape(%mul.2725), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=321} + %mul.2727 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2726), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=321} + %mul.2728 = bf16[1,1,8,128]{3,2,1,0} multiply(%dot_general.919, %mul.2727), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=321} + %split.337 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.919), slice={[0:1], [0:1], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/split" stack_frame_id=313} + %neg.345 = bf16[1,1,8,64]{3,2,1,0} negate(%split.337), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/neg" stack_frame_id=314} + %stack.1329 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%neg.345), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/stack" stack_frame_id=317} + %split.336 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.919), slice={[0:1], [0:1], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/split" stack_frame_id=313} + %stack.1330 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%split.336), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/stack" stack_frame_id=317} + %stack.1331 = bf16[1,1,8,2,64]{4,3,2,1,0} concatenate(%stack.1329, %stack.1330), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/stack" stack_frame_id=317} + %reshape.617 = bf16[1,1,8,128]{3,2,1,0} reshape(%stack.1331), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/reshape" stack_frame_id=320} + %sin.168 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.616), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/sin" stack_frame_id=305} + %convert_element_type.1508 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.168), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/convert_element_type" stack_frame_id=309} + %mul.2729 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1508), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=322} + %mul.2730 = bf16[1,1,128]{2,1,0} reshape(%mul.2729), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=322} + %mul.2731 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2730), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=322} + %mul.2732 = bf16[1,1,8,128]{3,2,1,0} multiply(%reshape.617, %mul.2731), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=322} + %add.966 = bf16[1,1,8,128]{3,2,1,0} add(%mul.2728, %mul.2732), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/rotary_embedding_20/add" stack_frame_id=323} + %lt.550 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/lt" stack_frame_id=326} + %add.967 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/add" stack_frame_id=326} + %select_n.131 = s32[] select(%lt.550, %add.967, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/select_n" stack_frame_id=326} + %dynamic_update_slice.105 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.161, %add.966, %constant.127, %select_n.131, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/dynamic_update_slice" stack_frame_id=326} + %broadcast_in_dim.1323 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.105), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/broadcast_in_dim" stack_frame_id=335} + %reshape.618 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1323), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/reshape" stack_frame_id=335} + %dot_general.922 = f32[1,32,1,1,41]{4,3,2,1,0} dot(%reshape.620, %reshape.618), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=83} + %mul.2733 = f32[1,32,1,1,41]{4,3,2,1,0} multiply(%dot_general.922, %broadcast.36), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/vmap()/mul" stack_frame_id=83} + %vmap_jit__where__.83 = f32[1,1,32,1,41]{4,3,2,1,0} call(%and.259, %mul.2733, %constant.114), to_apply=%_where_3.137, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/vmap(jit(_where))" stack_frame_id=83} + %reduce_max.468 = f32[1,1,32,1]{3,2,1,0} reduce(%vmap_jit__where__.83, %constant.113), dimensions={4}, to_apply=%region_211.219, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/vmap()/reduce_max" stack_frame_id=83} + %max.92 = f32[1,1,32,1]{3,2,1,0} maximum(%reduce_max.468, %broadcast.35), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/vmap()/max" stack_frame_id=83} + %broadcast_in_dim.1327 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%max.92), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %sub.351 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1327), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/vmap()/sub" stack_frame_id=83} + %sub.352 = f32[1,1,32,1]{3,2,1,0} reshape(%sub.351), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/vmap()/sub" stack_frame_id=83} + %sub.353 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%sub.352), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/vmap()/sub" stack_frame_id=83} + %sub.354 = f32[1,1,32,1,41]{4,3,2,1,0} subtract(%vmap_jit__where__.83, %sub.353), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/vmap()/sub" stack_frame_id=83} + %exp.88 = f32[1,1,32,1,41]{4,3,2,1,0} exponential(%sub.354), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/vmap()/exp" stack_frame_id=83} + %reduce_sum.1412 = f32[1,1,32,1]{3,2,1,0} reduce(%exp.88, %constant.123), dimensions={4}, to_apply=%region_212.220, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/vmap()/reduce_sum" stack_frame_id=83} + %broadcast_in_dim.1328 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%reduce_sum.1412), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %div.848 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1328), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/vmap()/div" stack_frame_id=83} + %div.849 = f32[1,1,32,1]{3,2,1,0} reshape(%div.848), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/vmap()/div" stack_frame_id=83} + %div.850 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%div.849), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/vmap()/div" stack_frame_id=83} + %div.851 = f32[1,1,32,1,41]{4,3,2,1,0} divide(%exp.88, %div.850), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/vmap()/div" stack_frame_id=83} + %convert_element_type.1510 = bf16[1,1,32,1,41]{4,3,2,1,0} convert(%div.851), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/vmap()/convert_element_type" stack_frame_id=83} + %dot_general.923 = bf16[1,32,128,1,1]{4,3,2,1,0} dot(%reshape.619, %convert_element_type.1510), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=83} + %transpose.87 = bf16[1,1,32,128]{3,2,1,0} reshape(%dot_general.923), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/reshape;jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/vmap()/transpose;jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %dot_general.924 = bf16[1,1,4096]{2,1,0} dot(%transpose.87, %arg_tuple.5#192), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=354} + %add.969 = bf16[1,1,4096]{2,1,0} add(%dot_general.924, %add.958), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/add" stack_frame_id=355} + %convert_element_type.1511 = f32[1,1,4096]{2,1,0} convert(%add.969), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/feedforward_layernorm/convert_element_type" stack_frame_id=359} + %pow.338 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1511, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/feedforward_layernorm/pow" stack_frame_id=362} + %reduce_sum.1413 = f32[1,1]{1,0} reduce(%pow.338, %constant.123), dimensions={2}, to_apply=%region_213.221, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/feedforward_layernorm/reduce_sum" stack_frame_id=92} + %broadcast_in_dim.1329 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1413), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/feedforward_layernorm/broadcast_in_dim" stack_frame_id=92} + %div.852 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1329, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/feedforward_layernorm/div" stack_frame_id=92} + %add.970 = f32[1,1,1]{2,1,0} add(%div.852, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/feedforward_layernorm/add" stack_frame_id=363} + %rsqrt.169 = f32[1,1,1]{2,1,0} rsqrt(%add.970), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/feedforward_layernorm/rsqrt" stack_frame_id=366} + %mul.2734 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.169), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2735 = f32[1,1]{1,0} reshape(%mul.2734), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2736 = f32[1,1,4096]{2,1,0} broadcast(%mul.2735), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2737 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1511, %mul.2736), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/feedforward_layernorm/mul" stack_frame_id=367} + %convert_element_type.1512 = f32[4096]{0} convert(%arg_tuple.5#193), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/feedforward_layernorm/convert_element_type" stack_frame_id=368} + %broadcast_in_dim.1330 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1512), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/feedforward_layernorm/broadcast_in_dim" stack_frame_id=368} + %mul.2738 = f32[1,1,4096]{2,1,0} multiply(%mul.2737, %broadcast_in_dim.1330), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/feedforward_layernorm/mul" stack_frame_id=368} + %convert_element_type.1513 = bf16[1,1,4096]{2,1,0} convert(%mul.2738), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/feedforward_layernorm/convert_element_type" stack_frame_id=372} + %dot_general.926 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1513, %arg_tuple.5#195), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/feedforward_intermediate_dense/dot_general" stack_frame_id=401} + %dot_general.925 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1513, %arg_tuple.5#194), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/feedforward_gate_dense/dot_general" stack_frame_id=381} + %convert_element_type.1514 = f32[1,1,14336]{2,1,0} convert(%dot_general.925), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/convert_element_type" stack_frame_id=385} + %jit_silu_.83 = f32[1,1,14336]{2,1,0} call(%convert_element_type.1514), to_apply=%silu_4.141, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/jit(silu)" stack_frame_id=388} + %convert_element_type.1515 = bf16[1,1,14336]{2,1,0} convert(%jit_silu_.83), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/convert_element_type" stack_frame_id=392} + %mul.2739 = bf16[1,1,14336]{2,1,0} multiply(%dot_general.926, %convert_element_type.1515), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/mul" stack_frame_id=404} + %dot_general.927 = bf16[1,1,4096]{2,1,0} dot(%mul.2739, %arg_tuple.5#196), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/feedforward_output_dense/dot_general" stack_frame_id=413} + %add.971 = bf16[1,1,4096]{2,1,0} add(%dot_general.927, %add.969), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/add" stack_frame_id=414} + %convert_element_type.1519 = f32[1,1,4096]{2,1,0} convert(%add.971), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention_layernorm/convert_element_type" stack_frame_id=164} + %pow.339 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1519, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention_layernorm/pow" stack_frame_id=167} + %reduce_sum.1414 = f32[1,1]{1,0} reduce(%pow.339, %constant.123), dimensions={2}, to_apply=%region_214.222, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention_layernorm/reduce_sum" stack_frame_id=73} + %broadcast_in_dim.1335 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1414), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention_layernorm/broadcast_in_dim" stack_frame_id=73} + %div.853 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1335, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention_layernorm/div" stack_frame_id=73} + %add.975 = f32[1,1,1]{2,1,0} add(%div.853, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention_layernorm/add" stack_frame_id=168} + %rsqrt.170 = f32[1,1,1]{2,1,0} rsqrt(%add.975), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention_layernorm/rsqrt" stack_frame_id=171} + %mul.2740 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.170), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2741 = f32[1,1]{1,0} reshape(%mul.2740), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2742 = f32[1,1,4096]{2,1,0} broadcast(%mul.2741), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2743 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1519, %mul.2742), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention_layernorm/mul" stack_frame_id=172} + %convert_element_type.1520 = f32[4096]{0} convert(%arg_tuple.5#197), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention_layernorm/convert_element_type" stack_frame_id=173} + %broadcast_in_dim.1336 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1520), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention_layernorm/broadcast_in_dim" stack_frame_id=173} + %mul.2744 = f32[1,1,4096]{2,1,0} multiply(%mul.2743, %broadcast_in_dim.1336), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention_layernorm/mul" stack_frame_id=173} + %convert_element_type.1521 = bf16[1,1,4096]{2,1,0} convert(%mul.2744), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention_layernorm/convert_element_type" stack_frame_id=177} + %dot_general.931 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1521, %arg_tuple.5#200), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=265} + %lt.557 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/lt" stack_frame_id=329} + %add.981 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/add" stack_frame_id=329} + %select_n.134 = s32[] select(%lt.557, %add.981, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/select_n" stack_frame_id=329} + %dynamic_update_slice.108 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.165, %dot_general.931, %constant.127, %select_n.134, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/dynamic_update_slice" stack_frame_id=329} + %broadcast_in_dim.1338 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.108), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/broadcast_in_dim" stack_frame_id=338} + %reshape.626 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1338), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/reshape" stack_frame_id=338} + %convert_element_type.1516 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/convert_element_type" stack_frame_id=139} + %add.972 = f32[1,1]{1,0} reshape(%convert_element_type.1516), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_21/add"} + %ge.437 = f32[1,1]{1,0} broadcast(%add.972), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/ge" stack_frame_id=143} + %ge.438 = f32[1]{0} reshape(%ge.437), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/ge" stack_frame_id=143} + %ge.439 = f32[1,41]{1,0} broadcast(%ge.438), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/ge" stack_frame_id=143} + %iota.466 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/iota" stack_frame_id=142} + %broadcast_in_dim.1331 = f32[1,41]{1,0} reshape(%iota.466), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/broadcast_in_dim" stack_frame_id=143} + %ge.440 = pred[1,41]{1,0} compare(%ge.439, %broadcast_in_dim.1331), direction=GE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/ge" stack_frame_id=143} + %broadcast_in_dim.1332 = pred[1,1,41]{2,1,0} reshape(%ge.440), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/broadcast_in_dim" stack_frame_id=146} + %convert_element_type.1518 = s32[1,1,41]{2,1,0} convert(%broadcast_in_dim.1332), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/convert_element_type" stack_frame_id=160} + %add.973 = s32[1,1]{1,0} reshape(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/add" stack_frame_id=147} + %lt.552 = s32[1,1]{1,0} broadcast(%add.973), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/lt" stack_frame_id=152} + %lt.553 = s32[1]{0} reshape(%lt.552), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/lt" stack_frame_id=152} + %lt.554 = s32[1,41]{1,0} broadcast(%lt.553), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/lt" stack_frame_id=152} + %iota.467 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/iota" stack_frame_id=150} + %broadcast_in_dim.1333 = s32[1,41]{1,0} reshape(%iota.467), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/broadcast_in_dim" stack_frame_id=148} + %add.974 = s32[1,41]{1,0} add(%broadcast_in_dim.1333, %broadcast.43), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/add" stack_frame_id=151} + %lt.555 = pred[1,41]{1,0} compare(%lt.554, %add.974), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/lt" stack_frame_id=152} + %convert_element_type.1517 = s32[1,41]{1,0} convert(%lt.555), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/convert_element_type" stack_frame_id=156} + %broadcast_in_dim.1334 = s32[1,1,41]{2,1,0} reshape(%convert_element_type.1517), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/broadcast_in_dim" stack_frame_id=160} + %min.84 = s32[1,1,41]{2,1,0} minimum(%convert_element_type.1518, %broadcast_in_dim.1334), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/min" stack_frame_id=160} + %broadcast_in_dim.1339 = s32[1,1,1,41]{3,2,1,0} reshape(%min.84), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/broadcast_in_dim" stack_frame_id=341} + %convert_element_type.1528 = pred[1,1,1,41]{3,2,1,0} compare(%broadcast_in_dim.1339, %broadcast.37), direction=NE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/convert_element_type" stack_frame_id=345} + %broadcast_in_dim.1340 = pred[1,1,1,1,41]{4,3,2,1,0} reshape(%convert_element_type.1528), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/broadcast_in_dim" stack_frame_id=83} + %and.260 = pred[1,1,1,1,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1340), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/vmap()/and" stack_frame_id=83} + %and.261 = pred[1,1,1,41]{3,2,1,0} reshape(%and.260), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/vmap()/and" stack_frame_id=83} + %and.262 = pred[1,1,32,1,41]{4,3,2,1,0} broadcast(%and.261), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/vmap()/and" stack_frame_id=83} + %dot_general.928 = bf16[1,1,32,128]{3,2,1,0} dot(%convert_element_type.1521, %arg_tuple.5#198), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=186} + %convert_element_type.1522 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/convert_element_type" stack_frame_id=208} + %add.976 = f32[1,1]{1,0} reshape(%convert_element_type.1522), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/add"} + %iota.468 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/iota" stack_frame_id=197} + %mul.2745 = f32[64]{0} multiply(%iota.468, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=197} + %div.854 = f32[64]{0} divide(%mul.2745, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/div" stack_frame_id=198} + %neg.346 = f32[64]{0} negate(%div.854), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/neg" stack_frame_id=199} + %pow.340 = f32[64]{0} power(%broadcast.39, %neg.346), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/pow" stack_frame_id=202} + %div.855 = f32[64]{0} divide(%pow.340, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/div" stack_frame_id=203} + %dot_general.929 = f32[1,1,64]{2,1,0} dot(%add.976, %div.855), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/bi,j->bij/dot_general" stack_frame_id=211} + %stack.1335 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.929), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/stack" stack_frame_id=214} + %stack.1336 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.929), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/stack" stack_frame_id=214} + %stack.1337 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1335, %stack.1336), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/stack" stack_frame_id=214} + %reshape.621 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1337), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/reshape"} + %cos.169 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.621), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/cos" stack_frame_id=218} + %convert_element_type.1523 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.169), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/convert_element_type" stack_frame_id=222} + %mul.2746 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1523), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=242} + %mul.2747 = bf16[1,1,128]{2,1,0} reshape(%mul.2746), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=242} + %mul.2748 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2747), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=242} + %mul.2749 = bf16[1,1,32,128]{3,2,1,0} multiply(%dot_general.928, %mul.2748), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=242} + %split.339 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.928), slice={[0:1], [0:1], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/split" stack_frame_id=234} + %neg.347 = bf16[1,1,32,64]{3,2,1,0} negate(%split.339), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/neg" stack_frame_id=235} + %stack.1338 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%neg.347), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/stack" stack_frame_id=238} + %split.338 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.928), slice={[0:1], [0:1], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/split" stack_frame_id=234} + %stack.1339 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%split.338), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/stack" stack_frame_id=238} + %stack.1340 = bf16[1,1,32,2,64]{4,3,2,1,0} concatenate(%stack.1338, %stack.1339), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/stack" stack_frame_id=238} + %reshape.622 = bf16[1,1,32,128]{3,2,1,0} reshape(%stack.1340), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/reshape" stack_frame_id=241} + %sin.169 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.621), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/sin" stack_frame_id=226} + %convert_element_type.1524 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.169), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/convert_element_type" stack_frame_id=230} + %mul.2750 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1524), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=243} + %mul.2751 = bf16[1,1,128]{2,1,0} reshape(%mul.2750), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=243} + %mul.2752 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2751), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=243} + %mul.2753 = bf16[1,1,32,128]{3,2,1,0} multiply(%reshape.622, %mul.2752), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=243} + %add.977 = bf16[1,1,32,128]{3,2,1,0} add(%mul.2749, %mul.2753), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/add" stack_frame_id=244} + %reshape.627 = bf16[1,1,32,1,128]{4,3,2,1,0} reshape(%add.977), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/reshape" stack_frame_id=83} + %slice.160 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.163), slice={[0:1], [0:1], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/slice" stack_frame_id=245} + %squeeze.164 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.160), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/squeeze" stack_frame_id=245} + %dot_general.930 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1521, %arg_tuple.5#199), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=256} + %convert_element_type.1525 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/convert_element_type" stack_frame_id=287} + %add.978 = f32[1,1]{1,0} reshape(%convert_element_type.1525), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/add"} + %iota.469 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/iota" stack_frame_id=276} + %mul.2754 = f32[64]{0} multiply(%iota.469, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=276} + %div.856 = f32[64]{0} divide(%mul.2754, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/div" stack_frame_id=277} + %neg.348 = f32[64]{0} negate(%div.856), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/neg" stack_frame_id=278} + %pow.341 = f32[64]{0} power(%broadcast.39, %neg.348), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/pow" stack_frame_id=281} + %div.857 = f32[64]{0} divide(%pow.341, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/div" stack_frame_id=282} + %dot_general.932 = f32[1,1,64]{2,1,0} dot(%add.978, %div.857), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/bi,j->bij/dot_general" stack_frame_id=290} + %stack.1341 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.932), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/stack" stack_frame_id=293} + %stack.1342 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.932), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/stack" stack_frame_id=293} + %stack.1343 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1341, %stack.1342), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/stack" stack_frame_id=293} + %reshape.623 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1343), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/reshape"} + %cos.170 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.623), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/cos" stack_frame_id=297} + %convert_element_type.1526 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.170), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/convert_element_type" stack_frame_id=301} + %mul.2755 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1526), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=321} + %mul.2756 = bf16[1,1,128]{2,1,0} reshape(%mul.2755), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=321} + %mul.2757 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2756), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=321} + %mul.2758 = bf16[1,1,8,128]{3,2,1,0} multiply(%dot_general.930, %mul.2757), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=321} + %split.341 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.930), slice={[0:1], [0:1], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/split" stack_frame_id=313} + %neg.349 = bf16[1,1,8,64]{3,2,1,0} negate(%split.341), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/neg" stack_frame_id=314} + %stack.1344 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%neg.349), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/stack" stack_frame_id=317} + %split.340 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.930), slice={[0:1], [0:1], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/split" stack_frame_id=313} + %stack.1345 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%split.340), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/stack" stack_frame_id=317} + %stack.1346 = bf16[1,1,8,2,64]{4,3,2,1,0} concatenate(%stack.1344, %stack.1345), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/stack" stack_frame_id=317} + %reshape.624 = bf16[1,1,8,128]{3,2,1,0} reshape(%stack.1346), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/reshape" stack_frame_id=320} + %sin.170 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.623), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/sin" stack_frame_id=305} + %convert_element_type.1527 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.170), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/convert_element_type" stack_frame_id=309} + %mul.2759 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1527), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=322} + %mul.2760 = bf16[1,1,128]{2,1,0} reshape(%mul.2759), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=322} + %mul.2761 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2760), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=322} + %mul.2762 = bf16[1,1,8,128]{3,2,1,0} multiply(%reshape.624, %mul.2761), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=322} + %add.979 = bf16[1,1,8,128]{3,2,1,0} add(%mul.2758, %mul.2762), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/rotary_embedding_21/add" stack_frame_id=323} + %lt.556 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/lt" stack_frame_id=326} + %add.980 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/add" stack_frame_id=326} + %select_n.133 = s32[] select(%lt.556, %add.980, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/select_n" stack_frame_id=326} + %dynamic_update_slice.107 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.164, %add.979, %constant.127, %select_n.133, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/dynamic_update_slice" stack_frame_id=326} + %broadcast_in_dim.1337 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.107), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/broadcast_in_dim" stack_frame_id=335} + %reshape.625 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1337), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/reshape" stack_frame_id=335} + %dot_general.933 = f32[1,32,1,1,41]{4,3,2,1,0} dot(%reshape.627, %reshape.625), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=83} + %mul.2763 = f32[1,32,1,1,41]{4,3,2,1,0} multiply(%dot_general.933, %broadcast.36), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/vmap()/mul" stack_frame_id=83} + %vmap_jit__where__.84 = f32[1,1,32,1,41]{4,3,2,1,0} call(%and.262, %mul.2763, %constant.114), to_apply=%_where_3.137, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/vmap(jit(_where))" stack_frame_id=83} + %reduce_max.469 = f32[1,1,32,1]{3,2,1,0} reduce(%vmap_jit__where__.84, %constant.113), dimensions={4}, to_apply=%region_215.223, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/vmap()/reduce_max" stack_frame_id=83} + %max.93 = f32[1,1,32,1]{3,2,1,0} maximum(%reduce_max.469, %broadcast.35), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/vmap()/max" stack_frame_id=83} + %broadcast_in_dim.1341 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%max.93), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %sub.355 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1341), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/vmap()/sub" stack_frame_id=83} + %sub.356 = f32[1,1,32,1]{3,2,1,0} reshape(%sub.355), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/vmap()/sub" stack_frame_id=83} + %sub.357 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%sub.356), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/vmap()/sub" stack_frame_id=83} + %sub.358 = f32[1,1,32,1,41]{4,3,2,1,0} subtract(%vmap_jit__where__.84, %sub.357), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/vmap()/sub" stack_frame_id=83} + %exp.89 = f32[1,1,32,1,41]{4,3,2,1,0} exponential(%sub.358), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/vmap()/exp" stack_frame_id=83} + %reduce_sum.1415 = f32[1,1,32,1]{3,2,1,0} reduce(%exp.89, %constant.123), dimensions={4}, to_apply=%region_216.224, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/vmap()/reduce_sum" stack_frame_id=83} + %broadcast_in_dim.1342 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%reduce_sum.1415), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %div.858 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1342), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/vmap()/div" stack_frame_id=83} + %div.859 = f32[1,1,32,1]{3,2,1,0} reshape(%div.858), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/vmap()/div" stack_frame_id=83} + %div.860 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%div.859), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/vmap()/div" stack_frame_id=83} + %div.861 = f32[1,1,32,1,41]{4,3,2,1,0} divide(%exp.89, %div.860), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/vmap()/div" stack_frame_id=83} + %convert_element_type.1529 = bf16[1,1,32,1,41]{4,3,2,1,0} convert(%div.861), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/vmap()/convert_element_type" stack_frame_id=83} + %dot_general.934 = bf16[1,32,128,1,1]{4,3,2,1,0} dot(%reshape.626, %convert_element_type.1529), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=83} + %transpose.88 = bf16[1,1,32,128]{3,2,1,0} reshape(%dot_general.934), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/reshape;jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/vmap()/transpose;jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %dot_general.935 = bf16[1,1,4096]{2,1,0} dot(%transpose.88, %arg_tuple.5#201), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=354} + %add.982 = bf16[1,1,4096]{2,1,0} add(%dot_general.935, %add.971), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/add" stack_frame_id=355} + %convert_element_type.1530 = f32[1,1,4096]{2,1,0} convert(%add.982), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/feedforward_layernorm/convert_element_type" stack_frame_id=359} + %pow.342 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1530, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/feedforward_layernorm/pow" stack_frame_id=362} + %reduce_sum.1416 = f32[1,1]{1,0} reduce(%pow.342, %constant.123), dimensions={2}, to_apply=%region_217.225, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/feedforward_layernorm/reduce_sum" stack_frame_id=92} + %broadcast_in_dim.1343 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1416), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/feedforward_layernorm/broadcast_in_dim" stack_frame_id=92} + %div.862 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1343, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/feedforward_layernorm/div" stack_frame_id=92} + %add.983 = f32[1,1,1]{2,1,0} add(%div.862, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/feedforward_layernorm/add" stack_frame_id=363} + %rsqrt.171 = f32[1,1,1]{2,1,0} rsqrt(%add.983), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/feedforward_layernorm/rsqrt" stack_frame_id=366} + %mul.2764 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.171), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2765 = f32[1,1]{1,0} reshape(%mul.2764), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2766 = f32[1,1,4096]{2,1,0} broadcast(%mul.2765), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2767 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1530, %mul.2766), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/feedforward_layernorm/mul" stack_frame_id=367} + %convert_element_type.1531 = f32[4096]{0} convert(%arg_tuple.5#202), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/feedforward_layernorm/convert_element_type" stack_frame_id=368} + %broadcast_in_dim.1344 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1531), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/feedforward_layernorm/broadcast_in_dim" stack_frame_id=368} + %mul.2768 = f32[1,1,4096]{2,1,0} multiply(%mul.2767, %broadcast_in_dim.1344), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/feedforward_layernorm/mul" stack_frame_id=368} + %convert_element_type.1532 = bf16[1,1,4096]{2,1,0} convert(%mul.2768), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/feedforward_layernorm/convert_element_type" stack_frame_id=372} + %dot_general.937 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1532, %arg_tuple.5#204), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/feedforward_intermediate_dense/dot_general" stack_frame_id=401} + %dot_general.936 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1532, %arg_tuple.5#203), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/feedforward_gate_dense/dot_general" stack_frame_id=381} + %convert_element_type.1533 = f32[1,1,14336]{2,1,0} convert(%dot_general.936), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/convert_element_type" stack_frame_id=385} + %jit_silu_.84 = f32[1,1,14336]{2,1,0} call(%convert_element_type.1533), to_apply=%silu_4.141, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/jit(silu)" stack_frame_id=388} + %convert_element_type.1534 = bf16[1,1,14336]{2,1,0} convert(%jit_silu_.84), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/convert_element_type" stack_frame_id=392} + %mul.2769 = bf16[1,1,14336]{2,1,0} multiply(%dot_general.937, %convert_element_type.1534), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/mul" stack_frame_id=404} + %dot_general.938 = bf16[1,1,4096]{2,1,0} dot(%mul.2769, %arg_tuple.5#205), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/feedforward_output_dense/dot_general" stack_frame_id=413} + %add.984 = bf16[1,1,4096]{2,1,0} add(%dot_general.938, %add.982), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/add" stack_frame_id=414} + %convert_element_type.1538 = f32[1,1,4096]{2,1,0} convert(%add.984), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention_layernorm/convert_element_type" stack_frame_id=164} + %pow.343 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1538, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention_layernorm/pow" stack_frame_id=167} + %reduce_sum.1417 = f32[1,1]{1,0} reduce(%pow.343, %constant.123), dimensions={2}, to_apply=%region_218.226, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention_layernorm/reduce_sum" stack_frame_id=73} + %broadcast_in_dim.1349 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1417), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention_layernorm/broadcast_in_dim" stack_frame_id=73} + %div.863 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1349, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention_layernorm/div" stack_frame_id=73} + %add.988 = f32[1,1,1]{2,1,0} add(%div.863, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention_layernorm/add" stack_frame_id=168} + %rsqrt.172 = f32[1,1,1]{2,1,0} rsqrt(%add.988), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention_layernorm/rsqrt" stack_frame_id=171} + %mul.2770 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.172), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2771 = f32[1,1]{1,0} reshape(%mul.2770), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2772 = f32[1,1,4096]{2,1,0} broadcast(%mul.2771), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2773 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1538, %mul.2772), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention_layernorm/mul" stack_frame_id=172} + %convert_element_type.1539 = f32[4096]{0} convert(%arg_tuple.5#206), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention_layernorm/convert_element_type" stack_frame_id=173} + %broadcast_in_dim.1350 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1539), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention_layernorm/broadcast_in_dim" stack_frame_id=173} + %mul.2774 = f32[1,1,4096]{2,1,0} multiply(%mul.2773, %broadcast_in_dim.1350), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention_layernorm/mul" stack_frame_id=173} + %convert_element_type.1540 = bf16[1,1,4096]{2,1,0} convert(%mul.2774), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention_layernorm/convert_element_type" stack_frame_id=177} + %dot_general.942 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1540, %arg_tuple.5#209), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=265} + %lt.563 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/lt" stack_frame_id=329} + %add.994 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/add" stack_frame_id=329} + %select_n.136 = s32[] select(%lt.563, %add.994, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/select_n" stack_frame_id=329} + %dynamic_update_slice.110 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.168, %dot_general.942, %constant.127, %select_n.136, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/dynamic_update_slice" stack_frame_id=329} + %broadcast_in_dim.1352 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.110), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/broadcast_in_dim" stack_frame_id=338} + %reshape.633 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1352), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/reshape" stack_frame_id=338} + %convert_element_type.1535 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/convert_element_type" stack_frame_id=139} + %add.985 = f32[1,1]{1,0} reshape(%convert_element_type.1535), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_22/add"} + %ge.441 = f32[1,1]{1,0} broadcast(%add.985), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/ge" stack_frame_id=143} + %ge.442 = f32[1]{0} reshape(%ge.441), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/ge" stack_frame_id=143} + %ge.443 = f32[1,41]{1,0} broadcast(%ge.442), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/ge" stack_frame_id=143} + %iota.470 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/iota" stack_frame_id=142} + %broadcast_in_dim.1345 = f32[1,41]{1,0} reshape(%iota.470), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/broadcast_in_dim" stack_frame_id=143} + %ge.444 = pred[1,41]{1,0} compare(%ge.443, %broadcast_in_dim.1345), direction=GE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/ge" stack_frame_id=143} + %broadcast_in_dim.1346 = pred[1,1,41]{2,1,0} reshape(%ge.444), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/broadcast_in_dim" stack_frame_id=146} + %convert_element_type.1537 = s32[1,1,41]{2,1,0} convert(%broadcast_in_dim.1346), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/convert_element_type" stack_frame_id=160} + %add.986 = s32[1,1]{1,0} reshape(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/add" stack_frame_id=147} + %lt.558 = s32[1,1]{1,0} broadcast(%add.986), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/lt" stack_frame_id=152} + %lt.559 = s32[1]{0} reshape(%lt.558), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/lt" stack_frame_id=152} + %lt.560 = s32[1,41]{1,0} broadcast(%lt.559), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/lt" stack_frame_id=152} + %iota.471 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/iota" stack_frame_id=150} + %broadcast_in_dim.1347 = s32[1,41]{1,0} reshape(%iota.471), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/broadcast_in_dim" stack_frame_id=148} + %add.987 = s32[1,41]{1,0} add(%broadcast_in_dim.1347, %broadcast.43), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/add" stack_frame_id=151} + %lt.561 = pred[1,41]{1,0} compare(%lt.560, %add.987), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/lt" stack_frame_id=152} + %convert_element_type.1536 = s32[1,41]{1,0} convert(%lt.561), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/convert_element_type" stack_frame_id=156} + %broadcast_in_dim.1348 = s32[1,1,41]{2,1,0} reshape(%convert_element_type.1536), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/broadcast_in_dim" stack_frame_id=160} + %min.85 = s32[1,1,41]{2,1,0} minimum(%convert_element_type.1537, %broadcast_in_dim.1348), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/min" stack_frame_id=160} + %broadcast_in_dim.1353 = s32[1,1,1,41]{3,2,1,0} reshape(%min.85), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/broadcast_in_dim" stack_frame_id=341} + %convert_element_type.1547 = pred[1,1,1,41]{3,2,1,0} compare(%broadcast_in_dim.1353, %broadcast.37), direction=NE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/convert_element_type" stack_frame_id=345} + %broadcast_in_dim.1354 = pred[1,1,1,1,41]{4,3,2,1,0} reshape(%convert_element_type.1547), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/broadcast_in_dim" stack_frame_id=83} + %and.263 = pred[1,1,1,1,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1354), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/vmap()/and" stack_frame_id=83} + %and.264 = pred[1,1,1,41]{3,2,1,0} reshape(%and.263), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/vmap()/and" stack_frame_id=83} + %and.265 = pred[1,1,32,1,41]{4,3,2,1,0} broadcast(%and.264), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/vmap()/and" stack_frame_id=83} + %dot_general.939 = bf16[1,1,32,128]{3,2,1,0} dot(%convert_element_type.1540, %arg_tuple.5#207), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=186} + %convert_element_type.1541 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/convert_element_type" stack_frame_id=208} + %add.989 = f32[1,1]{1,0} reshape(%convert_element_type.1541), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/add"} + %iota.472 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/iota" stack_frame_id=197} + %mul.2775 = f32[64]{0} multiply(%iota.472, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=197} + %div.864 = f32[64]{0} divide(%mul.2775, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/div" stack_frame_id=198} + %neg.350 = f32[64]{0} negate(%div.864), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/neg" stack_frame_id=199} + %pow.344 = f32[64]{0} power(%broadcast.39, %neg.350), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/pow" stack_frame_id=202} + %div.865 = f32[64]{0} divide(%pow.344, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/div" stack_frame_id=203} + %dot_general.940 = f32[1,1,64]{2,1,0} dot(%add.989, %div.865), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/bi,j->bij/dot_general" stack_frame_id=211} + %stack.1350 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.940), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/stack" stack_frame_id=214} + %stack.1351 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.940), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/stack" stack_frame_id=214} + %stack.1352 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1350, %stack.1351), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/stack" stack_frame_id=214} + %reshape.628 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1352), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/reshape"} + %cos.171 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.628), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/cos" stack_frame_id=218} + %convert_element_type.1542 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.171), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/convert_element_type" stack_frame_id=222} + %mul.2776 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1542), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=242} + %mul.2777 = bf16[1,1,128]{2,1,0} reshape(%mul.2776), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=242} + %mul.2778 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2777), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=242} + %mul.2779 = bf16[1,1,32,128]{3,2,1,0} multiply(%dot_general.939, %mul.2778), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=242} + %split.343 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.939), slice={[0:1], [0:1], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/split" stack_frame_id=234} + %neg.351 = bf16[1,1,32,64]{3,2,1,0} negate(%split.343), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/neg" stack_frame_id=235} + %stack.1353 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%neg.351), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/stack" stack_frame_id=238} + %split.342 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.939), slice={[0:1], [0:1], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/split" stack_frame_id=234} + %stack.1354 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%split.342), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/stack" stack_frame_id=238} + %stack.1355 = bf16[1,1,32,2,64]{4,3,2,1,0} concatenate(%stack.1353, %stack.1354), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/stack" stack_frame_id=238} + %reshape.629 = bf16[1,1,32,128]{3,2,1,0} reshape(%stack.1355), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/reshape" stack_frame_id=241} + %sin.171 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.628), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/sin" stack_frame_id=226} + %convert_element_type.1543 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.171), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/convert_element_type" stack_frame_id=230} + %mul.2780 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1543), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=243} + %mul.2781 = bf16[1,1,128]{2,1,0} reshape(%mul.2780), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=243} + %mul.2782 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2781), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=243} + %mul.2783 = bf16[1,1,32,128]{3,2,1,0} multiply(%reshape.629, %mul.2782), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=243} + %add.990 = bf16[1,1,32,128]{3,2,1,0} add(%mul.2779, %mul.2783), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/add" stack_frame_id=244} + %reshape.634 = bf16[1,1,32,1,128]{4,3,2,1,0} reshape(%add.990), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/reshape" stack_frame_id=83} + %slice.163 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.166), slice={[0:1], [0:1], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/slice" stack_frame_id=245} + %squeeze.167 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.163), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/squeeze" stack_frame_id=245} + %dot_general.941 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1540, %arg_tuple.5#208), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=256} + %convert_element_type.1544 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/convert_element_type" stack_frame_id=287} + %add.991 = f32[1,1]{1,0} reshape(%convert_element_type.1544), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/add"} + %iota.473 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/iota" stack_frame_id=276} + %mul.2784 = f32[64]{0} multiply(%iota.473, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=276} + %div.866 = f32[64]{0} divide(%mul.2784, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/div" stack_frame_id=277} + %neg.352 = f32[64]{0} negate(%div.866), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/neg" stack_frame_id=278} + %pow.345 = f32[64]{0} power(%broadcast.39, %neg.352), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/pow" stack_frame_id=281} + %div.867 = f32[64]{0} divide(%pow.345, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/div" stack_frame_id=282} + %dot_general.943 = f32[1,1,64]{2,1,0} dot(%add.991, %div.867), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/bi,j->bij/dot_general" stack_frame_id=290} + %stack.1356 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.943), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/stack" stack_frame_id=293} + %stack.1357 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.943), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/stack" stack_frame_id=293} + %stack.1358 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1356, %stack.1357), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/stack" stack_frame_id=293} + %reshape.630 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1358), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/reshape"} + %cos.172 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.630), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/cos" stack_frame_id=297} + %convert_element_type.1545 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.172), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/convert_element_type" stack_frame_id=301} + %mul.2785 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1545), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=321} + %mul.2786 = bf16[1,1,128]{2,1,0} reshape(%mul.2785), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=321} + %mul.2787 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2786), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=321} + %mul.2788 = bf16[1,1,8,128]{3,2,1,0} multiply(%dot_general.941, %mul.2787), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=321} + %split.345 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.941), slice={[0:1], [0:1], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/split" stack_frame_id=313} + %neg.353 = bf16[1,1,8,64]{3,2,1,0} negate(%split.345), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/neg" stack_frame_id=314} + %stack.1359 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%neg.353), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/stack" stack_frame_id=317} + %split.344 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.941), slice={[0:1], [0:1], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/split" stack_frame_id=313} + %stack.1360 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%split.344), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/stack" stack_frame_id=317} + %stack.1361 = bf16[1,1,8,2,64]{4,3,2,1,0} concatenate(%stack.1359, %stack.1360), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/stack" stack_frame_id=317} + %reshape.631 = bf16[1,1,8,128]{3,2,1,0} reshape(%stack.1361), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/reshape" stack_frame_id=320} + %sin.172 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.630), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/sin" stack_frame_id=305} + %convert_element_type.1546 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.172), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/convert_element_type" stack_frame_id=309} + %mul.2789 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1546), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=322} + %mul.2790 = bf16[1,1,128]{2,1,0} reshape(%mul.2789), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=322} + %mul.2791 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2790), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=322} + %mul.2792 = bf16[1,1,8,128]{3,2,1,0} multiply(%reshape.631, %mul.2791), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=322} + %add.992 = bf16[1,1,8,128]{3,2,1,0} add(%mul.2788, %mul.2792), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/rotary_embedding_22/add" stack_frame_id=323} + %lt.562 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/lt" stack_frame_id=326} + %add.993 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/add" stack_frame_id=326} + %select_n.135 = s32[] select(%lt.562, %add.993, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/select_n" stack_frame_id=326} + %dynamic_update_slice.109 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.167, %add.992, %constant.127, %select_n.135, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/dynamic_update_slice" stack_frame_id=326} + %broadcast_in_dim.1351 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.109), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/broadcast_in_dim" stack_frame_id=335} + %reshape.632 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1351), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/reshape" stack_frame_id=335} + %dot_general.944 = f32[1,32,1,1,41]{4,3,2,1,0} dot(%reshape.634, %reshape.632), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=83} + %mul.2793 = f32[1,32,1,1,41]{4,3,2,1,0} multiply(%dot_general.944, %broadcast.36), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/vmap()/mul" stack_frame_id=83} + %vmap_jit__where__.85 = f32[1,1,32,1,41]{4,3,2,1,0} call(%and.265, %mul.2793, %constant.114), to_apply=%_where_3.137, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/vmap(jit(_where))" stack_frame_id=83} + %reduce_max.470 = f32[1,1,32,1]{3,2,1,0} reduce(%vmap_jit__where__.85, %constant.113), dimensions={4}, to_apply=%region_219.227, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/vmap()/reduce_max" stack_frame_id=83} + %max.94 = f32[1,1,32,1]{3,2,1,0} maximum(%reduce_max.470, %broadcast.35), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/vmap()/max" stack_frame_id=83} + %broadcast_in_dim.1355 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%max.94), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %sub.359 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1355), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/vmap()/sub" stack_frame_id=83} + %sub.360 = f32[1,1,32,1]{3,2,1,0} reshape(%sub.359), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/vmap()/sub" stack_frame_id=83} + %sub.361 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%sub.360), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/vmap()/sub" stack_frame_id=83} + %sub.362 = f32[1,1,32,1,41]{4,3,2,1,0} subtract(%vmap_jit__where__.85, %sub.361), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/vmap()/sub" stack_frame_id=83} + %exp.90 = f32[1,1,32,1,41]{4,3,2,1,0} exponential(%sub.362), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/vmap()/exp" stack_frame_id=83} + %reduce_sum.1418 = f32[1,1,32,1]{3,2,1,0} reduce(%exp.90, %constant.123), dimensions={4}, to_apply=%region_220.228, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/vmap()/reduce_sum" stack_frame_id=83} + %broadcast_in_dim.1356 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%reduce_sum.1418), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %div.868 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1356), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/vmap()/div" stack_frame_id=83} + %div.869 = f32[1,1,32,1]{3,2,1,0} reshape(%div.868), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/vmap()/div" stack_frame_id=83} + %div.870 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%div.869), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/vmap()/div" stack_frame_id=83} + %div.871 = f32[1,1,32,1,41]{4,3,2,1,0} divide(%exp.90, %div.870), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/vmap()/div" stack_frame_id=83} + %convert_element_type.1548 = bf16[1,1,32,1,41]{4,3,2,1,0} convert(%div.871), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/vmap()/convert_element_type" stack_frame_id=83} + %dot_general.945 = bf16[1,32,128,1,1]{4,3,2,1,0} dot(%reshape.633, %convert_element_type.1548), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=83} + %transpose.89 = bf16[1,1,32,128]{3,2,1,0} reshape(%dot_general.945), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/reshape;jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/vmap()/transpose;jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %dot_general.946 = bf16[1,1,4096]{2,1,0} dot(%transpose.89, %arg_tuple.5#210), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=354} + %add.995 = bf16[1,1,4096]{2,1,0} add(%dot_general.946, %add.984), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/add" stack_frame_id=355} + %convert_element_type.1549 = f32[1,1,4096]{2,1,0} convert(%add.995), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/feedforward_layernorm/convert_element_type" stack_frame_id=359} + %pow.346 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1549, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/feedforward_layernorm/pow" stack_frame_id=362} + %reduce_sum.1419 = f32[1,1]{1,0} reduce(%pow.346, %constant.123), dimensions={2}, to_apply=%region_221.229, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/feedforward_layernorm/reduce_sum" stack_frame_id=92} + %broadcast_in_dim.1357 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1419), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/feedforward_layernorm/broadcast_in_dim" stack_frame_id=92} + %div.872 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1357, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/feedforward_layernorm/div" stack_frame_id=92} + %add.996 = f32[1,1,1]{2,1,0} add(%div.872, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/feedforward_layernorm/add" stack_frame_id=363} + %rsqrt.173 = f32[1,1,1]{2,1,0} rsqrt(%add.996), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/feedforward_layernorm/rsqrt" stack_frame_id=366} + %mul.2794 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.173), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2795 = f32[1,1]{1,0} reshape(%mul.2794), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2796 = f32[1,1,4096]{2,1,0} broadcast(%mul.2795), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2797 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1549, %mul.2796), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/feedforward_layernorm/mul" stack_frame_id=367} + %convert_element_type.1550 = f32[4096]{0} convert(%arg_tuple.5#211), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/feedforward_layernorm/convert_element_type" stack_frame_id=368} + %broadcast_in_dim.1358 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1550), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/feedforward_layernorm/broadcast_in_dim" stack_frame_id=368} + %mul.2798 = f32[1,1,4096]{2,1,0} multiply(%mul.2797, %broadcast_in_dim.1358), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/feedforward_layernorm/mul" stack_frame_id=368} + %convert_element_type.1551 = bf16[1,1,4096]{2,1,0} convert(%mul.2798), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/feedforward_layernorm/convert_element_type" stack_frame_id=372} + %dot_general.948 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1551, %arg_tuple.5#213), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/feedforward_intermediate_dense/dot_general" stack_frame_id=401} + %dot_general.947 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1551, %arg_tuple.5#212), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/feedforward_gate_dense/dot_general" stack_frame_id=381} + %convert_element_type.1552 = f32[1,1,14336]{2,1,0} convert(%dot_general.947), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/convert_element_type" stack_frame_id=385} + %jit_silu_.85 = f32[1,1,14336]{2,1,0} call(%convert_element_type.1552), to_apply=%silu_4.141, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/jit(silu)" stack_frame_id=388} + %convert_element_type.1553 = bf16[1,1,14336]{2,1,0} convert(%jit_silu_.85), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/convert_element_type" stack_frame_id=392} + %mul.2799 = bf16[1,1,14336]{2,1,0} multiply(%dot_general.948, %convert_element_type.1553), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/mul" stack_frame_id=404} + %dot_general.949 = bf16[1,1,4096]{2,1,0} dot(%mul.2799, %arg_tuple.5#214), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/feedforward_output_dense/dot_general" stack_frame_id=413} + %add.997 = bf16[1,1,4096]{2,1,0} add(%dot_general.949, %add.995), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/add" stack_frame_id=414} + %convert_element_type.1557 = f32[1,1,4096]{2,1,0} convert(%add.997), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention_layernorm/convert_element_type" stack_frame_id=164} + %pow.347 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1557, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention_layernorm/pow" stack_frame_id=167} + %reduce_sum.1420 = f32[1,1]{1,0} reduce(%pow.347, %constant.123), dimensions={2}, to_apply=%region_222.230, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention_layernorm/reduce_sum" stack_frame_id=73} + %broadcast_in_dim.1363 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1420), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention_layernorm/broadcast_in_dim" stack_frame_id=73} + %div.873 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1363, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention_layernorm/div" stack_frame_id=73} + %add.1001 = f32[1,1,1]{2,1,0} add(%div.873, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention_layernorm/add" stack_frame_id=168} + %rsqrt.174 = f32[1,1,1]{2,1,0} rsqrt(%add.1001), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention_layernorm/rsqrt" stack_frame_id=171} + %mul.2800 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.174), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2801 = f32[1,1]{1,0} reshape(%mul.2800), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2802 = f32[1,1,4096]{2,1,0} broadcast(%mul.2801), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2803 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1557, %mul.2802), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention_layernorm/mul" stack_frame_id=172} + %convert_element_type.1558 = f32[4096]{0} convert(%arg_tuple.5#215), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention_layernorm/convert_element_type" stack_frame_id=173} + %broadcast_in_dim.1364 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1558), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention_layernorm/broadcast_in_dim" stack_frame_id=173} + %mul.2804 = f32[1,1,4096]{2,1,0} multiply(%mul.2803, %broadcast_in_dim.1364), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention_layernorm/mul" stack_frame_id=173} + %convert_element_type.1559 = bf16[1,1,4096]{2,1,0} convert(%mul.2804), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention_layernorm/convert_element_type" stack_frame_id=177} + %dot_general.953 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1559, %arg_tuple.5#218), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=265} + %lt.569 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/lt" stack_frame_id=329} + %add.1007 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/add" stack_frame_id=329} + %select_n.138 = s32[] select(%lt.569, %add.1007, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/select_n" stack_frame_id=329} + %dynamic_update_slice.112 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.171, %dot_general.953, %constant.127, %select_n.138, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/dynamic_update_slice" stack_frame_id=329} + %broadcast_in_dim.1366 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.112), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/broadcast_in_dim" stack_frame_id=338} + %reshape.640 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1366), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/reshape" stack_frame_id=338} + %convert_element_type.1554 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/convert_element_type" stack_frame_id=139} + %add.998 = f32[1,1]{1,0} reshape(%convert_element_type.1554), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_23/add"} + %ge.445 = f32[1,1]{1,0} broadcast(%add.998), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/ge" stack_frame_id=143} + %ge.446 = f32[1]{0} reshape(%ge.445), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/ge" stack_frame_id=143} + %ge.447 = f32[1,41]{1,0} broadcast(%ge.446), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/ge" stack_frame_id=143} + %iota.474 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/iota" stack_frame_id=142} + %broadcast_in_dim.1359 = f32[1,41]{1,0} reshape(%iota.474), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/broadcast_in_dim" stack_frame_id=143} + %ge.448 = pred[1,41]{1,0} compare(%ge.447, %broadcast_in_dim.1359), direction=GE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/ge" stack_frame_id=143} + %broadcast_in_dim.1360 = pred[1,1,41]{2,1,0} reshape(%ge.448), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/broadcast_in_dim" stack_frame_id=146} + %convert_element_type.1556 = s32[1,1,41]{2,1,0} convert(%broadcast_in_dim.1360), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/convert_element_type" stack_frame_id=160} + %add.999 = s32[1,1]{1,0} reshape(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/add" stack_frame_id=147} + %lt.564 = s32[1,1]{1,0} broadcast(%add.999), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/lt" stack_frame_id=152} + %lt.565 = s32[1]{0} reshape(%lt.564), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/lt" stack_frame_id=152} + %lt.566 = s32[1,41]{1,0} broadcast(%lt.565), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/lt" stack_frame_id=152} + %iota.475 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/iota" stack_frame_id=150} + %broadcast_in_dim.1361 = s32[1,41]{1,0} reshape(%iota.475), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/broadcast_in_dim" stack_frame_id=148} + %add.1000 = s32[1,41]{1,0} add(%broadcast_in_dim.1361, %broadcast.43), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/add" stack_frame_id=151} + %lt.567 = pred[1,41]{1,0} compare(%lt.566, %add.1000), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/lt" stack_frame_id=152} + %convert_element_type.1555 = s32[1,41]{1,0} convert(%lt.567), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/convert_element_type" stack_frame_id=156} + %broadcast_in_dim.1362 = s32[1,1,41]{2,1,0} reshape(%convert_element_type.1555), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/broadcast_in_dim" stack_frame_id=160} + %min.86 = s32[1,1,41]{2,1,0} minimum(%convert_element_type.1556, %broadcast_in_dim.1362), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/min" stack_frame_id=160} + %broadcast_in_dim.1367 = s32[1,1,1,41]{3,2,1,0} reshape(%min.86), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/broadcast_in_dim" stack_frame_id=341} + %convert_element_type.1566 = pred[1,1,1,41]{3,2,1,0} compare(%broadcast_in_dim.1367, %broadcast.37), direction=NE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/convert_element_type" stack_frame_id=345} + %broadcast_in_dim.1368 = pred[1,1,1,1,41]{4,3,2,1,0} reshape(%convert_element_type.1566), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/broadcast_in_dim" stack_frame_id=83} + %and.266 = pred[1,1,1,1,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1368), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/vmap()/and" stack_frame_id=83} + %and.267 = pred[1,1,1,41]{3,2,1,0} reshape(%and.266), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/vmap()/and" stack_frame_id=83} + %and.268 = pred[1,1,32,1,41]{4,3,2,1,0} broadcast(%and.267), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/vmap()/and" stack_frame_id=83} + %dot_general.950 = bf16[1,1,32,128]{3,2,1,0} dot(%convert_element_type.1559, %arg_tuple.5#216), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=186} + %convert_element_type.1560 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/convert_element_type" stack_frame_id=208} + %add.1002 = f32[1,1]{1,0} reshape(%convert_element_type.1560), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/add"} + %iota.476 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/iota" stack_frame_id=197} + %mul.2805 = f32[64]{0} multiply(%iota.476, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=197} + %div.874 = f32[64]{0} divide(%mul.2805, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/div" stack_frame_id=198} + %neg.354 = f32[64]{0} negate(%div.874), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/neg" stack_frame_id=199} + %pow.348 = f32[64]{0} power(%broadcast.39, %neg.354), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/pow" stack_frame_id=202} + %div.875 = f32[64]{0} divide(%pow.348, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/div" stack_frame_id=203} + %dot_general.951 = f32[1,1,64]{2,1,0} dot(%add.1002, %div.875), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/bi,j->bij/dot_general" stack_frame_id=211} + %stack.1365 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.951), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/stack" stack_frame_id=214} + %stack.1366 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.951), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/stack" stack_frame_id=214} + %stack.1367 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1365, %stack.1366), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/stack" stack_frame_id=214} + %reshape.635 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1367), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/reshape"} + %cos.173 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.635), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/cos" stack_frame_id=218} + %convert_element_type.1561 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.173), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/convert_element_type" stack_frame_id=222} + %mul.2806 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1561), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=242} + %mul.2807 = bf16[1,1,128]{2,1,0} reshape(%mul.2806), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=242} + %mul.2808 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2807), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=242} + %mul.2809 = bf16[1,1,32,128]{3,2,1,0} multiply(%dot_general.950, %mul.2808), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=242} + %split.347 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.950), slice={[0:1], [0:1], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/split" stack_frame_id=234} + %neg.355 = bf16[1,1,32,64]{3,2,1,0} negate(%split.347), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/neg" stack_frame_id=235} + %stack.1368 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%neg.355), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/stack" stack_frame_id=238} + %split.346 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.950), slice={[0:1], [0:1], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/split" stack_frame_id=234} + %stack.1369 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%split.346), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/stack" stack_frame_id=238} + %stack.1370 = bf16[1,1,32,2,64]{4,3,2,1,0} concatenate(%stack.1368, %stack.1369), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/stack" stack_frame_id=238} + %reshape.636 = bf16[1,1,32,128]{3,2,1,0} reshape(%stack.1370), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/reshape" stack_frame_id=241} + %sin.173 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.635), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/sin" stack_frame_id=226} + %convert_element_type.1562 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.173), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/convert_element_type" stack_frame_id=230} + %mul.2810 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1562), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=243} + %mul.2811 = bf16[1,1,128]{2,1,0} reshape(%mul.2810), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=243} + %mul.2812 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2811), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=243} + %mul.2813 = bf16[1,1,32,128]{3,2,1,0} multiply(%reshape.636, %mul.2812), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=243} + %add.1003 = bf16[1,1,32,128]{3,2,1,0} add(%mul.2809, %mul.2813), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/add" stack_frame_id=244} + %reshape.641 = bf16[1,1,32,1,128]{4,3,2,1,0} reshape(%add.1003), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/reshape" stack_frame_id=83} + %slice.166 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.169), slice={[0:1], [0:1], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/slice" stack_frame_id=245} + %squeeze.170 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.166), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/squeeze" stack_frame_id=245} + %dot_general.952 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1559, %arg_tuple.5#217), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=256} + %convert_element_type.1563 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/convert_element_type" stack_frame_id=287} + %add.1004 = f32[1,1]{1,0} reshape(%convert_element_type.1563), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/add"} + %iota.477 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/iota" stack_frame_id=276} + %mul.2814 = f32[64]{0} multiply(%iota.477, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=276} + %div.876 = f32[64]{0} divide(%mul.2814, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/div" stack_frame_id=277} + %neg.356 = f32[64]{0} negate(%div.876), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/neg" stack_frame_id=278} + %pow.349 = f32[64]{0} power(%broadcast.39, %neg.356), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/pow" stack_frame_id=281} + %div.877 = f32[64]{0} divide(%pow.349, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/div" stack_frame_id=282} + %dot_general.954 = f32[1,1,64]{2,1,0} dot(%add.1004, %div.877), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/bi,j->bij/dot_general" stack_frame_id=290} + %stack.1371 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.954), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/stack" stack_frame_id=293} + %stack.1372 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.954), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/stack" stack_frame_id=293} + %stack.1373 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1371, %stack.1372), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/stack" stack_frame_id=293} + %reshape.637 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1373), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/reshape"} + %cos.174 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.637), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/cos" stack_frame_id=297} + %convert_element_type.1564 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.174), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/convert_element_type" stack_frame_id=301} + %mul.2815 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1564), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=321} + %mul.2816 = bf16[1,1,128]{2,1,0} reshape(%mul.2815), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=321} + %mul.2817 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2816), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=321} + %mul.2818 = bf16[1,1,8,128]{3,2,1,0} multiply(%dot_general.952, %mul.2817), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=321} + %split.349 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.952), slice={[0:1], [0:1], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/split" stack_frame_id=313} + %neg.357 = bf16[1,1,8,64]{3,2,1,0} negate(%split.349), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/neg" stack_frame_id=314} + %stack.1374 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%neg.357), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/stack" stack_frame_id=317} + %split.348 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.952), slice={[0:1], [0:1], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/split" stack_frame_id=313} + %stack.1375 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%split.348), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/stack" stack_frame_id=317} + %stack.1376 = bf16[1,1,8,2,64]{4,3,2,1,0} concatenate(%stack.1374, %stack.1375), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/stack" stack_frame_id=317} + %reshape.638 = bf16[1,1,8,128]{3,2,1,0} reshape(%stack.1376), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/reshape" stack_frame_id=320} + %sin.174 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.637), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/sin" stack_frame_id=305} + %convert_element_type.1565 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.174), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/convert_element_type" stack_frame_id=309} + %mul.2819 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1565), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=322} + %mul.2820 = bf16[1,1,128]{2,1,0} reshape(%mul.2819), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=322} + %mul.2821 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2820), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=322} + %mul.2822 = bf16[1,1,8,128]{3,2,1,0} multiply(%reshape.638, %mul.2821), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=322} + %add.1005 = bf16[1,1,8,128]{3,2,1,0} add(%mul.2818, %mul.2822), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/rotary_embedding_23/add" stack_frame_id=323} + %lt.568 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/lt" stack_frame_id=326} + %add.1006 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/add" stack_frame_id=326} + %select_n.137 = s32[] select(%lt.568, %add.1006, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/select_n" stack_frame_id=326} + %dynamic_update_slice.111 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.170, %add.1005, %constant.127, %select_n.137, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/dynamic_update_slice" stack_frame_id=326} + %broadcast_in_dim.1365 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.111), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/broadcast_in_dim" stack_frame_id=335} + %reshape.639 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1365), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/reshape" stack_frame_id=335} + %dot_general.955 = f32[1,32,1,1,41]{4,3,2,1,0} dot(%reshape.641, %reshape.639), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=83} + %mul.2823 = f32[1,32,1,1,41]{4,3,2,1,0} multiply(%dot_general.955, %broadcast.36), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/vmap()/mul" stack_frame_id=83} + %vmap_jit__where__.86 = f32[1,1,32,1,41]{4,3,2,1,0} call(%and.268, %mul.2823, %constant.114), to_apply=%_where_3.137, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/vmap(jit(_where))" stack_frame_id=83} + %reduce_max.471 = f32[1,1,32,1]{3,2,1,0} reduce(%vmap_jit__where__.86, %constant.113), dimensions={4}, to_apply=%region_223.231, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/vmap()/reduce_max" stack_frame_id=83} + %max.95 = f32[1,1,32,1]{3,2,1,0} maximum(%reduce_max.471, %broadcast.35), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/vmap()/max" stack_frame_id=83} + %broadcast_in_dim.1369 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%max.95), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %sub.363 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1369), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/vmap()/sub" stack_frame_id=83} + %sub.364 = f32[1,1,32,1]{3,2,1,0} reshape(%sub.363), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/vmap()/sub" stack_frame_id=83} + %sub.365 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%sub.364), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/vmap()/sub" stack_frame_id=83} + %sub.366 = f32[1,1,32,1,41]{4,3,2,1,0} subtract(%vmap_jit__where__.86, %sub.365), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/vmap()/sub" stack_frame_id=83} + %exp.91 = f32[1,1,32,1,41]{4,3,2,1,0} exponential(%sub.366), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/vmap()/exp" stack_frame_id=83} + %reduce_sum.1421 = f32[1,1,32,1]{3,2,1,0} reduce(%exp.91, %constant.123), dimensions={4}, to_apply=%region_224.232, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/vmap()/reduce_sum" stack_frame_id=83} + %broadcast_in_dim.1370 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%reduce_sum.1421), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %div.878 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1370), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/vmap()/div" stack_frame_id=83} + %div.879 = f32[1,1,32,1]{3,2,1,0} reshape(%div.878), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/vmap()/div" stack_frame_id=83} + %div.880 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%div.879), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/vmap()/div" stack_frame_id=83} + %div.881 = f32[1,1,32,1,41]{4,3,2,1,0} divide(%exp.91, %div.880), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/vmap()/div" stack_frame_id=83} + %convert_element_type.1567 = bf16[1,1,32,1,41]{4,3,2,1,0} convert(%div.881), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/vmap()/convert_element_type" stack_frame_id=83} + %dot_general.956 = bf16[1,32,128,1,1]{4,3,2,1,0} dot(%reshape.640, %convert_element_type.1567), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=83} + %transpose.90 = bf16[1,1,32,128]{3,2,1,0} reshape(%dot_general.956), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/reshape;jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/vmap()/transpose;jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %dot_general.957 = bf16[1,1,4096]{2,1,0} dot(%transpose.90, %arg_tuple.5#219), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=354} + %add.1008 = bf16[1,1,4096]{2,1,0} add(%dot_general.957, %add.997), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/add" stack_frame_id=355} + %convert_element_type.1568 = f32[1,1,4096]{2,1,0} convert(%add.1008), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/feedforward_layernorm/convert_element_type" stack_frame_id=359} + %pow.350 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1568, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/feedforward_layernorm/pow" stack_frame_id=362} + %reduce_sum.1422 = f32[1,1]{1,0} reduce(%pow.350, %constant.123), dimensions={2}, to_apply=%region_225.233, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/feedforward_layernorm/reduce_sum" stack_frame_id=92} + %broadcast_in_dim.1371 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1422), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/feedforward_layernorm/broadcast_in_dim" stack_frame_id=92} + %div.882 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1371, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/feedforward_layernorm/div" stack_frame_id=92} + %add.1009 = f32[1,1,1]{2,1,0} add(%div.882, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/feedforward_layernorm/add" stack_frame_id=363} + %rsqrt.175 = f32[1,1,1]{2,1,0} rsqrt(%add.1009), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/feedforward_layernorm/rsqrt" stack_frame_id=366} + %mul.2824 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.175), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2825 = f32[1,1]{1,0} reshape(%mul.2824), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2826 = f32[1,1,4096]{2,1,0} broadcast(%mul.2825), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2827 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1568, %mul.2826), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/feedforward_layernorm/mul" stack_frame_id=367} + %convert_element_type.1569 = f32[4096]{0} convert(%arg_tuple.5#220), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/feedforward_layernorm/convert_element_type" stack_frame_id=368} + %broadcast_in_dim.1372 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1569), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/feedforward_layernorm/broadcast_in_dim" stack_frame_id=368} + %mul.2828 = f32[1,1,4096]{2,1,0} multiply(%mul.2827, %broadcast_in_dim.1372), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/feedforward_layernorm/mul" stack_frame_id=368} + %convert_element_type.1570 = bf16[1,1,4096]{2,1,0} convert(%mul.2828), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/feedforward_layernorm/convert_element_type" stack_frame_id=372} + %dot_general.959 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1570, %arg_tuple.5#222), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/feedforward_intermediate_dense/dot_general" stack_frame_id=401} + %dot_general.958 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1570, %arg_tuple.5#221), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/feedforward_gate_dense/dot_general" stack_frame_id=381} + %convert_element_type.1571 = f32[1,1,14336]{2,1,0} convert(%dot_general.958), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/convert_element_type" stack_frame_id=385} + %jit_silu_.86 = f32[1,1,14336]{2,1,0} call(%convert_element_type.1571), to_apply=%silu_4.141, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/jit(silu)" stack_frame_id=388} + %convert_element_type.1572 = bf16[1,1,14336]{2,1,0} convert(%jit_silu_.86), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/convert_element_type" stack_frame_id=392} + %mul.2829 = bf16[1,1,14336]{2,1,0} multiply(%dot_general.959, %convert_element_type.1572), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/mul" stack_frame_id=404} + %dot_general.960 = bf16[1,1,4096]{2,1,0} dot(%mul.2829, %arg_tuple.5#223), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/feedforward_output_dense/dot_general" stack_frame_id=413} + %add.1010 = bf16[1,1,4096]{2,1,0} add(%dot_general.960, %add.1008), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/add" stack_frame_id=414} + %convert_element_type.1576 = f32[1,1,4096]{2,1,0} convert(%add.1010), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention_layernorm/convert_element_type" stack_frame_id=164} + %pow.351 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1576, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention_layernorm/pow" stack_frame_id=167} + %reduce_sum.1423 = f32[1,1]{1,0} reduce(%pow.351, %constant.123), dimensions={2}, to_apply=%region_226.234, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention_layernorm/reduce_sum" stack_frame_id=73} + %broadcast_in_dim.1377 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1423), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention_layernorm/broadcast_in_dim" stack_frame_id=73} + %div.883 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1377, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention_layernorm/div" stack_frame_id=73} + %add.1014 = f32[1,1,1]{2,1,0} add(%div.883, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention_layernorm/add" stack_frame_id=168} + %rsqrt.176 = f32[1,1,1]{2,1,0} rsqrt(%add.1014), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention_layernorm/rsqrt" stack_frame_id=171} + %mul.2830 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.176), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2831 = f32[1,1]{1,0} reshape(%mul.2830), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2832 = f32[1,1,4096]{2,1,0} broadcast(%mul.2831), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2833 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1576, %mul.2832), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention_layernorm/mul" stack_frame_id=172} + %convert_element_type.1577 = f32[4096]{0} convert(%arg_tuple.5#224), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention_layernorm/convert_element_type" stack_frame_id=173} + %broadcast_in_dim.1378 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1577), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention_layernorm/broadcast_in_dim" stack_frame_id=173} + %mul.2834 = f32[1,1,4096]{2,1,0} multiply(%mul.2833, %broadcast_in_dim.1378), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention_layernorm/mul" stack_frame_id=173} + %convert_element_type.1578 = bf16[1,1,4096]{2,1,0} convert(%mul.2834), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention_layernorm/convert_element_type" stack_frame_id=177} + %dot_general.964 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1578, %arg_tuple.5#227), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=265} + %lt.575 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/lt" stack_frame_id=329} + %add.1020 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/add" stack_frame_id=329} + %select_n.140 = s32[] select(%lt.575, %add.1020, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/select_n" stack_frame_id=329} + %dynamic_update_slice.114 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.174, %dot_general.964, %constant.127, %select_n.140, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/dynamic_update_slice" stack_frame_id=329} + %broadcast_in_dim.1380 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.114), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/broadcast_in_dim" stack_frame_id=338} + %reshape.647 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1380), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/reshape" stack_frame_id=338} + %convert_element_type.1573 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/convert_element_type" stack_frame_id=139} + %add.1011 = f32[1,1]{1,0} reshape(%convert_element_type.1573), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_24/add"} + %ge.449 = f32[1,1]{1,0} broadcast(%add.1011), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/ge" stack_frame_id=143} + %ge.450 = f32[1]{0} reshape(%ge.449), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/ge" stack_frame_id=143} + %ge.451 = f32[1,41]{1,0} broadcast(%ge.450), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/ge" stack_frame_id=143} + %iota.478 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/iota" stack_frame_id=142} + %broadcast_in_dim.1373 = f32[1,41]{1,0} reshape(%iota.478), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/broadcast_in_dim" stack_frame_id=143} + %ge.452 = pred[1,41]{1,0} compare(%ge.451, %broadcast_in_dim.1373), direction=GE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/ge" stack_frame_id=143} + %broadcast_in_dim.1374 = pred[1,1,41]{2,1,0} reshape(%ge.452), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/broadcast_in_dim" stack_frame_id=146} + %convert_element_type.1575 = s32[1,1,41]{2,1,0} convert(%broadcast_in_dim.1374), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/convert_element_type" stack_frame_id=160} + %add.1012 = s32[1,1]{1,0} reshape(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/add" stack_frame_id=147} + %lt.570 = s32[1,1]{1,0} broadcast(%add.1012), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/lt" stack_frame_id=152} + %lt.571 = s32[1]{0} reshape(%lt.570), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/lt" stack_frame_id=152} + %lt.572 = s32[1,41]{1,0} broadcast(%lt.571), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/lt" stack_frame_id=152} + %iota.479 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/iota" stack_frame_id=150} + %broadcast_in_dim.1375 = s32[1,41]{1,0} reshape(%iota.479), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/broadcast_in_dim" stack_frame_id=148} + %add.1013 = s32[1,41]{1,0} add(%broadcast_in_dim.1375, %broadcast.43), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/add" stack_frame_id=151} + %lt.573 = pred[1,41]{1,0} compare(%lt.572, %add.1013), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/lt" stack_frame_id=152} + %convert_element_type.1574 = s32[1,41]{1,0} convert(%lt.573), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/convert_element_type" stack_frame_id=156} + %broadcast_in_dim.1376 = s32[1,1,41]{2,1,0} reshape(%convert_element_type.1574), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/broadcast_in_dim" stack_frame_id=160} + %min.87 = s32[1,1,41]{2,1,0} minimum(%convert_element_type.1575, %broadcast_in_dim.1376), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/min" stack_frame_id=160} + %broadcast_in_dim.1381 = s32[1,1,1,41]{3,2,1,0} reshape(%min.87), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/broadcast_in_dim" stack_frame_id=341} + %convert_element_type.1585 = pred[1,1,1,41]{3,2,1,0} compare(%broadcast_in_dim.1381, %broadcast.37), direction=NE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/convert_element_type" stack_frame_id=345} + %broadcast_in_dim.1382 = pred[1,1,1,1,41]{4,3,2,1,0} reshape(%convert_element_type.1585), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/broadcast_in_dim" stack_frame_id=83} + %and.269 = pred[1,1,1,1,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1382), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/vmap()/and" stack_frame_id=83} + %and.270 = pred[1,1,1,41]{3,2,1,0} reshape(%and.269), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/vmap()/and" stack_frame_id=83} + %and.271 = pred[1,1,32,1,41]{4,3,2,1,0} broadcast(%and.270), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/vmap()/and" stack_frame_id=83} + %dot_general.961 = bf16[1,1,32,128]{3,2,1,0} dot(%convert_element_type.1578, %arg_tuple.5#225), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=186} + %convert_element_type.1579 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/convert_element_type" stack_frame_id=208} + %add.1015 = f32[1,1]{1,0} reshape(%convert_element_type.1579), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/add"} + %iota.480 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/iota" stack_frame_id=197} + %mul.2835 = f32[64]{0} multiply(%iota.480, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=197} + %div.884 = f32[64]{0} divide(%mul.2835, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/div" stack_frame_id=198} + %neg.358 = f32[64]{0} negate(%div.884), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/neg" stack_frame_id=199} + %pow.352 = f32[64]{0} power(%broadcast.39, %neg.358), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/pow" stack_frame_id=202} + %div.885 = f32[64]{0} divide(%pow.352, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/div" stack_frame_id=203} + %dot_general.962 = f32[1,1,64]{2,1,0} dot(%add.1015, %div.885), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/bi,j->bij/dot_general" stack_frame_id=211} + %stack.1380 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.962), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/stack" stack_frame_id=214} + %stack.1381 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.962), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/stack" stack_frame_id=214} + %stack.1382 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1380, %stack.1381), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/stack" stack_frame_id=214} + %reshape.642 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1382), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/reshape"} + %cos.175 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.642), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/cos" stack_frame_id=218} + %convert_element_type.1580 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.175), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/convert_element_type" stack_frame_id=222} + %mul.2836 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1580), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=242} + %mul.2837 = bf16[1,1,128]{2,1,0} reshape(%mul.2836), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=242} + %mul.2838 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2837), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=242} + %mul.2839 = bf16[1,1,32,128]{3,2,1,0} multiply(%dot_general.961, %mul.2838), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=242} + %split.351 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.961), slice={[0:1], [0:1], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/split" stack_frame_id=234} + %neg.359 = bf16[1,1,32,64]{3,2,1,0} negate(%split.351), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/neg" stack_frame_id=235} + %stack.1383 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%neg.359), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/stack" stack_frame_id=238} + %split.350 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.961), slice={[0:1], [0:1], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/split" stack_frame_id=234} + %stack.1384 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%split.350), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/stack" stack_frame_id=238} + %stack.1385 = bf16[1,1,32,2,64]{4,3,2,1,0} concatenate(%stack.1383, %stack.1384), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/stack" stack_frame_id=238} + %reshape.643 = bf16[1,1,32,128]{3,2,1,0} reshape(%stack.1385), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/reshape" stack_frame_id=241} + %sin.175 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.642), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/sin" stack_frame_id=226} + %convert_element_type.1581 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.175), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/convert_element_type" stack_frame_id=230} + %mul.2840 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1581), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=243} + %mul.2841 = bf16[1,1,128]{2,1,0} reshape(%mul.2840), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=243} + %mul.2842 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2841), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=243} + %mul.2843 = bf16[1,1,32,128]{3,2,1,0} multiply(%reshape.643, %mul.2842), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=243} + %add.1016 = bf16[1,1,32,128]{3,2,1,0} add(%mul.2839, %mul.2843), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/add" stack_frame_id=244} + %reshape.648 = bf16[1,1,32,1,128]{4,3,2,1,0} reshape(%add.1016), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/reshape" stack_frame_id=83} + %slice.169 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.172), slice={[0:1], [0:1], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/slice" stack_frame_id=245} + %squeeze.173 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.169), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/squeeze" stack_frame_id=245} + %dot_general.963 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1578, %arg_tuple.5#226), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=256} + %convert_element_type.1582 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/convert_element_type" stack_frame_id=287} + %add.1017 = f32[1,1]{1,0} reshape(%convert_element_type.1582), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/add"} + %iota.481 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/iota" stack_frame_id=276} + %mul.2844 = f32[64]{0} multiply(%iota.481, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=276} + %div.886 = f32[64]{0} divide(%mul.2844, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/div" stack_frame_id=277} + %neg.360 = f32[64]{0} negate(%div.886), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/neg" stack_frame_id=278} + %pow.353 = f32[64]{0} power(%broadcast.39, %neg.360), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/pow" stack_frame_id=281} + %div.887 = f32[64]{0} divide(%pow.353, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/div" stack_frame_id=282} + %dot_general.965 = f32[1,1,64]{2,1,0} dot(%add.1017, %div.887), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/bi,j->bij/dot_general" stack_frame_id=290} + %stack.1386 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.965), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/stack" stack_frame_id=293} + %stack.1387 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.965), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/stack" stack_frame_id=293} + %stack.1388 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1386, %stack.1387), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/stack" stack_frame_id=293} + %reshape.644 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1388), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/reshape"} + %cos.176 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.644), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/cos" stack_frame_id=297} + %convert_element_type.1583 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.176), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/convert_element_type" stack_frame_id=301} + %mul.2845 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1583), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=321} + %mul.2846 = bf16[1,1,128]{2,1,0} reshape(%mul.2845), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=321} + %mul.2847 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2846), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=321} + %mul.2848 = bf16[1,1,8,128]{3,2,1,0} multiply(%dot_general.963, %mul.2847), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=321} + %split.353 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.963), slice={[0:1], [0:1], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/split" stack_frame_id=313} + %neg.361 = bf16[1,1,8,64]{3,2,1,0} negate(%split.353), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/neg" stack_frame_id=314} + %stack.1389 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%neg.361), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/stack" stack_frame_id=317} + %split.352 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.963), slice={[0:1], [0:1], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/split" stack_frame_id=313} + %stack.1390 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%split.352), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/stack" stack_frame_id=317} + %stack.1391 = bf16[1,1,8,2,64]{4,3,2,1,0} concatenate(%stack.1389, %stack.1390), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/stack" stack_frame_id=317} + %reshape.645 = bf16[1,1,8,128]{3,2,1,0} reshape(%stack.1391), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/reshape" stack_frame_id=320} + %sin.176 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.644), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/sin" stack_frame_id=305} + %convert_element_type.1584 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.176), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/convert_element_type" stack_frame_id=309} + %mul.2849 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1584), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=322} + %mul.2850 = bf16[1,1,128]{2,1,0} reshape(%mul.2849), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=322} + %mul.2851 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2850), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=322} + %mul.2852 = bf16[1,1,8,128]{3,2,1,0} multiply(%reshape.645, %mul.2851), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=322} + %add.1018 = bf16[1,1,8,128]{3,2,1,0} add(%mul.2848, %mul.2852), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/rotary_embedding_24/add" stack_frame_id=323} + %lt.574 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/lt" stack_frame_id=326} + %add.1019 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/add" stack_frame_id=326} + %select_n.139 = s32[] select(%lt.574, %add.1019, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/select_n" stack_frame_id=326} + %dynamic_update_slice.113 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.173, %add.1018, %constant.127, %select_n.139, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/dynamic_update_slice" stack_frame_id=326} + %broadcast_in_dim.1379 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.113), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/broadcast_in_dim" stack_frame_id=335} + %reshape.646 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1379), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/reshape" stack_frame_id=335} + %dot_general.966 = f32[1,32,1,1,41]{4,3,2,1,0} dot(%reshape.648, %reshape.646), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=83} + %mul.2853 = f32[1,32,1,1,41]{4,3,2,1,0} multiply(%dot_general.966, %broadcast.36), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/vmap()/mul" stack_frame_id=83} + %vmap_jit__where__.87 = f32[1,1,32,1,41]{4,3,2,1,0} call(%and.271, %mul.2853, %constant.114), to_apply=%_where_3.137, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/vmap(jit(_where))" stack_frame_id=83} + %reduce_max.472 = f32[1,1,32,1]{3,2,1,0} reduce(%vmap_jit__where__.87, %constant.113), dimensions={4}, to_apply=%region_227.235, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/vmap()/reduce_max" stack_frame_id=83} + %max.96 = f32[1,1,32,1]{3,2,1,0} maximum(%reduce_max.472, %broadcast.35), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/vmap()/max" stack_frame_id=83} + %broadcast_in_dim.1383 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%max.96), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %sub.367 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1383), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/vmap()/sub" stack_frame_id=83} + %sub.368 = f32[1,1,32,1]{3,2,1,0} reshape(%sub.367), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/vmap()/sub" stack_frame_id=83} + %sub.369 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%sub.368), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/vmap()/sub" stack_frame_id=83} + %sub.370 = f32[1,1,32,1,41]{4,3,2,1,0} subtract(%vmap_jit__where__.87, %sub.369), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/vmap()/sub" stack_frame_id=83} + %exp.92 = f32[1,1,32,1,41]{4,3,2,1,0} exponential(%sub.370), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/vmap()/exp" stack_frame_id=83} + %reduce_sum.1424 = f32[1,1,32,1]{3,2,1,0} reduce(%exp.92, %constant.123), dimensions={4}, to_apply=%region_228.236, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/vmap()/reduce_sum" stack_frame_id=83} + %broadcast_in_dim.1384 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%reduce_sum.1424), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %div.888 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1384), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/vmap()/div" stack_frame_id=83} + %div.889 = f32[1,1,32,1]{3,2,1,0} reshape(%div.888), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/vmap()/div" stack_frame_id=83} + %div.890 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%div.889), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/vmap()/div" stack_frame_id=83} + %div.891 = f32[1,1,32,1,41]{4,3,2,1,0} divide(%exp.92, %div.890), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/vmap()/div" stack_frame_id=83} + %convert_element_type.1586 = bf16[1,1,32,1,41]{4,3,2,1,0} convert(%div.891), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/vmap()/convert_element_type" stack_frame_id=83} + %dot_general.967 = bf16[1,32,128,1,1]{4,3,2,1,0} dot(%reshape.647, %convert_element_type.1586), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=83} + %transpose.91 = bf16[1,1,32,128]{3,2,1,0} reshape(%dot_general.967), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/reshape;jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/vmap()/transpose;jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %dot_general.968 = bf16[1,1,4096]{2,1,0} dot(%transpose.91, %arg_tuple.5#228), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=354} + %add.1021 = bf16[1,1,4096]{2,1,0} add(%dot_general.968, %add.1010), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/add" stack_frame_id=355} + %convert_element_type.1587 = f32[1,1,4096]{2,1,0} convert(%add.1021), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/feedforward_layernorm/convert_element_type" stack_frame_id=359} + %pow.354 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1587, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/feedforward_layernorm/pow" stack_frame_id=362} + %reduce_sum.1425 = f32[1,1]{1,0} reduce(%pow.354, %constant.123), dimensions={2}, to_apply=%region_229.237, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/feedforward_layernorm/reduce_sum" stack_frame_id=92} + %broadcast_in_dim.1385 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1425), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/feedforward_layernorm/broadcast_in_dim" stack_frame_id=92} + %div.892 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1385, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/feedforward_layernorm/div" stack_frame_id=92} + %add.1022 = f32[1,1,1]{2,1,0} add(%div.892, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/feedforward_layernorm/add" stack_frame_id=363} + %rsqrt.177 = f32[1,1,1]{2,1,0} rsqrt(%add.1022), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/feedforward_layernorm/rsqrt" stack_frame_id=366} + %mul.2854 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.177), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2855 = f32[1,1]{1,0} reshape(%mul.2854), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2856 = f32[1,1,4096]{2,1,0} broadcast(%mul.2855), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2857 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1587, %mul.2856), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/feedforward_layernorm/mul" stack_frame_id=367} + %convert_element_type.1588 = f32[4096]{0} convert(%arg_tuple.5#229), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/feedforward_layernorm/convert_element_type" stack_frame_id=368} + %broadcast_in_dim.1386 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1588), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/feedforward_layernorm/broadcast_in_dim" stack_frame_id=368} + %mul.2858 = f32[1,1,4096]{2,1,0} multiply(%mul.2857, %broadcast_in_dim.1386), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/feedforward_layernorm/mul" stack_frame_id=368} + %convert_element_type.1589 = bf16[1,1,4096]{2,1,0} convert(%mul.2858), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/feedforward_layernorm/convert_element_type" stack_frame_id=372} + %dot_general.970 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1589, %arg_tuple.5#231), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/feedforward_intermediate_dense/dot_general" stack_frame_id=401} + %dot_general.969 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1589, %arg_tuple.5#230), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/feedforward_gate_dense/dot_general" stack_frame_id=381} + %convert_element_type.1590 = f32[1,1,14336]{2,1,0} convert(%dot_general.969), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/convert_element_type" stack_frame_id=385} + %jit_silu_.87 = f32[1,1,14336]{2,1,0} call(%convert_element_type.1590), to_apply=%silu_4.141, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/jit(silu)" stack_frame_id=388} + %convert_element_type.1591 = bf16[1,1,14336]{2,1,0} convert(%jit_silu_.87), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/convert_element_type" stack_frame_id=392} + %mul.2859 = bf16[1,1,14336]{2,1,0} multiply(%dot_general.970, %convert_element_type.1591), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/mul" stack_frame_id=404} + %dot_general.971 = bf16[1,1,4096]{2,1,0} dot(%mul.2859, %arg_tuple.5#232), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/feedforward_output_dense/dot_general" stack_frame_id=413} + %add.1023 = bf16[1,1,4096]{2,1,0} add(%dot_general.971, %add.1021), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/add" stack_frame_id=414} + %convert_element_type.1595 = f32[1,1,4096]{2,1,0} convert(%add.1023), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention_layernorm/convert_element_type" stack_frame_id=164} + %pow.355 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1595, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention_layernorm/pow" stack_frame_id=167} + %reduce_sum.1426 = f32[1,1]{1,0} reduce(%pow.355, %constant.123), dimensions={2}, to_apply=%region_230.238, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention_layernorm/reduce_sum" stack_frame_id=73} + %broadcast_in_dim.1391 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1426), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention_layernorm/broadcast_in_dim" stack_frame_id=73} + %div.893 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1391, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention_layernorm/div" stack_frame_id=73} + %add.1027 = f32[1,1,1]{2,1,0} add(%div.893, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention_layernorm/add" stack_frame_id=168} + %rsqrt.178 = f32[1,1,1]{2,1,0} rsqrt(%add.1027), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention_layernorm/rsqrt" stack_frame_id=171} + %mul.2860 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.178), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2861 = f32[1,1]{1,0} reshape(%mul.2860), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2862 = f32[1,1,4096]{2,1,0} broadcast(%mul.2861), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2863 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1595, %mul.2862), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention_layernorm/mul" stack_frame_id=172} + %convert_element_type.1596 = f32[4096]{0} convert(%arg_tuple.5#233), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention_layernorm/convert_element_type" stack_frame_id=173} + %broadcast_in_dim.1392 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1596), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention_layernorm/broadcast_in_dim" stack_frame_id=173} + %mul.2864 = f32[1,1,4096]{2,1,0} multiply(%mul.2863, %broadcast_in_dim.1392), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention_layernorm/mul" stack_frame_id=173} + %convert_element_type.1597 = bf16[1,1,4096]{2,1,0} convert(%mul.2864), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention_layernorm/convert_element_type" stack_frame_id=177} + %dot_general.975 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1597, %arg_tuple.5#236), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=265} + %lt.581 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/lt" stack_frame_id=329} + %add.1033 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/add" stack_frame_id=329} + %select_n.142 = s32[] select(%lt.581, %add.1033, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/select_n" stack_frame_id=329} + %dynamic_update_slice.116 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.177, %dot_general.975, %constant.127, %select_n.142, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/dynamic_update_slice" stack_frame_id=329} + %broadcast_in_dim.1394 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.116), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/broadcast_in_dim" stack_frame_id=338} + %reshape.654 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1394), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/reshape" stack_frame_id=338} + %convert_element_type.1592 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/convert_element_type" stack_frame_id=139} + %add.1024 = f32[1,1]{1,0} reshape(%convert_element_type.1592), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_25/add"} + %ge.453 = f32[1,1]{1,0} broadcast(%add.1024), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/ge" stack_frame_id=143} + %ge.454 = f32[1]{0} reshape(%ge.453), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/ge" stack_frame_id=143} + %ge.455 = f32[1,41]{1,0} broadcast(%ge.454), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/ge" stack_frame_id=143} + %iota.482 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/iota" stack_frame_id=142} + %broadcast_in_dim.1387 = f32[1,41]{1,0} reshape(%iota.482), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/broadcast_in_dim" stack_frame_id=143} + %ge.456 = pred[1,41]{1,0} compare(%ge.455, %broadcast_in_dim.1387), direction=GE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/ge" stack_frame_id=143} + %broadcast_in_dim.1388 = pred[1,1,41]{2,1,0} reshape(%ge.456), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/broadcast_in_dim" stack_frame_id=146} + %convert_element_type.1594 = s32[1,1,41]{2,1,0} convert(%broadcast_in_dim.1388), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/convert_element_type" stack_frame_id=160} + %add.1025 = s32[1,1]{1,0} reshape(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/add" stack_frame_id=147} + %lt.576 = s32[1,1]{1,0} broadcast(%add.1025), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/lt" stack_frame_id=152} + %lt.577 = s32[1]{0} reshape(%lt.576), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/lt" stack_frame_id=152} + %lt.578 = s32[1,41]{1,0} broadcast(%lt.577), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/lt" stack_frame_id=152} + %iota.483 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/iota" stack_frame_id=150} + %broadcast_in_dim.1389 = s32[1,41]{1,0} reshape(%iota.483), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/broadcast_in_dim" stack_frame_id=148} + %add.1026 = s32[1,41]{1,0} add(%broadcast_in_dim.1389, %broadcast.43), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/add" stack_frame_id=151} + %lt.579 = pred[1,41]{1,0} compare(%lt.578, %add.1026), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/lt" stack_frame_id=152} + %convert_element_type.1593 = s32[1,41]{1,0} convert(%lt.579), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/convert_element_type" stack_frame_id=156} + %broadcast_in_dim.1390 = s32[1,1,41]{2,1,0} reshape(%convert_element_type.1593), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/broadcast_in_dim" stack_frame_id=160} + %min.88 = s32[1,1,41]{2,1,0} minimum(%convert_element_type.1594, %broadcast_in_dim.1390), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/min" stack_frame_id=160} + %broadcast_in_dim.1395 = s32[1,1,1,41]{3,2,1,0} reshape(%min.88), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/broadcast_in_dim" stack_frame_id=341} + %convert_element_type.1604 = pred[1,1,1,41]{3,2,1,0} compare(%broadcast_in_dim.1395, %broadcast.37), direction=NE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/convert_element_type" stack_frame_id=345} + %broadcast_in_dim.1396 = pred[1,1,1,1,41]{4,3,2,1,0} reshape(%convert_element_type.1604), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/broadcast_in_dim" stack_frame_id=83} + %and.272 = pred[1,1,1,1,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1396), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/vmap()/and" stack_frame_id=83} + %and.273 = pred[1,1,1,41]{3,2,1,0} reshape(%and.272), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/vmap()/and" stack_frame_id=83} + %and.274 = pred[1,1,32,1,41]{4,3,2,1,0} broadcast(%and.273), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/vmap()/and" stack_frame_id=83} + %dot_general.972 = bf16[1,1,32,128]{3,2,1,0} dot(%convert_element_type.1597, %arg_tuple.5#234), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=186} + %convert_element_type.1598 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/convert_element_type" stack_frame_id=208} + %add.1028 = f32[1,1]{1,0} reshape(%convert_element_type.1598), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/add"} + %iota.484 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/iota" stack_frame_id=197} + %mul.2865 = f32[64]{0} multiply(%iota.484, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=197} + %div.894 = f32[64]{0} divide(%mul.2865, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/div" stack_frame_id=198} + %neg.362 = f32[64]{0} negate(%div.894), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/neg" stack_frame_id=199} + %pow.356 = f32[64]{0} power(%broadcast.39, %neg.362), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/pow" stack_frame_id=202} + %div.895 = f32[64]{0} divide(%pow.356, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/div" stack_frame_id=203} + %dot_general.973 = f32[1,1,64]{2,1,0} dot(%add.1028, %div.895), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/bi,j->bij/dot_general" stack_frame_id=211} + %stack.1395 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.973), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/stack" stack_frame_id=214} + %stack.1396 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.973), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/stack" stack_frame_id=214} + %stack.1397 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1395, %stack.1396), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/stack" stack_frame_id=214} + %reshape.649 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1397), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/reshape"} + %cos.177 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.649), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/cos" stack_frame_id=218} + %convert_element_type.1599 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.177), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/convert_element_type" stack_frame_id=222} + %mul.2866 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1599), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=242} + %mul.2867 = bf16[1,1,128]{2,1,0} reshape(%mul.2866), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=242} + %mul.2868 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2867), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=242} + %mul.2869 = bf16[1,1,32,128]{3,2,1,0} multiply(%dot_general.972, %mul.2868), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=242} + %split.355 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.972), slice={[0:1], [0:1], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/split" stack_frame_id=234} + %neg.363 = bf16[1,1,32,64]{3,2,1,0} negate(%split.355), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/neg" stack_frame_id=235} + %stack.1398 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%neg.363), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/stack" stack_frame_id=238} + %split.354 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.972), slice={[0:1], [0:1], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/split" stack_frame_id=234} + %stack.1399 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%split.354), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/stack" stack_frame_id=238} + %stack.1400 = bf16[1,1,32,2,64]{4,3,2,1,0} concatenate(%stack.1398, %stack.1399), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/stack" stack_frame_id=238} + %reshape.650 = bf16[1,1,32,128]{3,2,1,0} reshape(%stack.1400), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/reshape" stack_frame_id=241} + %sin.177 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.649), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/sin" stack_frame_id=226} + %convert_element_type.1600 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.177), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/convert_element_type" stack_frame_id=230} + %mul.2870 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1600), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=243} + %mul.2871 = bf16[1,1,128]{2,1,0} reshape(%mul.2870), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=243} + %mul.2872 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2871), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=243} + %mul.2873 = bf16[1,1,32,128]{3,2,1,0} multiply(%reshape.650, %mul.2872), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=243} + %add.1029 = bf16[1,1,32,128]{3,2,1,0} add(%mul.2869, %mul.2873), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/add" stack_frame_id=244} + %reshape.655 = bf16[1,1,32,1,128]{4,3,2,1,0} reshape(%add.1029), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/reshape" stack_frame_id=83} + %slice.172 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.175), slice={[0:1], [0:1], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/slice" stack_frame_id=245} + %squeeze.176 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.172), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/squeeze" stack_frame_id=245} + %dot_general.974 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1597, %arg_tuple.5#235), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=256} + %convert_element_type.1601 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/convert_element_type" stack_frame_id=287} + %add.1030 = f32[1,1]{1,0} reshape(%convert_element_type.1601), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/add"} + %iota.485 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/iota" stack_frame_id=276} + %mul.2874 = f32[64]{0} multiply(%iota.485, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=276} + %div.896 = f32[64]{0} divide(%mul.2874, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/div" stack_frame_id=277} + %neg.364 = f32[64]{0} negate(%div.896), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/neg" stack_frame_id=278} + %pow.357 = f32[64]{0} power(%broadcast.39, %neg.364), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/pow" stack_frame_id=281} + %div.897 = f32[64]{0} divide(%pow.357, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/div" stack_frame_id=282} + %dot_general.976 = f32[1,1,64]{2,1,0} dot(%add.1030, %div.897), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/bi,j->bij/dot_general" stack_frame_id=290} + %stack.1401 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.976), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/stack" stack_frame_id=293} + %stack.1402 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.976), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/stack" stack_frame_id=293} + %stack.1403 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1401, %stack.1402), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/stack" stack_frame_id=293} + %reshape.651 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1403), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/reshape"} + %cos.178 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.651), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/cos" stack_frame_id=297} + %convert_element_type.1602 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.178), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/convert_element_type" stack_frame_id=301} + %mul.2875 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1602), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=321} + %mul.2876 = bf16[1,1,128]{2,1,0} reshape(%mul.2875), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=321} + %mul.2877 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2876), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=321} + %mul.2878 = bf16[1,1,8,128]{3,2,1,0} multiply(%dot_general.974, %mul.2877), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=321} + %split.357 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.974), slice={[0:1], [0:1], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/split" stack_frame_id=313} + %neg.365 = bf16[1,1,8,64]{3,2,1,0} negate(%split.357), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/neg" stack_frame_id=314} + %stack.1404 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%neg.365), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/stack" stack_frame_id=317} + %split.356 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.974), slice={[0:1], [0:1], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/split" stack_frame_id=313} + %stack.1405 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%split.356), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/stack" stack_frame_id=317} + %stack.1406 = bf16[1,1,8,2,64]{4,3,2,1,0} concatenate(%stack.1404, %stack.1405), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/stack" stack_frame_id=317} + %reshape.652 = bf16[1,1,8,128]{3,2,1,0} reshape(%stack.1406), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/reshape" stack_frame_id=320} + %sin.178 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.651), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/sin" stack_frame_id=305} + %convert_element_type.1603 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.178), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/convert_element_type" stack_frame_id=309} + %mul.2879 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1603), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=322} + %mul.2880 = bf16[1,1,128]{2,1,0} reshape(%mul.2879), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=322} + %mul.2881 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2880), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=322} + %mul.2882 = bf16[1,1,8,128]{3,2,1,0} multiply(%reshape.652, %mul.2881), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=322} + %add.1031 = bf16[1,1,8,128]{3,2,1,0} add(%mul.2878, %mul.2882), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/rotary_embedding_25/add" stack_frame_id=323} + %lt.580 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/lt" stack_frame_id=326} + %add.1032 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/add" stack_frame_id=326} + %select_n.141 = s32[] select(%lt.580, %add.1032, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/select_n" stack_frame_id=326} + %dynamic_update_slice.115 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.176, %add.1031, %constant.127, %select_n.141, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/dynamic_update_slice" stack_frame_id=326} + %broadcast_in_dim.1393 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.115), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/broadcast_in_dim" stack_frame_id=335} + %reshape.653 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1393), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/reshape" stack_frame_id=335} + %dot_general.977 = f32[1,32,1,1,41]{4,3,2,1,0} dot(%reshape.655, %reshape.653), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=83} + %mul.2883 = f32[1,32,1,1,41]{4,3,2,1,0} multiply(%dot_general.977, %broadcast.36), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/vmap()/mul" stack_frame_id=83} + %vmap_jit__where__.88 = f32[1,1,32,1,41]{4,3,2,1,0} call(%and.274, %mul.2883, %constant.114), to_apply=%_where_3.137, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/vmap(jit(_where))" stack_frame_id=83} + %reduce_max.473 = f32[1,1,32,1]{3,2,1,0} reduce(%vmap_jit__where__.88, %constant.113), dimensions={4}, to_apply=%region_231.239, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/vmap()/reduce_max" stack_frame_id=83} + %max.97 = f32[1,1,32,1]{3,2,1,0} maximum(%reduce_max.473, %broadcast.35), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/vmap()/max" stack_frame_id=83} + %broadcast_in_dim.1397 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%max.97), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %sub.371 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1397), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/vmap()/sub" stack_frame_id=83} + %sub.372 = f32[1,1,32,1]{3,2,1,0} reshape(%sub.371), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/vmap()/sub" stack_frame_id=83} + %sub.373 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%sub.372), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/vmap()/sub" stack_frame_id=83} + %sub.374 = f32[1,1,32,1,41]{4,3,2,1,0} subtract(%vmap_jit__where__.88, %sub.373), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/vmap()/sub" stack_frame_id=83} + %exp.93 = f32[1,1,32,1,41]{4,3,2,1,0} exponential(%sub.374), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/vmap()/exp" stack_frame_id=83} + %reduce_sum.1427 = f32[1,1,32,1]{3,2,1,0} reduce(%exp.93, %constant.123), dimensions={4}, to_apply=%region_232.240, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/vmap()/reduce_sum" stack_frame_id=83} + %broadcast_in_dim.1398 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%reduce_sum.1427), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %div.898 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1398), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/vmap()/div" stack_frame_id=83} + %div.899 = f32[1,1,32,1]{3,2,1,0} reshape(%div.898), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/vmap()/div" stack_frame_id=83} + %div.900 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%div.899), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/vmap()/div" stack_frame_id=83} + %div.901 = f32[1,1,32,1,41]{4,3,2,1,0} divide(%exp.93, %div.900), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/vmap()/div" stack_frame_id=83} + %convert_element_type.1605 = bf16[1,1,32,1,41]{4,3,2,1,0} convert(%div.901), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/vmap()/convert_element_type" stack_frame_id=83} + %dot_general.978 = bf16[1,32,128,1,1]{4,3,2,1,0} dot(%reshape.654, %convert_element_type.1605), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=83} + %transpose.92 = bf16[1,1,32,128]{3,2,1,0} reshape(%dot_general.978), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/reshape;jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/vmap()/transpose;jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %dot_general.979 = bf16[1,1,4096]{2,1,0} dot(%transpose.92, %arg_tuple.5#237), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=354} + %add.1034 = bf16[1,1,4096]{2,1,0} add(%dot_general.979, %add.1023), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/add" stack_frame_id=355} + %convert_element_type.1606 = f32[1,1,4096]{2,1,0} convert(%add.1034), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/feedforward_layernorm/convert_element_type" stack_frame_id=359} + %pow.358 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1606, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/feedforward_layernorm/pow" stack_frame_id=362} + %reduce_sum.1428 = f32[1,1]{1,0} reduce(%pow.358, %constant.123), dimensions={2}, to_apply=%region_233.241, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/feedforward_layernorm/reduce_sum" stack_frame_id=92} + %broadcast_in_dim.1399 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1428), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/feedforward_layernorm/broadcast_in_dim" stack_frame_id=92} + %div.902 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1399, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/feedforward_layernorm/div" stack_frame_id=92} + %add.1035 = f32[1,1,1]{2,1,0} add(%div.902, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/feedforward_layernorm/add" stack_frame_id=363} + %rsqrt.179 = f32[1,1,1]{2,1,0} rsqrt(%add.1035), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/feedforward_layernorm/rsqrt" stack_frame_id=366} + %mul.2884 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.179), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2885 = f32[1,1]{1,0} reshape(%mul.2884), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2886 = f32[1,1,4096]{2,1,0} broadcast(%mul.2885), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2887 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1606, %mul.2886), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/feedforward_layernorm/mul" stack_frame_id=367} + %convert_element_type.1607 = f32[4096]{0} convert(%arg_tuple.5#238), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/feedforward_layernorm/convert_element_type" stack_frame_id=368} + %broadcast_in_dim.1400 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1607), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/feedforward_layernorm/broadcast_in_dim" stack_frame_id=368} + %mul.2888 = f32[1,1,4096]{2,1,0} multiply(%mul.2887, %broadcast_in_dim.1400), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/feedforward_layernorm/mul" stack_frame_id=368} + %convert_element_type.1608 = bf16[1,1,4096]{2,1,0} convert(%mul.2888), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/feedforward_layernorm/convert_element_type" stack_frame_id=372} + %dot_general.981 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1608, %arg_tuple.5#240), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/feedforward_intermediate_dense/dot_general" stack_frame_id=401} + %dot_general.980 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1608, %arg_tuple.5#239), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/feedforward_gate_dense/dot_general" stack_frame_id=381} + %convert_element_type.1609 = f32[1,1,14336]{2,1,0} convert(%dot_general.980), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/convert_element_type" stack_frame_id=385} + %jit_silu_.88 = f32[1,1,14336]{2,1,0} call(%convert_element_type.1609), to_apply=%silu_4.141, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/jit(silu)" stack_frame_id=388} + %convert_element_type.1610 = bf16[1,1,14336]{2,1,0} convert(%jit_silu_.88), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/convert_element_type" stack_frame_id=392} + %mul.2889 = bf16[1,1,14336]{2,1,0} multiply(%dot_general.981, %convert_element_type.1610), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/mul" stack_frame_id=404} + %dot_general.982 = bf16[1,1,4096]{2,1,0} dot(%mul.2889, %arg_tuple.5#241), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/feedforward_output_dense/dot_general" stack_frame_id=413} + %add.1036 = bf16[1,1,4096]{2,1,0} add(%dot_general.982, %add.1034), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/add" stack_frame_id=414} + %convert_element_type.1614 = f32[1,1,4096]{2,1,0} convert(%add.1036), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention_layernorm/convert_element_type" stack_frame_id=164} + %pow.359 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1614, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention_layernorm/pow" stack_frame_id=167} + %reduce_sum.1429 = f32[1,1]{1,0} reduce(%pow.359, %constant.123), dimensions={2}, to_apply=%region_234.242, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention_layernorm/reduce_sum" stack_frame_id=73} + %broadcast_in_dim.1405 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1429), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention_layernorm/broadcast_in_dim" stack_frame_id=73} + %div.903 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1405, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention_layernorm/div" stack_frame_id=73} + %add.1040 = f32[1,1,1]{2,1,0} add(%div.903, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention_layernorm/add" stack_frame_id=168} + %rsqrt.180 = f32[1,1,1]{2,1,0} rsqrt(%add.1040), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention_layernorm/rsqrt" stack_frame_id=171} + %mul.2890 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.180), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2891 = f32[1,1]{1,0} reshape(%mul.2890), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2892 = f32[1,1,4096]{2,1,0} broadcast(%mul.2891), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2893 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1614, %mul.2892), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention_layernorm/mul" stack_frame_id=172} + %convert_element_type.1615 = f32[4096]{0} convert(%arg_tuple.5#242), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention_layernorm/convert_element_type" stack_frame_id=173} + %broadcast_in_dim.1406 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1615), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention_layernorm/broadcast_in_dim" stack_frame_id=173} + %mul.2894 = f32[1,1,4096]{2,1,0} multiply(%mul.2893, %broadcast_in_dim.1406), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention_layernorm/mul" stack_frame_id=173} + %convert_element_type.1616 = bf16[1,1,4096]{2,1,0} convert(%mul.2894), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention_layernorm/convert_element_type" stack_frame_id=177} + %dot_general.986 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1616, %arg_tuple.5#245), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=265} + %lt.587 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/lt" stack_frame_id=329} + %add.1046 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/add" stack_frame_id=329} + %select_n.144 = s32[] select(%lt.587, %add.1046, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/select_n" stack_frame_id=329} + %dynamic_update_slice.118 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.180, %dot_general.986, %constant.127, %select_n.144, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/dynamic_update_slice" stack_frame_id=329} + %broadcast_in_dim.1408 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.118), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/broadcast_in_dim" stack_frame_id=338} + %reshape.661 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1408), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/reshape" stack_frame_id=338} + %convert_element_type.1611 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/convert_element_type" stack_frame_id=139} + %add.1037 = f32[1,1]{1,0} reshape(%convert_element_type.1611), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_26/add"} + %ge.457 = f32[1,1]{1,0} broadcast(%add.1037), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/ge" stack_frame_id=143} + %ge.458 = f32[1]{0} reshape(%ge.457), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/ge" stack_frame_id=143} + %ge.459 = f32[1,41]{1,0} broadcast(%ge.458), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/ge" stack_frame_id=143} + %iota.486 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/iota" stack_frame_id=142} + %broadcast_in_dim.1401 = f32[1,41]{1,0} reshape(%iota.486), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/broadcast_in_dim" stack_frame_id=143} + %ge.460 = pred[1,41]{1,0} compare(%ge.459, %broadcast_in_dim.1401), direction=GE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/ge" stack_frame_id=143} + %broadcast_in_dim.1402 = pred[1,1,41]{2,1,0} reshape(%ge.460), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/broadcast_in_dim" stack_frame_id=146} + %convert_element_type.1613 = s32[1,1,41]{2,1,0} convert(%broadcast_in_dim.1402), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/convert_element_type" stack_frame_id=160} + %add.1038 = s32[1,1]{1,0} reshape(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/add" stack_frame_id=147} + %lt.582 = s32[1,1]{1,0} broadcast(%add.1038), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/lt" stack_frame_id=152} + %lt.583 = s32[1]{0} reshape(%lt.582), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/lt" stack_frame_id=152} + %lt.584 = s32[1,41]{1,0} broadcast(%lt.583), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/lt" stack_frame_id=152} + %iota.487 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/iota" stack_frame_id=150} + %broadcast_in_dim.1403 = s32[1,41]{1,0} reshape(%iota.487), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/broadcast_in_dim" stack_frame_id=148} + %add.1039 = s32[1,41]{1,0} add(%broadcast_in_dim.1403, %broadcast.43), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/add" stack_frame_id=151} + %lt.585 = pred[1,41]{1,0} compare(%lt.584, %add.1039), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/lt" stack_frame_id=152} + %convert_element_type.1612 = s32[1,41]{1,0} convert(%lt.585), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/convert_element_type" stack_frame_id=156} + %broadcast_in_dim.1404 = s32[1,1,41]{2,1,0} reshape(%convert_element_type.1612), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/broadcast_in_dim" stack_frame_id=160} + %min.89 = s32[1,1,41]{2,1,0} minimum(%convert_element_type.1613, %broadcast_in_dim.1404), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/min" stack_frame_id=160} + %broadcast_in_dim.1409 = s32[1,1,1,41]{3,2,1,0} reshape(%min.89), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/broadcast_in_dim" stack_frame_id=341} + %convert_element_type.1623 = pred[1,1,1,41]{3,2,1,0} compare(%broadcast_in_dim.1409, %broadcast.37), direction=NE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/convert_element_type" stack_frame_id=345} + %broadcast_in_dim.1410 = pred[1,1,1,1,41]{4,3,2,1,0} reshape(%convert_element_type.1623), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/broadcast_in_dim" stack_frame_id=83} + %and.275 = pred[1,1,1,1,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1410), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/vmap()/and" stack_frame_id=83} + %and.276 = pred[1,1,1,41]{3,2,1,0} reshape(%and.275), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/vmap()/and" stack_frame_id=83} + %and.277 = pred[1,1,32,1,41]{4,3,2,1,0} broadcast(%and.276), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/vmap()/and" stack_frame_id=83} + %dot_general.983 = bf16[1,1,32,128]{3,2,1,0} dot(%convert_element_type.1616, %arg_tuple.5#243), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=186} + %convert_element_type.1617 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/convert_element_type" stack_frame_id=208} + %add.1041 = f32[1,1]{1,0} reshape(%convert_element_type.1617), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/add"} + %iota.488 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/iota" stack_frame_id=197} + %mul.2895 = f32[64]{0} multiply(%iota.488, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=197} + %div.904 = f32[64]{0} divide(%mul.2895, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/div" stack_frame_id=198} + %neg.366 = f32[64]{0} negate(%div.904), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/neg" stack_frame_id=199} + %pow.360 = f32[64]{0} power(%broadcast.39, %neg.366), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/pow" stack_frame_id=202} + %div.905 = f32[64]{0} divide(%pow.360, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/div" stack_frame_id=203} + %dot_general.984 = f32[1,1,64]{2,1,0} dot(%add.1041, %div.905), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/bi,j->bij/dot_general" stack_frame_id=211} + %stack.1410 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.984), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/stack" stack_frame_id=214} + %stack.1411 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.984), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/stack" stack_frame_id=214} + %stack.1412 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1410, %stack.1411), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/stack" stack_frame_id=214} + %reshape.656 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1412), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/reshape"} + %cos.179 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.656), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/cos" stack_frame_id=218} + %convert_element_type.1618 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.179), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/convert_element_type" stack_frame_id=222} + %mul.2896 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1618), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=242} + %mul.2897 = bf16[1,1,128]{2,1,0} reshape(%mul.2896), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=242} + %mul.2898 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2897), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=242} + %mul.2899 = bf16[1,1,32,128]{3,2,1,0} multiply(%dot_general.983, %mul.2898), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=242} + %split.359 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.983), slice={[0:1], [0:1], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/split" stack_frame_id=234} + %neg.367 = bf16[1,1,32,64]{3,2,1,0} negate(%split.359), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/neg" stack_frame_id=235} + %stack.1413 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%neg.367), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/stack" stack_frame_id=238} + %split.358 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.983), slice={[0:1], [0:1], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/split" stack_frame_id=234} + %stack.1414 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%split.358), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/stack" stack_frame_id=238} + %stack.1415 = bf16[1,1,32,2,64]{4,3,2,1,0} concatenate(%stack.1413, %stack.1414), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/stack" stack_frame_id=238} + %reshape.657 = bf16[1,1,32,128]{3,2,1,0} reshape(%stack.1415), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/reshape" stack_frame_id=241} + %sin.179 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.656), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/sin" stack_frame_id=226} + %convert_element_type.1619 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.179), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/convert_element_type" stack_frame_id=230} + %mul.2900 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1619), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=243} + %mul.2901 = bf16[1,1,128]{2,1,0} reshape(%mul.2900), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=243} + %mul.2902 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2901), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=243} + %mul.2903 = bf16[1,1,32,128]{3,2,1,0} multiply(%reshape.657, %mul.2902), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=243} + %add.1042 = bf16[1,1,32,128]{3,2,1,0} add(%mul.2899, %mul.2903), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/add" stack_frame_id=244} + %reshape.662 = bf16[1,1,32,1,128]{4,3,2,1,0} reshape(%add.1042), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/reshape" stack_frame_id=83} + %slice.175 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.178), slice={[0:1], [0:1], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/slice" stack_frame_id=245} + %squeeze.179 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.175), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/squeeze" stack_frame_id=245} + %dot_general.985 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1616, %arg_tuple.5#244), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=256} + %convert_element_type.1620 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/convert_element_type" stack_frame_id=287} + %add.1043 = f32[1,1]{1,0} reshape(%convert_element_type.1620), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/add"} + %iota.489 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/iota" stack_frame_id=276} + %mul.2904 = f32[64]{0} multiply(%iota.489, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=276} + %div.906 = f32[64]{0} divide(%mul.2904, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/div" stack_frame_id=277} + %neg.368 = f32[64]{0} negate(%div.906), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/neg" stack_frame_id=278} + %pow.361 = f32[64]{0} power(%broadcast.39, %neg.368), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/pow" stack_frame_id=281} + %div.907 = f32[64]{0} divide(%pow.361, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/div" stack_frame_id=282} + %dot_general.987 = f32[1,1,64]{2,1,0} dot(%add.1043, %div.907), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/bi,j->bij/dot_general" stack_frame_id=290} + %stack.1416 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.987), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/stack" stack_frame_id=293} + %stack.1417 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.987), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/stack" stack_frame_id=293} + %stack.1418 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1416, %stack.1417), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/stack" stack_frame_id=293} + %reshape.658 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1418), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/reshape"} + %cos.180 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.658), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/cos" stack_frame_id=297} + %convert_element_type.1621 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.180), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/convert_element_type" stack_frame_id=301} + %mul.2905 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1621), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=321} + %mul.2906 = bf16[1,1,128]{2,1,0} reshape(%mul.2905), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=321} + %mul.2907 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2906), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=321} + %mul.2908 = bf16[1,1,8,128]{3,2,1,0} multiply(%dot_general.985, %mul.2907), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=321} + %split.361 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.985), slice={[0:1], [0:1], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/split" stack_frame_id=313} + %neg.369 = bf16[1,1,8,64]{3,2,1,0} negate(%split.361), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/neg" stack_frame_id=314} + %stack.1419 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%neg.369), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/stack" stack_frame_id=317} + %split.360 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.985), slice={[0:1], [0:1], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/split" stack_frame_id=313} + %stack.1420 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%split.360), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/stack" stack_frame_id=317} + %stack.1421 = bf16[1,1,8,2,64]{4,3,2,1,0} concatenate(%stack.1419, %stack.1420), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/stack" stack_frame_id=317} + %reshape.659 = bf16[1,1,8,128]{3,2,1,0} reshape(%stack.1421), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/reshape" stack_frame_id=320} + %sin.180 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.658), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/sin" stack_frame_id=305} + %convert_element_type.1622 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.180), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/convert_element_type" stack_frame_id=309} + %mul.2909 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1622), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=322} + %mul.2910 = bf16[1,1,128]{2,1,0} reshape(%mul.2909), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=322} + %mul.2911 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2910), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=322} + %mul.2912 = bf16[1,1,8,128]{3,2,1,0} multiply(%reshape.659, %mul.2911), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=322} + %add.1044 = bf16[1,1,8,128]{3,2,1,0} add(%mul.2908, %mul.2912), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/rotary_embedding_26/add" stack_frame_id=323} + %lt.586 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/lt" stack_frame_id=326} + %add.1045 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/add" stack_frame_id=326} + %select_n.143 = s32[] select(%lt.586, %add.1045, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/select_n" stack_frame_id=326} + %dynamic_update_slice.117 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.179, %add.1044, %constant.127, %select_n.143, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/dynamic_update_slice" stack_frame_id=326} + %broadcast_in_dim.1407 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.117), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/broadcast_in_dim" stack_frame_id=335} + %reshape.660 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1407), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/reshape" stack_frame_id=335} + %dot_general.988 = f32[1,32,1,1,41]{4,3,2,1,0} dot(%reshape.662, %reshape.660), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=83} + %mul.2913 = f32[1,32,1,1,41]{4,3,2,1,0} multiply(%dot_general.988, %broadcast.36), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/vmap()/mul" stack_frame_id=83} + %vmap_jit__where__.89 = f32[1,1,32,1,41]{4,3,2,1,0} call(%and.277, %mul.2913, %constant.114), to_apply=%_where_3.137, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/vmap(jit(_where))" stack_frame_id=83} + %reduce_max.474 = f32[1,1,32,1]{3,2,1,0} reduce(%vmap_jit__where__.89, %constant.113), dimensions={4}, to_apply=%region_235.243, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/vmap()/reduce_max" stack_frame_id=83} + %max.98 = f32[1,1,32,1]{3,2,1,0} maximum(%reduce_max.474, %broadcast.35), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/vmap()/max" stack_frame_id=83} + %broadcast_in_dim.1411 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%max.98), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %sub.375 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1411), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/vmap()/sub" stack_frame_id=83} + %sub.376 = f32[1,1,32,1]{3,2,1,0} reshape(%sub.375), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/vmap()/sub" stack_frame_id=83} + %sub.377 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%sub.376), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/vmap()/sub" stack_frame_id=83} + %sub.378 = f32[1,1,32,1,41]{4,3,2,1,0} subtract(%vmap_jit__where__.89, %sub.377), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/vmap()/sub" stack_frame_id=83} + %exp.94 = f32[1,1,32,1,41]{4,3,2,1,0} exponential(%sub.378), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/vmap()/exp" stack_frame_id=83} + %reduce_sum.1430 = f32[1,1,32,1]{3,2,1,0} reduce(%exp.94, %constant.123), dimensions={4}, to_apply=%region_236.244, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/vmap()/reduce_sum" stack_frame_id=83} + %broadcast_in_dim.1412 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%reduce_sum.1430), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %div.908 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1412), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/vmap()/div" stack_frame_id=83} + %div.909 = f32[1,1,32,1]{3,2,1,0} reshape(%div.908), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/vmap()/div" stack_frame_id=83} + %div.910 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%div.909), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/vmap()/div" stack_frame_id=83} + %div.911 = f32[1,1,32,1,41]{4,3,2,1,0} divide(%exp.94, %div.910), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/vmap()/div" stack_frame_id=83} + %convert_element_type.1624 = bf16[1,1,32,1,41]{4,3,2,1,0} convert(%div.911), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/vmap()/convert_element_type" stack_frame_id=83} + %dot_general.989 = bf16[1,32,128,1,1]{4,3,2,1,0} dot(%reshape.661, %convert_element_type.1624), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=83} + %transpose.93 = bf16[1,1,32,128]{3,2,1,0} reshape(%dot_general.989), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/reshape;jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/vmap()/transpose;jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %dot_general.990 = bf16[1,1,4096]{2,1,0} dot(%transpose.93, %arg_tuple.5#246), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=354} + %add.1047 = bf16[1,1,4096]{2,1,0} add(%dot_general.990, %add.1036), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/add" stack_frame_id=355} + %convert_element_type.1625 = f32[1,1,4096]{2,1,0} convert(%add.1047), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/feedforward_layernorm/convert_element_type" stack_frame_id=359} + %pow.362 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1625, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/feedforward_layernorm/pow" stack_frame_id=362} + %reduce_sum.1431 = f32[1,1]{1,0} reduce(%pow.362, %constant.123), dimensions={2}, to_apply=%region_237.245, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/feedforward_layernorm/reduce_sum" stack_frame_id=92} + %broadcast_in_dim.1413 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1431), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/feedforward_layernorm/broadcast_in_dim" stack_frame_id=92} + %div.912 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1413, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/feedforward_layernorm/div" stack_frame_id=92} + %add.1048 = f32[1,1,1]{2,1,0} add(%div.912, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/feedforward_layernorm/add" stack_frame_id=363} + %rsqrt.181 = f32[1,1,1]{2,1,0} rsqrt(%add.1048), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/feedforward_layernorm/rsqrt" stack_frame_id=366} + %mul.2914 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.181), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2915 = f32[1,1]{1,0} reshape(%mul.2914), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2916 = f32[1,1,4096]{2,1,0} broadcast(%mul.2915), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2917 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1625, %mul.2916), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/feedforward_layernorm/mul" stack_frame_id=367} + %convert_element_type.1626 = f32[4096]{0} convert(%arg_tuple.5#247), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/feedforward_layernorm/convert_element_type" stack_frame_id=368} + %broadcast_in_dim.1414 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1626), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/feedforward_layernorm/broadcast_in_dim" stack_frame_id=368} + %mul.2918 = f32[1,1,4096]{2,1,0} multiply(%mul.2917, %broadcast_in_dim.1414), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/feedforward_layernorm/mul" stack_frame_id=368} + %convert_element_type.1627 = bf16[1,1,4096]{2,1,0} convert(%mul.2918), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/feedforward_layernorm/convert_element_type" stack_frame_id=372} + %dot_general.992 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1627, %arg_tuple.5#249), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/feedforward_intermediate_dense/dot_general" stack_frame_id=401} + %dot_general.991 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1627, %arg_tuple.5#248), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/feedforward_gate_dense/dot_general" stack_frame_id=381} + %convert_element_type.1628 = f32[1,1,14336]{2,1,0} convert(%dot_general.991), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/convert_element_type" stack_frame_id=385} + %jit_silu_.89 = f32[1,1,14336]{2,1,0} call(%convert_element_type.1628), to_apply=%silu_4.141, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/jit(silu)" stack_frame_id=388} + %convert_element_type.1629 = bf16[1,1,14336]{2,1,0} convert(%jit_silu_.89), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/convert_element_type" stack_frame_id=392} + %mul.2919 = bf16[1,1,14336]{2,1,0} multiply(%dot_general.992, %convert_element_type.1629), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/mul" stack_frame_id=404} + %dot_general.993 = bf16[1,1,4096]{2,1,0} dot(%mul.2919, %arg_tuple.5#250), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/feedforward_output_dense/dot_general" stack_frame_id=413} + %add.1049 = bf16[1,1,4096]{2,1,0} add(%dot_general.993, %add.1047), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/add" stack_frame_id=414} + %convert_element_type.1633 = f32[1,1,4096]{2,1,0} convert(%add.1049), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention_layernorm/convert_element_type" stack_frame_id=164} + %pow.363 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1633, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention_layernorm/pow" stack_frame_id=167} + %reduce_sum.1432 = f32[1,1]{1,0} reduce(%pow.363, %constant.123), dimensions={2}, to_apply=%region_238.246, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention_layernorm/reduce_sum" stack_frame_id=73} + %broadcast_in_dim.1419 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1432), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention_layernorm/broadcast_in_dim" stack_frame_id=73} + %div.913 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1419, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention_layernorm/div" stack_frame_id=73} + %add.1053 = f32[1,1,1]{2,1,0} add(%div.913, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention_layernorm/add" stack_frame_id=168} + %rsqrt.182 = f32[1,1,1]{2,1,0} rsqrt(%add.1053), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention_layernorm/rsqrt" stack_frame_id=171} + %mul.2920 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.182), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2921 = f32[1,1]{1,0} reshape(%mul.2920), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2922 = f32[1,1,4096]{2,1,0} broadcast(%mul.2921), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2923 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1633, %mul.2922), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention_layernorm/mul" stack_frame_id=172} + %convert_element_type.1634 = f32[4096]{0} convert(%arg_tuple.5#251), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention_layernorm/convert_element_type" stack_frame_id=173} + %broadcast_in_dim.1420 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1634), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention_layernorm/broadcast_in_dim" stack_frame_id=173} + %mul.2924 = f32[1,1,4096]{2,1,0} multiply(%mul.2923, %broadcast_in_dim.1420), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention_layernorm/mul" stack_frame_id=173} + %convert_element_type.1635 = bf16[1,1,4096]{2,1,0} convert(%mul.2924), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention_layernorm/convert_element_type" stack_frame_id=177} + %dot_general.997 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1635, %arg_tuple.5#254), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=265} + %lt.593 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/lt" stack_frame_id=329} + %add.1059 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/add" stack_frame_id=329} + %select_n.146 = s32[] select(%lt.593, %add.1059, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/select_n" stack_frame_id=329} + %dynamic_update_slice.120 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.183, %dot_general.997, %constant.127, %select_n.146, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/dynamic_update_slice" stack_frame_id=329} + %broadcast_in_dim.1422 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.120), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/broadcast_in_dim" stack_frame_id=338} + %reshape.668 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1422), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/reshape" stack_frame_id=338} + %convert_element_type.1630 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/convert_element_type" stack_frame_id=139} + %add.1050 = f32[1,1]{1,0} reshape(%convert_element_type.1630), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_27/add"} + %ge.461 = f32[1,1]{1,0} broadcast(%add.1050), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/ge" stack_frame_id=143} + %ge.462 = f32[1]{0} reshape(%ge.461), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/ge" stack_frame_id=143} + %ge.463 = f32[1,41]{1,0} broadcast(%ge.462), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/ge" stack_frame_id=143} + %iota.490 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/iota" stack_frame_id=142} + %broadcast_in_dim.1415 = f32[1,41]{1,0} reshape(%iota.490), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/broadcast_in_dim" stack_frame_id=143} + %ge.464 = pred[1,41]{1,0} compare(%ge.463, %broadcast_in_dim.1415), direction=GE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/ge" stack_frame_id=143} + %broadcast_in_dim.1416 = pred[1,1,41]{2,1,0} reshape(%ge.464), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/broadcast_in_dim" stack_frame_id=146} + %convert_element_type.1632 = s32[1,1,41]{2,1,0} convert(%broadcast_in_dim.1416), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/convert_element_type" stack_frame_id=160} + %add.1051 = s32[1,1]{1,0} reshape(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/add" stack_frame_id=147} + %lt.588 = s32[1,1]{1,0} broadcast(%add.1051), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/lt" stack_frame_id=152} + %lt.589 = s32[1]{0} reshape(%lt.588), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/lt" stack_frame_id=152} + %lt.590 = s32[1,41]{1,0} broadcast(%lt.589), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/lt" stack_frame_id=152} + %iota.491 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/iota" stack_frame_id=150} + %broadcast_in_dim.1417 = s32[1,41]{1,0} reshape(%iota.491), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/broadcast_in_dim" stack_frame_id=148} + %add.1052 = s32[1,41]{1,0} add(%broadcast_in_dim.1417, %broadcast.43), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/add" stack_frame_id=151} + %lt.591 = pred[1,41]{1,0} compare(%lt.590, %add.1052), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/lt" stack_frame_id=152} + %convert_element_type.1631 = s32[1,41]{1,0} convert(%lt.591), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/convert_element_type" stack_frame_id=156} + %broadcast_in_dim.1418 = s32[1,1,41]{2,1,0} reshape(%convert_element_type.1631), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/broadcast_in_dim" stack_frame_id=160} + %min.90 = s32[1,1,41]{2,1,0} minimum(%convert_element_type.1632, %broadcast_in_dim.1418), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/min" stack_frame_id=160} + %broadcast_in_dim.1423 = s32[1,1,1,41]{3,2,1,0} reshape(%min.90), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/broadcast_in_dim" stack_frame_id=341} + %convert_element_type.1642 = pred[1,1,1,41]{3,2,1,0} compare(%broadcast_in_dim.1423, %broadcast.37), direction=NE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/convert_element_type" stack_frame_id=345} + %broadcast_in_dim.1424 = pred[1,1,1,1,41]{4,3,2,1,0} reshape(%convert_element_type.1642), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/broadcast_in_dim" stack_frame_id=83} + %and.278 = pred[1,1,1,1,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1424), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/vmap()/and" stack_frame_id=83} + %and.279 = pred[1,1,1,41]{3,2,1,0} reshape(%and.278), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/vmap()/and" stack_frame_id=83} + %and.280 = pred[1,1,32,1,41]{4,3,2,1,0} broadcast(%and.279), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/vmap()/and" stack_frame_id=83} + %dot_general.994 = bf16[1,1,32,128]{3,2,1,0} dot(%convert_element_type.1635, %arg_tuple.5#252), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=186} + %convert_element_type.1636 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/convert_element_type" stack_frame_id=208} + %add.1054 = f32[1,1]{1,0} reshape(%convert_element_type.1636), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/add"} + %iota.492 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/iota" stack_frame_id=197} + %mul.2925 = f32[64]{0} multiply(%iota.492, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=197} + %div.914 = f32[64]{0} divide(%mul.2925, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/div" stack_frame_id=198} + %neg.370 = f32[64]{0} negate(%div.914), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/neg" stack_frame_id=199} + %pow.364 = f32[64]{0} power(%broadcast.39, %neg.370), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/pow" stack_frame_id=202} + %div.915 = f32[64]{0} divide(%pow.364, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/div" stack_frame_id=203} + %dot_general.995 = f32[1,1,64]{2,1,0} dot(%add.1054, %div.915), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/bi,j->bij/dot_general" stack_frame_id=211} + %stack.1425 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.995), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/stack" stack_frame_id=214} + %stack.1426 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.995), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/stack" stack_frame_id=214} + %stack.1427 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1425, %stack.1426), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/stack" stack_frame_id=214} + %reshape.663 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1427), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/reshape"} + %cos.181 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.663), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/cos" stack_frame_id=218} + %convert_element_type.1637 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.181), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/convert_element_type" stack_frame_id=222} + %mul.2926 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1637), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=242} + %mul.2927 = bf16[1,1,128]{2,1,0} reshape(%mul.2926), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=242} + %mul.2928 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2927), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=242} + %mul.2929 = bf16[1,1,32,128]{3,2,1,0} multiply(%dot_general.994, %mul.2928), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=242} + %split.363 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.994), slice={[0:1], [0:1], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/split" stack_frame_id=234} + %neg.371 = bf16[1,1,32,64]{3,2,1,0} negate(%split.363), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/neg" stack_frame_id=235} + %stack.1428 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%neg.371), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/stack" stack_frame_id=238} + %split.362 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.994), slice={[0:1], [0:1], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/split" stack_frame_id=234} + %stack.1429 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%split.362), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/stack" stack_frame_id=238} + %stack.1430 = bf16[1,1,32,2,64]{4,3,2,1,0} concatenate(%stack.1428, %stack.1429), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/stack" stack_frame_id=238} + %reshape.664 = bf16[1,1,32,128]{3,2,1,0} reshape(%stack.1430), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/reshape" stack_frame_id=241} + %sin.181 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.663), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/sin" stack_frame_id=226} + %convert_element_type.1638 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.181), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/convert_element_type" stack_frame_id=230} + %mul.2930 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1638), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=243} + %mul.2931 = bf16[1,1,128]{2,1,0} reshape(%mul.2930), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=243} + %mul.2932 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2931), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=243} + %mul.2933 = bf16[1,1,32,128]{3,2,1,0} multiply(%reshape.664, %mul.2932), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=243} + %add.1055 = bf16[1,1,32,128]{3,2,1,0} add(%mul.2929, %mul.2933), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/add" stack_frame_id=244} + %reshape.669 = bf16[1,1,32,1,128]{4,3,2,1,0} reshape(%add.1055), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/reshape" stack_frame_id=83} + %slice.178 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.181), slice={[0:1], [0:1], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/slice" stack_frame_id=245} + %squeeze.182 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.178), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/squeeze" stack_frame_id=245} + %dot_general.996 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1635, %arg_tuple.5#253), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=256} + %convert_element_type.1639 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/convert_element_type" stack_frame_id=287} + %add.1056 = f32[1,1]{1,0} reshape(%convert_element_type.1639), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/add"} + %iota.493 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/iota" stack_frame_id=276} + %mul.2934 = f32[64]{0} multiply(%iota.493, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=276} + %div.916 = f32[64]{0} divide(%mul.2934, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/div" stack_frame_id=277} + %neg.372 = f32[64]{0} negate(%div.916), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/neg" stack_frame_id=278} + %pow.365 = f32[64]{0} power(%broadcast.39, %neg.372), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/pow" stack_frame_id=281} + %div.917 = f32[64]{0} divide(%pow.365, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/div" stack_frame_id=282} + %dot_general.998 = f32[1,1,64]{2,1,0} dot(%add.1056, %div.917), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/bi,j->bij/dot_general" stack_frame_id=290} + %stack.1431 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.998), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/stack" stack_frame_id=293} + %stack.1432 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.998), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/stack" stack_frame_id=293} + %stack.1433 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1431, %stack.1432), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/stack" stack_frame_id=293} + %reshape.665 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1433), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/reshape"} + %cos.182 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.665), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/cos" stack_frame_id=297} + %convert_element_type.1640 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.182), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/convert_element_type" stack_frame_id=301} + %mul.2935 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1640), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=321} + %mul.2936 = bf16[1,1,128]{2,1,0} reshape(%mul.2935), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=321} + %mul.2937 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2936), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=321} + %mul.2938 = bf16[1,1,8,128]{3,2,1,0} multiply(%dot_general.996, %mul.2937), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=321} + %split.365 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.996), slice={[0:1], [0:1], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/split" stack_frame_id=313} + %neg.373 = bf16[1,1,8,64]{3,2,1,0} negate(%split.365), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/neg" stack_frame_id=314} + %stack.1434 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%neg.373), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/stack" stack_frame_id=317} + %split.364 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.996), slice={[0:1], [0:1], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/split" stack_frame_id=313} + %stack.1435 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%split.364), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/stack" stack_frame_id=317} + %stack.1436 = bf16[1,1,8,2,64]{4,3,2,1,0} concatenate(%stack.1434, %stack.1435), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/stack" stack_frame_id=317} + %reshape.666 = bf16[1,1,8,128]{3,2,1,0} reshape(%stack.1436), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/reshape" stack_frame_id=320} + %sin.182 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.665), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/sin" stack_frame_id=305} + %convert_element_type.1641 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.182), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/convert_element_type" stack_frame_id=309} + %mul.2939 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1641), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=322} + %mul.2940 = bf16[1,1,128]{2,1,0} reshape(%mul.2939), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=322} + %mul.2941 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2940), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=322} + %mul.2942 = bf16[1,1,8,128]{3,2,1,0} multiply(%reshape.666, %mul.2941), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=322} + %add.1057 = bf16[1,1,8,128]{3,2,1,0} add(%mul.2938, %mul.2942), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/rotary_embedding_27/add" stack_frame_id=323} + %lt.592 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/lt" stack_frame_id=326} + %add.1058 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/add" stack_frame_id=326} + %select_n.145 = s32[] select(%lt.592, %add.1058, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/select_n" stack_frame_id=326} + %dynamic_update_slice.119 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.182, %add.1057, %constant.127, %select_n.145, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/dynamic_update_slice" stack_frame_id=326} + %broadcast_in_dim.1421 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.119), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/broadcast_in_dim" stack_frame_id=335} + %reshape.667 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1421), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/reshape" stack_frame_id=335} + %dot_general.999 = f32[1,32,1,1,41]{4,3,2,1,0} dot(%reshape.669, %reshape.667), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=83} + %mul.2943 = f32[1,32,1,1,41]{4,3,2,1,0} multiply(%dot_general.999, %broadcast.36), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/vmap()/mul" stack_frame_id=83} + %vmap_jit__where__.90 = f32[1,1,32,1,41]{4,3,2,1,0} call(%and.280, %mul.2943, %constant.114), to_apply=%_where_3.137, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/vmap(jit(_where))" stack_frame_id=83} + %reduce_max.475 = f32[1,1,32,1]{3,2,1,0} reduce(%vmap_jit__where__.90, %constant.113), dimensions={4}, to_apply=%region_239.247, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/vmap()/reduce_max" stack_frame_id=83} + %max.99 = f32[1,1,32,1]{3,2,1,0} maximum(%reduce_max.475, %broadcast.35), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/vmap()/max" stack_frame_id=83} + %broadcast_in_dim.1425 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%max.99), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %sub.379 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1425), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/vmap()/sub" stack_frame_id=83} + %sub.380 = f32[1,1,32,1]{3,2,1,0} reshape(%sub.379), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/vmap()/sub" stack_frame_id=83} + %sub.381 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%sub.380), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/vmap()/sub" stack_frame_id=83} + %sub.382 = f32[1,1,32,1,41]{4,3,2,1,0} subtract(%vmap_jit__where__.90, %sub.381), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/vmap()/sub" stack_frame_id=83} + %exp.95 = f32[1,1,32,1,41]{4,3,2,1,0} exponential(%sub.382), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/vmap()/exp" stack_frame_id=83} + %reduce_sum.1433 = f32[1,1,32,1]{3,2,1,0} reduce(%exp.95, %constant.123), dimensions={4}, to_apply=%region_240.248, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/vmap()/reduce_sum" stack_frame_id=83} + %broadcast_in_dim.1426 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%reduce_sum.1433), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %div.918 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1426), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/vmap()/div" stack_frame_id=83} + %div.919 = f32[1,1,32,1]{3,2,1,0} reshape(%div.918), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/vmap()/div" stack_frame_id=83} + %div.920 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%div.919), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/vmap()/div" stack_frame_id=83} + %div.921 = f32[1,1,32,1,41]{4,3,2,1,0} divide(%exp.95, %div.920), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/vmap()/div" stack_frame_id=83} + %convert_element_type.1643 = bf16[1,1,32,1,41]{4,3,2,1,0} convert(%div.921), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/vmap()/convert_element_type" stack_frame_id=83} + %dot_general.1000 = bf16[1,32,128,1,1]{4,3,2,1,0} dot(%reshape.668, %convert_element_type.1643), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=83} + %transpose.94 = bf16[1,1,32,128]{3,2,1,0} reshape(%dot_general.1000), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/reshape;jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/vmap()/transpose;jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %dot_general.1001 = bf16[1,1,4096]{2,1,0} dot(%transpose.94, %arg_tuple.5#255), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=354} + %add.1060 = bf16[1,1,4096]{2,1,0} add(%dot_general.1001, %add.1049), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/add" stack_frame_id=355} + %convert_element_type.1644 = f32[1,1,4096]{2,1,0} convert(%add.1060), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/feedforward_layernorm/convert_element_type" stack_frame_id=359} + %pow.366 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1644, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/feedforward_layernorm/pow" stack_frame_id=362} + %reduce_sum.1434 = f32[1,1]{1,0} reduce(%pow.366, %constant.123), dimensions={2}, to_apply=%region_241.249, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/feedforward_layernorm/reduce_sum" stack_frame_id=92} + %broadcast_in_dim.1427 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1434), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/feedforward_layernorm/broadcast_in_dim" stack_frame_id=92} + %div.922 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1427, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/feedforward_layernorm/div" stack_frame_id=92} + %add.1061 = f32[1,1,1]{2,1,0} add(%div.922, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/feedforward_layernorm/add" stack_frame_id=363} + %rsqrt.183 = f32[1,1,1]{2,1,0} rsqrt(%add.1061), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/feedforward_layernorm/rsqrt" stack_frame_id=366} + %mul.2944 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.183), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2945 = f32[1,1]{1,0} reshape(%mul.2944), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2946 = f32[1,1,4096]{2,1,0} broadcast(%mul.2945), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2947 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1644, %mul.2946), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/feedforward_layernorm/mul" stack_frame_id=367} + %convert_element_type.1645 = f32[4096]{0} convert(%arg_tuple.5#256), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/feedforward_layernorm/convert_element_type" stack_frame_id=368} + %broadcast_in_dim.1428 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1645), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/feedforward_layernorm/broadcast_in_dim" stack_frame_id=368} + %mul.2948 = f32[1,1,4096]{2,1,0} multiply(%mul.2947, %broadcast_in_dim.1428), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/feedforward_layernorm/mul" stack_frame_id=368} + %convert_element_type.1646 = bf16[1,1,4096]{2,1,0} convert(%mul.2948), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/feedforward_layernorm/convert_element_type" stack_frame_id=372} + %dot_general.1003 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1646, %arg_tuple.5#258), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/feedforward_intermediate_dense/dot_general" stack_frame_id=401} + %dot_general.1002 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1646, %arg_tuple.5#257), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/feedforward_gate_dense/dot_general" stack_frame_id=381} + %convert_element_type.1647 = f32[1,1,14336]{2,1,0} convert(%dot_general.1002), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/convert_element_type" stack_frame_id=385} + %jit_silu_.90 = f32[1,1,14336]{2,1,0} call(%convert_element_type.1647), to_apply=%silu_4.141, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/jit(silu)" stack_frame_id=388} + %convert_element_type.1648 = bf16[1,1,14336]{2,1,0} convert(%jit_silu_.90), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/convert_element_type" stack_frame_id=392} + %mul.2949 = bf16[1,1,14336]{2,1,0} multiply(%dot_general.1003, %convert_element_type.1648), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/mul" stack_frame_id=404} + %dot_general.1004 = bf16[1,1,4096]{2,1,0} dot(%mul.2949, %arg_tuple.5#259), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/feedforward_output_dense/dot_general" stack_frame_id=413} + %add.1062 = bf16[1,1,4096]{2,1,0} add(%dot_general.1004, %add.1060), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/add" stack_frame_id=414} + %convert_element_type.1652 = f32[1,1,4096]{2,1,0} convert(%add.1062), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention_layernorm/convert_element_type" stack_frame_id=164} + %pow.367 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1652, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention_layernorm/pow" stack_frame_id=167} + %reduce_sum.1435 = f32[1,1]{1,0} reduce(%pow.367, %constant.123), dimensions={2}, to_apply=%region_242.250, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention_layernorm/reduce_sum" stack_frame_id=73} + %broadcast_in_dim.1433 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1435), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention_layernorm/broadcast_in_dim" stack_frame_id=73} + %div.923 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1433, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention_layernorm/div" stack_frame_id=73} + %add.1066 = f32[1,1,1]{2,1,0} add(%div.923, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention_layernorm/add" stack_frame_id=168} + %rsqrt.184 = f32[1,1,1]{2,1,0} rsqrt(%add.1066), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention_layernorm/rsqrt" stack_frame_id=171} + %mul.2950 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.184), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2951 = f32[1,1]{1,0} reshape(%mul.2950), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2952 = f32[1,1,4096]{2,1,0} broadcast(%mul.2951), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2953 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1652, %mul.2952), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention_layernorm/mul" stack_frame_id=172} + %convert_element_type.1653 = f32[4096]{0} convert(%arg_tuple.5#260), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention_layernorm/convert_element_type" stack_frame_id=173} + %broadcast_in_dim.1434 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1653), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention_layernorm/broadcast_in_dim" stack_frame_id=173} + %mul.2954 = f32[1,1,4096]{2,1,0} multiply(%mul.2953, %broadcast_in_dim.1434), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention_layernorm/mul" stack_frame_id=173} + %convert_element_type.1654 = bf16[1,1,4096]{2,1,0} convert(%mul.2954), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention_layernorm/convert_element_type" stack_frame_id=177} + %dot_general.1008 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1654, %arg_tuple.5#263), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=265} + %lt.599 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/lt" stack_frame_id=329} + %add.1072 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/add" stack_frame_id=329} + %select_n.148 = s32[] select(%lt.599, %add.1072, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/select_n" stack_frame_id=329} + %dynamic_update_slice.122 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.186, %dot_general.1008, %constant.127, %select_n.148, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/dynamic_update_slice" stack_frame_id=329} + %broadcast_in_dim.1436 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.122), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/broadcast_in_dim" stack_frame_id=338} + %reshape.675 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1436), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/reshape" stack_frame_id=338} + %convert_element_type.1649 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/convert_element_type" stack_frame_id=139} + %add.1063 = f32[1,1]{1,0} reshape(%convert_element_type.1649), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_28/add"} + %ge.465 = f32[1,1]{1,0} broadcast(%add.1063), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/ge" stack_frame_id=143} + %ge.466 = f32[1]{0} reshape(%ge.465), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/ge" stack_frame_id=143} + %ge.467 = f32[1,41]{1,0} broadcast(%ge.466), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/ge" stack_frame_id=143} + %iota.494 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/iota" stack_frame_id=142} + %broadcast_in_dim.1429 = f32[1,41]{1,0} reshape(%iota.494), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/broadcast_in_dim" stack_frame_id=143} + %ge.468 = pred[1,41]{1,0} compare(%ge.467, %broadcast_in_dim.1429), direction=GE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/ge" stack_frame_id=143} + %broadcast_in_dim.1430 = pred[1,1,41]{2,1,0} reshape(%ge.468), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/broadcast_in_dim" stack_frame_id=146} + %convert_element_type.1651 = s32[1,1,41]{2,1,0} convert(%broadcast_in_dim.1430), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/convert_element_type" stack_frame_id=160} + %add.1064 = s32[1,1]{1,0} reshape(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/add" stack_frame_id=147} + %lt.594 = s32[1,1]{1,0} broadcast(%add.1064), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/lt" stack_frame_id=152} + %lt.595 = s32[1]{0} reshape(%lt.594), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/lt" stack_frame_id=152} + %lt.596 = s32[1,41]{1,0} broadcast(%lt.595), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/lt" stack_frame_id=152} + %iota.495 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/iota" stack_frame_id=150} + %broadcast_in_dim.1431 = s32[1,41]{1,0} reshape(%iota.495), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/broadcast_in_dim" stack_frame_id=148} + %add.1065 = s32[1,41]{1,0} add(%broadcast_in_dim.1431, %broadcast.43), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/add" stack_frame_id=151} + %lt.597 = pred[1,41]{1,0} compare(%lt.596, %add.1065), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/lt" stack_frame_id=152} + %convert_element_type.1650 = s32[1,41]{1,0} convert(%lt.597), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/convert_element_type" stack_frame_id=156} + %broadcast_in_dim.1432 = s32[1,1,41]{2,1,0} reshape(%convert_element_type.1650), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/broadcast_in_dim" stack_frame_id=160} + %min.91 = s32[1,1,41]{2,1,0} minimum(%convert_element_type.1651, %broadcast_in_dim.1432), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/min" stack_frame_id=160} + %broadcast_in_dim.1437 = s32[1,1,1,41]{3,2,1,0} reshape(%min.91), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/broadcast_in_dim" stack_frame_id=341} + %convert_element_type.1661 = pred[1,1,1,41]{3,2,1,0} compare(%broadcast_in_dim.1437, %broadcast.37), direction=NE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/convert_element_type" stack_frame_id=345} + %broadcast_in_dim.1438 = pred[1,1,1,1,41]{4,3,2,1,0} reshape(%convert_element_type.1661), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/broadcast_in_dim" stack_frame_id=83} + %and.281 = pred[1,1,1,1,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1438), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/vmap()/and" stack_frame_id=83} + %and.282 = pred[1,1,1,41]{3,2,1,0} reshape(%and.281), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/vmap()/and" stack_frame_id=83} + %and.283 = pred[1,1,32,1,41]{4,3,2,1,0} broadcast(%and.282), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/vmap()/and" stack_frame_id=83} + %dot_general.1005 = bf16[1,1,32,128]{3,2,1,0} dot(%convert_element_type.1654, %arg_tuple.5#261), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=186} + %convert_element_type.1655 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/convert_element_type" stack_frame_id=208} + %add.1067 = f32[1,1]{1,0} reshape(%convert_element_type.1655), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/add"} + %iota.496 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/iota" stack_frame_id=197} + %mul.2955 = f32[64]{0} multiply(%iota.496, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=197} + %div.924 = f32[64]{0} divide(%mul.2955, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/div" stack_frame_id=198} + %neg.374 = f32[64]{0} negate(%div.924), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/neg" stack_frame_id=199} + %pow.368 = f32[64]{0} power(%broadcast.39, %neg.374), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/pow" stack_frame_id=202} + %div.925 = f32[64]{0} divide(%pow.368, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/div" stack_frame_id=203} + %dot_general.1006 = f32[1,1,64]{2,1,0} dot(%add.1067, %div.925), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/bi,j->bij/dot_general" stack_frame_id=211} + %stack.1440 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.1006), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/stack" stack_frame_id=214} + %stack.1441 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.1006), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/stack" stack_frame_id=214} + %stack.1442 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1440, %stack.1441), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/stack" stack_frame_id=214} + %reshape.670 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1442), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/reshape"} + %cos.183 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.670), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/cos" stack_frame_id=218} + %convert_element_type.1656 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.183), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/convert_element_type" stack_frame_id=222} + %mul.2956 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1656), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=242} + %mul.2957 = bf16[1,1,128]{2,1,0} reshape(%mul.2956), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=242} + %mul.2958 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2957), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=242} + %mul.2959 = bf16[1,1,32,128]{3,2,1,0} multiply(%dot_general.1005, %mul.2958), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=242} + %split.367 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.1005), slice={[0:1], [0:1], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/split" stack_frame_id=234} + %neg.375 = bf16[1,1,32,64]{3,2,1,0} negate(%split.367), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/neg" stack_frame_id=235} + %stack.1443 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%neg.375), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/stack" stack_frame_id=238} + %split.366 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.1005), slice={[0:1], [0:1], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/split" stack_frame_id=234} + %stack.1444 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%split.366), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/stack" stack_frame_id=238} + %stack.1445 = bf16[1,1,32,2,64]{4,3,2,1,0} concatenate(%stack.1443, %stack.1444), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/stack" stack_frame_id=238} + %reshape.671 = bf16[1,1,32,128]{3,2,1,0} reshape(%stack.1445), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/reshape" stack_frame_id=241} + %sin.183 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.670), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/sin" stack_frame_id=226} + %convert_element_type.1657 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.183), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/convert_element_type" stack_frame_id=230} + %mul.2960 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1657), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=243} + %mul.2961 = bf16[1,1,128]{2,1,0} reshape(%mul.2960), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=243} + %mul.2962 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2961), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=243} + %mul.2963 = bf16[1,1,32,128]{3,2,1,0} multiply(%reshape.671, %mul.2962), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=243} + %add.1068 = bf16[1,1,32,128]{3,2,1,0} add(%mul.2959, %mul.2963), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/add" stack_frame_id=244} + %reshape.676 = bf16[1,1,32,1,128]{4,3,2,1,0} reshape(%add.1068), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/reshape" stack_frame_id=83} + %slice.181 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.184), slice={[0:1], [0:1], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/slice" stack_frame_id=245} + %squeeze.185 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.181), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/squeeze" stack_frame_id=245} + %dot_general.1007 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1654, %arg_tuple.5#262), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=256} + %convert_element_type.1658 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/convert_element_type" stack_frame_id=287} + %add.1069 = f32[1,1]{1,0} reshape(%convert_element_type.1658), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/add"} + %iota.497 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/iota" stack_frame_id=276} + %mul.2964 = f32[64]{0} multiply(%iota.497, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=276} + %div.926 = f32[64]{0} divide(%mul.2964, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/div" stack_frame_id=277} + %neg.376 = f32[64]{0} negate(%div.926), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/neg" stack_frame_id=278} + %pow.369 = f32[64]{0} power(%broadcast.39, %neg.376), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/pow" stack_frame_id=281} + %div.927 = f32[64]{0} divide(%pow.369, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/div" stack_frame_id=282} + %dot_general.1009 = f32[1,1,64]{2,1,0} dot(%add.1069, %div.927), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/bi,j->bij/dot_general" stack_frame_id=290} + %stack.1446 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.1009), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/stack" stack_frame_id=293} + %stack.1447 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.1009), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/stack" stack_frame_id=293} + %stack.1448 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1446, %stack.1447), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/stack" stack_frame_id=293} + %reshape.672 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1448), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/reshape"} + %cos.184 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.672), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/cos" stack_frame_id=297} + %convert_element_type.1659 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.184), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/convert_element_type" stack_frame_id=301} + %mul.2965 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1659), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=321} + %mul.2966 = bf16[1,1,128]{2,1,0} reshape(%mul.2965), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=321} + %mul.2967 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2966), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=321} + %mul.2968 = bf16[1,1,8,128]{3,2,1,0} multiply(%dot_general.1007, %mul.2967), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=321} + %split.369 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.1007), slice={[0:1], [0:1], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/split" stack_frame_id=313} + %neg.377 = bf16[1,1,8,64]{3,2,1,0} negate(%split.369), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/neg" stack_frame_id=314} + %stack.1449 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%neg.377), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/stack" stack_frame_id=317} + %split.368 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.1007), slice={[0:1], [0:1], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/split" stack_frame_id=313} + %stack.1450 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%split.368), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/stack" stack_frame_id=317} + %stack.1451 = bf16[1,1,8,2,64]{4,3,2,1,0} concatenate(%stack.1449, %stack.1450), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/stack" stack_frame_id=317} + %reshape.673 = bf16[1,1,8,128]{3,2,1,0} reshape(%stack.1451), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/reshape" stack_frame_id=320} + %sin.184 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.672), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/sin" stack_frame_id=305} + %convert_element_type.1660 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.184), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/convert_element_type" stack_frame_id=309} + %mul.2969 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1660), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=322} + %mul.2970 = bf16[1,1,128]{2,1,0} reshape(%mul.2969), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=322} + %mul.2971 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2970), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=322} + %mul.2972 = bf16[1,1,8,128]{3,2,1,0} multiply(%reshape.673, %mul.2971), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=322} + %add.1070 = bf16[1,1,8,128]{3,2,1,0} add(%mul.2968, %mul.2972), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/rotary_embedding_28/add" stack_frame_id=323} + %lt.598 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/lt" stack_frame_id=326} + %add.1071 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/add" stack_frame_id=326} + %select_n.147 = s32[] select(%lt.598, %add.1071, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/select_n" stack_frame_id=326} + %dynamic_update_slice.121 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.185, %add.1070, %constant.127, %select_n.147, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/dynamic_update_slice" stack_frame_id=326} + %broadcast_in_dim.1435 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.121), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/broadcast_in_dim" stack_frame_id=335} + %reshape.674 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1435), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/reshape" stack_frame_id=335} + %dot_general.1010 = f32[1,32,1,1,41]{4,3,2,1,0} dot(%reshape.676, %reshape.674), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=83} + %mul.2973 = f32[1,32,1,1,41]{4,3,2,1,0} multiply(%dot_general.1010, %broadcast.36), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/vmap()/mul" stack_frame_id=83} + %vmap_jit__where__.91 = f32[1,1,32,1,41]{4,3,2,1,0} call(%and.283, %mul.2973, %constant.114), to_apply=%_where_3.137, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/vmap(jit(_where))" stack_frame_id=83} + %reduce_max.476 = f32[1,1,32,1]{3,2,1,0} reduce(%vmap_jit__where__.91, %constant.113), dimensions={4}, to_apply=%region_243.251, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/vmap()/reduce_max" stack_frame_id=83} + %max.100 = f32[1,1,32,1]{3,2,1,0} maximum(%reduce_max.476, %broadcast.35), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/vmap()/max" stack_frame_id=83} + %broadcast_in_dim.1439 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%max.100), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %sub.383 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1439), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/vmap()/sub" stack_frame_id=83} + %sub.384 = f32[1,1,32,1]{3,2,1,0} reshape(%sub.383), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/vmap()/sub" stack_frame_id=83} + %sub.385 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%sub.384), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/vmap()/sub" stack_frame_id=83} + %sub.386 = f32[1,1,32,1,41]{4,3,2,1,0} subtract(%vmap_jit__where__.91, %sub.385), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/vmap()/sub" stack_frame_id=83} + %exp.96 = f32[1,1,32,1,41]{4,3,2,1,0} exponential(%sub.386), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/vmap()/exp" stack_frame_id=83} + %reduce_sum.1436 = f32[1,1,32,1]{3,2,1,0} reduce(%exp.96, %constant.123), dimensions={4}, to_apply=%region_244.252, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/vmap()/reduce_sum" stack_frame_id=83} + %broadcast_in_dim.1440 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%reduce_sum.1436), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %div.928 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1440), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/vmap()/div" stack_frame_id=83} + %div.929 = f32[1,1,32,1]{3,2,1,0} reshape(%div.928), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/vmap()/div" stack_frame_id=83} + %div.930 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%div.929), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/vmap()/div" stack_frame_id=83} + %div.931 = f32[1,1,32,1,41]{4,3,2,1,0} divide(%exp.96, %div.930), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/vmap()/div" stack_frame_id=83} + %convert_element_type.1662 = bf16[1,1,32,1,41]{4,3,2,1,0} convert(%div.931), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/vmap()/convert_element_type" stack_frame_id=83} + %dot_general.1011 = bf16[1,32,128,1,1]{4,3,2,1,0} dot(%reshape.675, %convert_element_type.1662), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=83} + %transpose.95 = bf16[1,1,32,128]{3,2,1,0} reshape(%dot_general.1011), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/reshape;jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/vmap()/transpose;jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %dot_general.1012 = bf16[1,1,4096]{2,1,0} dot(%transpose.95, %arg_tuple.5#264), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=354} + %add.1073 = bf16[1,1,4096]{2,1,0} add(%dot_general.1012, %add.1062), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/add" stack_frame_id=355} + %convert_element_type.1663 = f32[1,1,4096]{2,1,0} convert(%add.1073), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/feedforward_layernorm/convert_element_type" stack_frame_id=359} + %pow.370 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1663, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/feedforward_layernorm/pow" stack_frame_id=362} + %reduce_sum.1437 = f32[1,1]{1,0} reduce(%pow.370, %constant.123), dimensions={2}, to_apply=%region_245.253, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/feedforward_layernorm/reduce_sum" stack_frame_id=92} + %broadcast_in_dim.1441 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1437), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/feedforward_layernorm/broadcast_in_dim" stack_frame_id=92} + %div.932 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1441, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/feedforward_layernorm/div" stack_frame_id=92} + %add.1074 = f32[1,1,1]{2,1,0} add(%div.932, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/feedforward_layernorm/add" stack_frame_id=363} + %rsqrt.185 = f32[1,1,1]{2,1,0} rsqrt(%add.1074), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/feedforward_layernorm/rsqrt" stack_frame_id=366} + %mul.2974 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.185), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2975 = f32[1,1]{1,0} reshape(%mul.2974), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2976 = f32[1,1,4096]{2,1,0} broadcast(%mul.2975), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/feedforward_layernorm/mul" stack_frame_id=367} + %mul.2977 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1663, %mul.2976), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/feedforward_layernorm/mul" stack_frame_id=367} + %convert_element_type.1664 = f32[4096]{0} convert(%arg_tuple.5#265), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/feedforward_layernorm/convert_element_type" stack_frame_id=368} + %broadcast_in_dim.1442 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1664), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/feedforward_layernorm/broadcast_in_dim" stack_frame_id=368} + %mul.2978 = f32[1,1,4096]{2,1,0} multiply(%mul.2977, %broadcast_in_dim.1442), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/feedforward_layernorm/mul" stack_frame_id=368} + %convert_element_type.1665 = bf16[1,1,4096]{2,1,0} convert(%mul.2978), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/feedforward_layernorm/convert_element_type" stack_frame_id=372} + %dot_general.1014 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1665, %arg_tuple.5#267), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/feedforward_intermediate_dense/dot_general" stack_frame_id=401} + %dot_general.1013 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1665, %arg_tuple.5#266), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/feedforward_gate_dense/dot_general" stack_frame_id=381} + %convert_element_type.1666 = f32[1,1,14336]{2,1,0} convert(%dot_general.1013), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/convert_element_type" stack_frame_id=385} + %jit_silu_.91 = f32[1,1,14336]{2,1,0} call(%convert_element_type.1666), to_apply=%silu_4.141, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/jit(silu)" stack_frame_id=388} + %convert_element_type.1667 = bf16[1,1,14336]{2,1,0} convert(%jit_silu_.91), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/convert_element_type" stack_frame_id=392} + %mul.2979 = bf16[1,1,14336]{2,1,0} multiply(%dot_general.1014, %convert_element_type.1667), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/mul" stack_frame_id=404} + %dot_general.1015 = bf16[1,1,4096]{2,1,0} dot(%mul.2979, %arg_tuple.5#268), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/feedforward_output_dense/dot_general" stack_frame_id=413} + %add.1075 = bf16[1,1,4096]{2,1,0} add(%dot_general.1015, %add.1073), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/add" stack_frame_id=414} + %convert_element_type.1671 = f32[1,1,4096]{2,1,0} convert(%add.1075), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention_layernorm/convert_element_type" stack_frame_id=164} + %pow.371 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1671, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention_layernorm/pow" stack_frame_id=167} + %reduce_sum.1438 = f32[1,1]{1,0} reduce(%pow.371, %constant.123), dimensions={2}, to_apply=%region_246.254, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention_layernorm/reduce_sum" stack_frame_id=73} + %broadcast_in_dim.1447 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1438), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention_layernorm/broadcast_in_dim" stack_frame_id=73} + %div.933 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1447, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention_layernorm/div" stack_frame_id=73} + %add.1079 = f32[1,1,1]{2,1,0} add(%div.933, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention_layernorm/add" stack_frame_id=168} + %rsqrt.186 = f32[1,1,1]{2,1,0} rsqrt(%add.1079), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention_layernorm/rsqrt" stack_frame_id=171} + %mul.2980 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.186), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2981 = f32[1,1]{1,0} reshape(%mul.2980), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2982 = f32[1,1,4096]{2,1,0} broadcast(%mul.2981), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention_layernorm/mul" stack_frame_id=172} + %mul.2983 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1671, %mul.2982), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention_layernorm/mul" stack_frame_id=172} + %convert_element_type.1672 = f32[4096]{0} convert(%arg_tuple.5#269), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention_layernorm/convert_element_type" stack_frame_id=173} + %broadcast_in_dim.1448 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1672), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention_layernorm/broadcast_in_dim" stack_frame_id=173} + %mul.2984 = f32[1,1,4096]{2,1,0} multiply(%mul.2983, %broadcast_in_dim.1448), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention_layernorm/mul" stack_frame_id=173} + %convert_element_type.1673 = bf16[1,1,4096]{2,1,0} convert(%mul.2984), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention_layernorm/convert_element_type" stack_frame_id=177} + %dot_general.1019 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1673, %arg_tuple.5#272), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=265} + %lt.605 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/lt" stack_frame_id=329} + %add.1085 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/add" stack_frame_id=329} + %select_n.150 = s32[] select(%lt.605, %add.1085, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/select_n" stack_frame_id=329} + %dynamic_update_slice.124 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.189, %dot_general.1019, %constant.127, %select_n.150, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/dynamic_update_slice" stack_frame_id=329} + %broadcast_in_dim.1450 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.124), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/broadcast_in_dim" stack_frame_id=338} + %reshape.682 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1450), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/reshape" stack_frame_id=338} + %convert_element_type.1668 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/convert_element_type" stack_frame_id=139} + %add.1076 = f32[1,1]{1,0} reshape(%convert_element_type.1668), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_29/add"} + %ge.469 = f32[1,1]{1,0} broadcast(%add.1076), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/ge" stack_frame_id=143} + %ge.470 = f32[1]{0} reshape(%ge.469), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/ge" stack_frame_id=143} + %ge.471 = f32[1,41]{1,0} broadcast(%ge.470), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/ge" stack_frame_id=143} + %iota.498 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/iota" stack_frame_id=142} + %broadcast_in_dim.1443 = f32[1,41]{1,0} reshape(%iota.498), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/broadcast_in_dim" stack_frame_id=143} + %ge.472 = pred[1,41]{1,0} compare(%ge.471, %broadcast_in_dim.1443), direction=GE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/ge" stack_frame_id=143} + %broadcast_in_dim.1444 = pred[1,1,41]{2,1,0} reshape(%ge.472), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/broadcast_in_dim" stack_frame_id=146} + %convert_element_type.1670 = s32[1,1,41]{2,1,0} convert(%broadcast_in_dim.1444), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/convert_element_type" stack_frame_id=160} + %add.1077 = s32[1,1]{1,0} reshape(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/add" stack_frame_id=147} + %lt.600 = s32[1,1]{1,0} broadcast(%add.1077), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/lt" stack_frame_id=152} + %lt.601 = s32[1]{0} reshape(%lt.600), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/lt" stack_frame_id=152} + %lt.602 = s32[1,41]{1,0} broadcast(%lt.601), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/lt" stack_frame_id=152} + %iota.499 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/iota" stack_frame_id=150} + %broadcast_in_dim.1445 = s32[1,41]{1,0} reshape(%iota.499), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/broadcast_in_dim" stack_frame_id=148} + %add.1078 = s32[1,41]{1,0} add(%broadcast_in_dim.1445, %broadcast.43), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/add" stack_frame_id=151} + %lt.603 = pred[1,41]{1,0} compare(%lt.602, %add.1078), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/lt" stack_frame_id=152} + %convert_element_type.1669 = s32[1,41]{1,0} convert(%lt.603), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/convert_element_type" stack_frame_id=156} + %broadcast_in_dim.1446 = s32[1,1,41]{2,1,0} reshape(%convert_element_type.1669), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/broadcast_in_dim" stack_frame_id=160} + %min.92 = s32[1,1,41]{2,1,0} minimum(%convert_element_type.1670, %broadcast_in_dim.1446), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/min" stack_frame_id=160} + %broadcast_in_dim.1451 = s32[1,1,1,41]{3,2,1,0} reshape(%min.92), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/broadcast_in_dim" stack_frame_id=341} + %convert_element_type.1680 = pred[1,1,1,41]{3,2,1,0} compare(%broadcast_in_dim.1451, %broadcast.37), direction=NE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/convert_element_type" stack_frame_id=345} + %broadcast_in_dim.1452 = pred[1,1,1,1,41]{4,3,2,1,0} reshape(%convert_element_type.1680), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/broadcast_in_dim" stack_frame_id=83} + %and.284 = pred[1,1,1,1,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1452), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/vmap()/and" stack_frame_id=83} + %and.285 = pred[1,1,1,41]{3,2,1,0} reshape(%and.284), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/vmap()/and" stack_frame_id=83} + %and.286 = pred[1,1,32,1,41]{4,3,2,1,0} broadcast(%and.285), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/vmap()/and" stack_frame_id=83} + %dot_general.1016 = bf16[1,1,32,128]{3,2,1,0} dot(%convert_element_type.1673, %arg_tuple.5#270), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=186} + %convert_element_type.1674 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/convert_element_type" stack_frame_id=208} + %add.1080 = f32[1,1]{1,0} reshape(%convert_element_type.1674), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/add"} + %iota.500 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/iota" stack_frame_id=197} + %mul.2985 = f32[64]{0} multiply(%iota.500, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=197} + %div.934 = f32[64]{0} divide(%mul.2985, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/div" stack_frame_id=198} + %neg.378 = f32[64]{0} negate(%div.934), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/neg" stack_frame_id=199} + %pow.372 = f32[64]{0} power(%broadcast.39, %neg.378), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/pow" stack_frame_id=202} + %div.935 = f32[64]{0} divide(%pow.372, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/div" stack_frame_id=203} + %dot_general.1017 = f32[1,1,64]{2,1,0} dot(%add.1080, %div.935), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/bi,j->bij/dot_general" stack_frame_id=211} + %stack.1455 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.1017), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/stack" stack_frame_id=214} + %stack.1456 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.1017), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/stack" stack_frame_id=214} + %stack.1457 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1455, %stack.1456), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/stack" stack_frame_id=214} + %reshape.677 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1457), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/reshape"} + %cos.185 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.677), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/cos" stack_frame_id=218} + %convert_element_type.1675 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.185), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/convert_element_type" stack_frame_id=222} + %mul.2986 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1675), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=242} + %mul.2987 = bf16[1,1,128]{2,1,0} reshape(%mul.2986), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=242} + %mul.2988 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2987), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=242} + %mul.2989 = bf16[1,1,32,128]{3,2,1,0} multiply(%dot_general.1016, %mul.2988), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=242} + %split.371 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.1016), slice={[0:1], [0:1], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/split" stack_frame_id=234} + %neg.379 = bf16[1,1,32,64]{3,2,1,0} negate(%split.371), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/neg" stack_frame_id=235} + %stack.1458 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%neg.379), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/stack" stack_frame_id=238} + %split.370 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.1016), slice={[0:1], [0:1], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/split" stack_frame_id=234} + %stack.1459 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%split.370), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/stack" stack_frame_id=238} + %stack.1460 = bf16[1,1,32,2,64]{4,3,2,1,0} concatenate(%stack.1458, %stack.1459), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/stack" stack_frame_id=238} + %reshape.678 = bf16[1,1,32,128]{3,2,1,0} reshape(%stack.1460), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/reshape" stack_frame_id=241} + %sin.185 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.677), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/sin" stack_frame_id=226} + %convert_element_type.1676 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.185), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/convert_element_type" stack_frame_id=230} + %mul.2990 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1676), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=243} + %mul.2991 = bf16[1,1,128]{2,1,0} reshape(%mul.2990), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=243} + %mul.2992 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.2991), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=243} + %mul.2993 = bf16[1,1,32,128]{3,2,1,0} multiply(%reshape.678, %mul.2992), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=243} + %add.1081 = bf16[1,1,32,128]{3,2,1,0} add(%mul.2989, %mul.2993), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/add" stack_frame_id=244} + %reshape.683 = bf16[1,1,32,1,128]{4,3,2,1,0} reshape(%add.1081), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/reshape" stack_frame_id=83} + %slice.184 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.187), slice={[0:1], [0:1], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/slice" stack_frame_id=245} + %squeeze.188 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.184), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/squeeze" stack_frame_id=245} + %dot_general.1018 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1673, %arg_tuple.5#271), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=256} + %convert_element_type.1677 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/convert_element_type" stack_frame_id=287} + %add.1082 = f32[1,1]{1,0} reshape(%convert_element_type.1677), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/add"} + %iota.501 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/iota" stack_frame_id=276} + %mul.2994 = f32[64]{0} multiply(%iota.501, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=276} + %div.936 = f32[64]{0} divide(%mul.2994, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/div" stack_frame_id=277} + %neg.380 = f32[64]{0} negate(%div.936), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/neg" stack_frame_id=278} + %pow.373 = f32[64]{0} power(%broadcast.39, %neg.380), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/pow" stack_frame_id=281} + %div.937 = f32[64]{0} divide(%pow.373, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/div" stack_frame_id=282} + %dot_general.1020 = f32[1,1,64]{2,1,0} dot(%add.1082, %div.937), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/bi,j->bij/dot_general" stack_frame_id=290} + %stack.1461 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.1020), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/stack" stack_frame_id=293} + %stack.1462 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.1020), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/stack" stack_frame_id=293} + %stack.1463 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1461, %stack.1462), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/stack" stack_frame_id=293} + %reshape.679 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1463), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/reshape"} + %cos.186 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.679), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/cos" stack_frame_id=297} + %convert_element_type.1678 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.186), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/convert_element_type" stack_frame_id=301} + %mul.2995 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1678), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=321} + %mul.2996 = bf16[1,1,128]{2,1,0} reshape(%mul.2995), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=321} + %mul.2997 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.2996), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=321} + %mul.2998 = bf16[1,1,8,128]{3,2,1,0} multiply(%dot_general.1018, %mul.2997), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=321} + %split.373 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.1018), slice={[0:1], [0:1], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/split" stack_frame_id=313} + %neg.381 = bf16[1,1,8,64]{3,2,1,0} negate(%split.373), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/neg" stack_frame_id=314} + %stack.1464 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%neg.381), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/stack" stack_frame_id=317} + %split.372 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.1018), slice={[0:1], [0:1], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/split" stack_frame_id=313} + %stack.1465 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%split.372), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/stack" stack_frame_id=317} + %stack.1466 = bf16[1,1,8,2,64]{4,3,2,1,0} concatenate(%stack.1464, %stack.1465), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/stack" stack_frame_id=317} + %reshape.680 = bf16[1,1,8,128]{3,2,1,0} reshape(%stack.1466), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/reshape" stack_frame_id=320} + %sin.186 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.679), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/sin" stack_frame_id=305} + %convert_element_type.1679 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.186), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/convert_element_type" stack_frame_id=309} + %mul.2999 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1679), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=322} + %mul.3000 = bf16[1,1,128]{2,1,0} reshape(%mul.2999), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=322} + %mul.3001 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.3000), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=322} + %mul.3002 = bf16[1,1,8,128]{3,2,1,0} multiply(%reshape.680, %mul.3001), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=322} + %add.1083 = bf16[1,1,8,128]{3,2,1,0} add(%mul.2998, %mul.3002), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/rotary_embedding_29/add" stack_frame_id=323} + %lt.604 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/lt" stack_frame_id=326} + %add.1084 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/add" stack_frame_id=326} + %select_n.149 = s32[] select(%lt.604, %add.1084, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/select_n" stack_frame_id=326} + %dynamic_update_slice.123 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.188, %add.1083, %constant.127, %select_n.149, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/dynamic_update_slice" stack_frame_id=326} + %broadcast_in_dim.1449 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.123), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/broadcast_in_dim" stack_frame_id=335} + %reshape.681 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1449), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/reshape" stack_frame_id=335} + %dot_general.1021 = f32[1,32,1,1,41]{4,3,2,1,0} dot(%reshape.683, %reshape.681), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=83} + %mul.3003 = f32[1,32,1,1,41]{4,3,2,1,0} multiply(%dot_general.1021, %broadcast.36), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/vmap()/mul" stack_frame_id=83} + %vmap_jit__where__.92 = f32[1,1,32,1,41]{4,3,2,1,0} call(%and.286, %mul.3003, %constant.114), to_apply=%_where_3.137, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/vmap(jit(_where))" stack_frame_id=83} + %reduce_max.477 = f32[1,1,32,1]{3,2,1,0} reduce(%vmap_jit__where__.92, %constant.113), dimensions={4}, to_apply=%region_247.255, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/vmap()/reduce_max" stack_frame_id=83} + %max.101 = f32[1,1,32,1]{3,2,1,0} maximum(%reduce_max.477, %broadcast.35), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/vmap()/max" stack_frame_id=83} + %broadcast_in_dim.1453 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%max.101), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %sub.387 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1453), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/vmap()/sub" stack_frame_id=83} + %sub.388 = f32[1,1,32,1]{3,2,1,0} reshape(%sub.387), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/vmap()/sub" stack_frame_id=83} + %sub.389 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%sub.388), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/vmap()/sub" stack_frame_id=83} + %sub.390 = f32[1,1,32,1,41]{4,3,2,1,0} subtract(%vmap_jit__where__.92, %sub.389), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/vmap()/sub" stack_frame_id=83} + %exp.97 = f32[1,1,32,1,41]{4,3,2,1,0} exponential(%sub.390), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/vmap()/exp" stack_frame_id=83} + %reduce_sum.1439 = f32[1,1,32,1]{3,2,1,0} reduce(%exp.97, %constant.123), dimensions={4}, to_apply=%region_248.256, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/vmap()/reduce_sum" stack_frame_id=83} + %broadcast_in_dim.1454 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%reduce_sum.1439), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %div.938 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1454), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/vmap()/div" stack_frame_id=83} + %div.939 = f32[1,1,32,1]{3,2,1,0} reshape(%div.938), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/vmap()/div" stack_frame_id=83} + %div.940 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%div.939), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/vmap()/div" stack_frame_id=83} + %div.941 = f32[1,1,32,1,41]{4,3,2,1,0} divide(%exp.97, %div.940), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/vmap()/div" stack_frame_id=83} + %convert_element_type.1681 = bf16[1,1,32,1,41]{4,3,2,1,0} convert(%div.941), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/vmap()/convert_element_type" stack_frame_id=83} + %dot_general.1022 = bf16[1,32,128,1,1]{4,3,2,1,0} dot(%reshape.682, %convert_element_type.1681), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=83} + %transpose.96 = bf16[1,1,32,128]{3,2,1,0} reshape(%dot_general.1022), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/reshape;jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/vmap()/transpose;jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %dot_general.1023 = bf16[1,1,4096]{2,1,0} dot(%transpose.96, %arg_tuple.5#273), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=354} + %add.1086 = bf16[1,1,4096]{2,1,0} add(%dot_general.1023, %add.1075), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/add" stack_frame_id=355} + %convert_element_type.1682 = f32[1,1,4096]{2,1,0} convert(%add.1086), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/feedforward_layernorm/convert_element_type" stack_frame_id=359} + %pow.374 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1682, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/feedforward_layernorm/pow" stack_frame_id=362} + %reduce_sum.1440 = f32[1,1]{1,0} reduce(%pow.374, %constant.123), dimensions={2}, to_apply=%region_249.257, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/feedforward_layernorm/reduce_sum" stack_frame_id=92} + %broadcast_in_dim.1455 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1440), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/feedforward_layernorm/broadcast_in_dim" stack_frame_id=92} + %div.942 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1455, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/feedforward_layernorm/div" stack_frame_id=92} + %add.1087 = f32[1,1,1]{2,1,0} add(%div.942, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/feedforward_layernorm/add" stack_frame_id=363} + %rsqrt.187 = f32[1,1,1]{2,1,0} rsqrt(%add.1087), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/feedforward_layernorm/rsqrt" stack_frame_id=366} + %mul.3004 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.187), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/feedforward_layernorm/mul" stack_frame_id=367} + %mul.3005 = f32[1,1]{1,0} reshape(%mul.3004), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/feedforward_layernorm/mul" stack_frame_id=367} + %mul.3006 = f32[1,1,4096]{2,1,0} broadcast(%mul.3005), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/feedforward_layernorm/mul" stack_frame_id=367} + %mul.3007 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1682, %mul.3006), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/feedforward_layernorm/mul" stack_frame_id=367} + %convert_element_type.1683 = f32[4096]{0} convert(%arg_tuple.5#274), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/feedforward_layernorm/convert_element_type" stack_frame_id=368} + %broadcast_in_dim.1456 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1683), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/feedforward_layernorm/broadcast_in_dim" stack_frame_id=368} + %mul.3008 = f32[1,1,4096]{2,1,0} multiply(%mul.3007, %broadcast_in_dim.1456), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/feedforward_layernorm/mul" stack_frame_id=368} + %convert_element_type.1684 = bf16[1,1,4096]{2,1,0} convert(%mul.3008), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/feedforward_layernorm/convert_element_type" stack_frame_id=372} + %dot_general.1025 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1684, %arg_tuple.5#276), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/feedforward_intermediate_dense/dot_general" stack_frame_id=401} + %dot_general.1024 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1684, %arg_tuple.5#275), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/feedforward_gate_dense/dot_general" stack_frame_id=381} + %convert_element_type.1685 = f32[1,1,14336]{2,1,0} convert(%dot_general.1024), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/convert_element_type" stack_frame_id=385} + %jit_silu_.92 = f32[1,1,14336]{2,1,0} call(%convert_element_type.1685), to_apply=%silu_4.141, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/jit(silu)" stack_frame_id=388} + %convert_element_type.1686 = bf16[1,1,14336]{2,1,0} convert(%jit_silu_.92), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/convert_element_type" stack_frame_id=392} + %mul.3009 = bf16[1,1,14336]{2,1,0} multiply(%dot_general.1025, %convert_element_type.1686), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/mul" stack_frame_id=404} + %dot_general.1026 = bf16[1,1,4096]{2,1,0} dot(%mul.3009, %arg_tuple.5#277), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/feedforward_output_dense/dot_general" stack_frame_id=413} + %add.1088 = bf16[1,1,4096]{2,1,0} add(%dot_general.1026, %add.1086), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/add" stack_frame_id=414} + %convert_element_type.1690 = f32[1,1,4096]{2,1,0} convert(%add.1088), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention_layernorm/convert_element_type" stack_frame_id=164} + %pow.375 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1690, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention_layernorm/pow" stack_frame_id=167} + %reduce_sum.1441 = f32[1,1]{1,0} reduce(%pow.375, %constant.123), dimensions={2}, to_apply=%region_250.258, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention_layernorm/reduce_sum" stack_frame_id=73} + %broadcast_in_dim.1461 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1441), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention_layernorm/broadcast_in_dim" stack_frame_id=73} + %div.943 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1461, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention_layernorm/div" stack_frame_id=73} + %add.1092 = f32[1,1,1]{2,1,0} add(%div.943, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention_layernorm/add" stack_frame_id=168} + %rsqrt.188 = f32[1,1,1]{2,1,0} rsqrt(%add.1092), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention_layernorm/rsqrt" stack_frame_id=171} + %mul.3010 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.188), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention_layernorm/mul" stack_frame_id=172} + %mul.3011 = f32[1,1]{1,0} reshape(%mul.3010), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention_layernorm/mul" stack_frame_id=172} + %mul.3012 = f32[1,1,4096]{2,1,0} broadcast(%mul.3011), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention_layernorm/mul" stack_frame_id=172} + %mul.3013 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1690, %mul.3012), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention_layernorm/mul" stack_frame_id=172} + %convert_element_type.1691 = f32[4096]{0} convert(%arg_tuple.5#278), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention_layernorm/convert_element_type" stack_frame_id=173} + %broadcast_in_dim.1462 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1691), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention_layernorm/broadcast_in_dim" stack_frame_id=173} + %mul.3014 = f32[1,1,4096]{2,1,0} multiply(%mul.3013, %broadcast_in_dim.1462), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention_layernorm/mul" stack_frame_id=173} + %convert_element_type.1692 = bf16[1,1,4096]{2,1,0} convert(%mul.3014), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention_layernorm/convert_element_type" stack_frame_id=177} + %dot_general.1030 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1692, %arg_tuple.5#281), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=265} + %lt.611 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/lt" stack_frame_id=329} + %add.1098 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/add" stack_frame_id=329} + %select_n.152 = s32[] select(%lt.611, %add.1098, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/select_n" stack_frame_id=329} + %dynamic_update_slice.126 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.192, %dot_general.1030, %constant.127, %select_n.152, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/dynamic_update_slice" stack_frame_id=329} + %broadcast_in_dim.1464 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.126), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/broadcast_in_dim" stack_frame_id=338} + %reshape.689 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1464), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/reshape" stack_frame_id=338} + %convert_element_type.1687 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/convert_element_type" stack_frame_id=139} + %add.1089 = f32[1,1]{1,0} reshape(%convert_element_type.1687), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_30/add"} + %ge.473 = f32[1,1]{1,0} broadcast(%add.1089), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/ge" stack_frame_id=143} + %ge.474 = f32[1]{0} reshape(%ge.473), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/ge" stack_frame_id=143} + %ge.475 = f32[1,41]{1,0} broadcast(%ge.474), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/ge" stack_frame_id=143} + %iota.502 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/iota" stack_frame_id=142} + %broadcast_in_dim.1457 = f32[1,41]{1,0} reshape(%iota.502), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/broadcast_in_dim" stack_frame_id=143} + %ge.476 = pred[1,41]{1,0} compare(%ge.475, %broadcast_in_dim.1457), direction=GE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/ge" stack_frame_id=143} + %broadcast_in_dim.1458 = pred[1,1,41]{2,1,0} reshape(%ge.476), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/broadcast_in_dim" stack_frame_id=146} + %convert_element_type.1689 = s32[1,1,41]{2,1,0} convert(%broadcast_in_dim.1458), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/convert_element_type" stack_frame_id=160} + %add.1090 = s32[1,1]{1,0} reshape(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/add" stack_frame_id=147} + %lt.606 = s32[1,1]{1,0} broadcast(%add.1090), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/lt" stack_frame_id=152} + %lt.607 = s32[1]{0} reshape(%lt.606), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/lt" stack_frame_id=152} + %lt.608 = s32[1,41]{1,0} broadcast(%lt.607), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/lt" stack_frame_id=152} + %iota.503 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/iota" stack_frame_id=150} + %broadcast_in_dim.1459 = s32[1,41]{1,0} reshape(%iota.503), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/broadcast_in_dim" stack_frame_id=148} + %add.1091 = s32[1,41]{1,0} add(%broadcast_in_dim.1459, %broadcast.43), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/add" stack_frame_id=151} + %lt.609 = pred[1,41]{1,0} compare(%lt.608, %add.1091), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/lt" stack_frame_id=152} + %convert_element_type.1688 = s32[1,41]{1,0} convert(%lt.609), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/convert_element_type" stack_frame_id=156} + %broadcast_in_dim.1460 = s32[1,1,41]{2,1,0} reshape(%convert_element_type.1688), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/broadcast_in_dim" stack_frame_id=160} + %min.93 = s32[1,1,41]{2,1,0} minimum(%convert_element_type.1689, %broadcast_in_dim.1460), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/min" stack_frame_id=160} + %broadcast_in_dim.1465 = s32[1,1,1,41]{3,2,1,0} reshape(%min.93), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/broadcast_in_dim" stack_frame_id=341} + %convert_element_type.1699 = pred[1,1,1,41]{3,2,1,0} compare(%broadcast_in_dim.1465, %broadcast.37), direction=NE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/convert_element_type" stack_frame_id=345} + %broadcast_in_dim.1466 = pred[1,1,1,1,41]{4,3,2,1,0} reshape(%convert_element_type.1699), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/broadcast_in_dim" stack_frame_id=83} + %and.287 = pred[1,1,1,1,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1466), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/vmap()/and" stack_frame_id=83} + %and.288 = pred[1,1,1,41]{3,2,1,0} reshape(%and.287), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/vmap()/and" stack_frame_id=83} + %and.289 = pred[1,1,32,1,41]{4,3,2,1,0} broadcast(%and.288), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/vmap()/and" stack_frame_id=83} + %dot_general.1027 = bf16[1,1,32,128]{3,2,1,0} dot(%convert_element_type.1692, %arg_tuple.5#279), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=186} + %convert_element_type.1693 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/convert_element_type" stack_frame_id=208} + %add.1093 = f32[1,1]{1,0} reshape(%convert_element_type.1693), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/add"} + %iota.504 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/iota" stack_frame_id=197} + %mul.3015 = f32[64]{0} multiply(%iota.504, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=197} + %div.944 = f32[64]{0} divide(%mul.3015, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/div" stack_frame_id=198} + %neg.382 = f32[64]{0} negate(%div.944), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/neg" stack_frame_id=199} + %pow.376 = f32[64]{0} power(%broadcast.39, %neg.382), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/pow" stack_frame_id=202} + %div.945 = f32[64]{0} divide(%pow.376, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/div" stack_frame_id=203} + %dot_general.1028 = f32[1,1,64]{2,1,0} dot(%add.1093, %div.945), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/bi,j->bij/dot_general" stack_frame_id=211} + %stack.1470 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.1028), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/stack" stack_frame_id=214} + %stack.1471 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.1028), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/stack" stack_frame_id=214} + %stack.1472 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1470, %stack.1471), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/stack" stack_frame_id=214} + %reshape.684 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1472), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/reshape"} + %cos.187 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.684), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/cos" stack_frame_id=218} + %convert_element_type.1694 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.187), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/convert_element_type" stack_frame_id=222} + %mul.3016 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1694), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=242} + %mul.3017 = bf16[1,1,128]{2,1,0} reshape(%mul.3016), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=242} + %mul.3018 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.3017), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=242} + %mul.3019 = bf16[1,1,32,128]{3,2,1,0} multiply(%dot_general.1027, %mul.3018), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=242} + %split.375 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.1027), slice={[0:1], [0:1], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/split" stack_frame_id=234} + %neg.383 = bf16[1,1,32,64]{3,2,1,0} negate(%split.375), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/neg" stack_frame_id=235} + %stack.1473 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%neg.383), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/stack" stack_frame_id=238} + %split.374 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.1027), slice={[0:1], [0:1], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/split" stack_frame_id=234} + %stack.1474 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%split.374), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/stack" stack_frame_id=238} + %stack.1475 = bf16[1,1,32,2,64]{4,3,2,1,0} concatenate(%stack.1473, %stack.1474), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/stack" stack_frame_id=238} + %reshape.685 = bf16[1,1,32,128]{3,2,1,0} reshape(%stack.1475), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/reshape" stack_frame_id=241} + %sin.187 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.684), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/sin" stack_frame_id=226} + %convert_element_type.1695 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.187), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/convert_element_type" stack_frame_id=230} + %mul.3020 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1695), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=243} + %mul.3021 = bf16[1,1,128]{2,1,0} reshape(%mul.3020), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=243} + %mul.3022 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.3021), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=243} + %mul.3023 = bf16[1,1,32,128]{3,2,1,0} multiply(%reshape.685, %mul.3022), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=243} + %add.1094 = bf16[1,1,32,128]{3,2,1,0} add(%mul.3019, %mul.3023), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/add" stack_frame_id=244} + %reshape.690 = bf16[1,1,32,1,128]{4,3,2,1,0} reshape(%add.1094), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/reshape" stack_frame_id=83} + %slice.187 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.190), slice={[0:1], [0:1], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/slice" stack_frame_id=245} + %squeeze.191 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.187), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/squeeze" stack_frame_id=245} + %dot_general.1029 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1692, %arg_tuple.5#280), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=256} + %convert_element_type.1696 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/convert_element_type" stack_frame_id=287} + %add.1095 = f32[1,1]{1,0} reshape(%convert_element_type.1696), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/add"} + %iota.505 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/iota" stack_frame_id=276} + %mul.3024 = f32[64]{0} multiply(%iota.505, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=276} + %div.946 = f32[64]{0} divide(%mul.3024, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/div" stack_frame_id=277} + %neg.384 = f32[64]{0} negate(%div.946), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/neg" stack_frame_id=278} + %pow.377 = f32[64]{0} power(%broadcast.39, %neg.384), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/pow" stack_frame_id=281} + %div.947 = f32[64]{0} divide(%pow.377, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/div" stack_frame_id=282} + %dot_general.1031 = f32[1,1,64]{2,1,0} dot(%add.1095, %div.947), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/bi,j->bij/dot_general" stack_frame_id=290} + %stack.1476 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.1031), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/stack" stack_frame_id=293} + %stack.1477 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.1031), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/stack" stack_frame_id=293} + %stack.1478 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1476, %stack.1477), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/stack" stack_frame_id=293} + %reshape.686 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1478), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/reshape"} + %cos.188 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.686), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/cos" stack_frame_id=297} + %convert_element_type.1697 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.188), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/convert_element_type" stack_frame_id=301} + %mul.3025 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1697), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=321} + %mul.3026 = bf16[1,1,128]{2,1,0} reshape(%mul.3025), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=321} + %mul.3027 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.3026), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=321} + %mul.3028 = bf16[1,1,8,128]{3,2,1,0} multiply(%dot_general.1029, %mul.3027), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=321} + %split.377 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.1029), slice={[0:1], [0:1], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/split" stack_frame_id=313} + %neg.385 = bf16[1,1,8,64]{3,2,1,0} negate(%split.377), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/neg" stack_frame_id=314} + %stack.1479 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%neg.385), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/stack" stack_frame_id=317} + %split.376 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.1029), slice={[0:1], [0:1], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/split" stack_frame_id=313} + %stack.1480 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%split.376), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/stack" stack_frame_id=317} + %stack.1481 = bf16[1,1,8,2,64]{4,3,2,1,0} concatenate(%stack.1479, %stack.1480), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/stack" stack_frame_id=317} + %reshape.687 = bf16[1,1,8,128]{3,2,1,0} reshape(%stack.1481), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/reshape" stack_frame_id=320} + %sin.188 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.686), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/sin" stack_frame_id=305} + %convert_element_type.1698 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.188), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/convert_element_type" stack_frame_id=309} + %mul.3029 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1698), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=322} + %mul.3030 = bf16[1,1,128]{2,1,0} reshape(%mul.3029), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=322} + %mul.3031 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.3030), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=322} + %mul.3032 = bf16[1,1,8,128]{3,2,1,0} multiply(%reshape.687, %mul.3031), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=322} + %add.1096 = bf16[1,1,8,128]{3,2,1,0} add(%mul.3028, %mul.3032), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/rotary_embedding_30/add" stack_frame_id=323} + %lt.610 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/lt" stack_frame_id=326} + %add.1097 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/add" stack_frame_id=326} + %select_n.151 = s32[] select(%lt.610, %add.1097, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/select_n" stack_frame_id=326} + %dynamic_update_slice.125 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.191, %add.1096, %constant.127, %select_n.151, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/dynamic_update_slice" stack_frame_id=326} + %broadcast_in_dim.1463 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.125), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/broadcast_in_dim" stack_frame_id=335} + %reshape.688 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1463), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/reshape" stack_frame_id=335} + %dot_general.1032 = f32[1,32,1,1,41]{4,3,2,1,0} dot(%reshape.690, %reshape.688), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=83} + %mul.3033 = f32[1,32,1,1,41]{4,3,2,1,0} multiply(%dot_general.1032, %broadcast.36), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/vmap()/mul" stack_frame_id=83} + %vmap_jit__where__.93 = f32[1,1,32,1,41]{4,3,2,1,0} call(%and.289, %mul.3033, %constant.114), to_apply=%_where_3.137, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/vmap(jit(_where))" stack_frame_id=83} + %reduce_max.478 = f32[1,1,32,1]{3,2,1,0} reduce(%vmap_jit__where__.93, %constant.113), dimensions={4}, to_apply=%region_251.259, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/vmap()/reduce_max" stack_frame_id=83} + %max.102 = f32[1,1,32,1]{3,2,1,0} maximum(%reduce_max.478, %broadcast.35), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/vmap()/max" stack_frame_id=83} + %broadcast_in_dim.1467 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%max.102), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %sub.391 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1467), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/vmap()/sub" stack_frame_id=83} + %sub.392 = f32[1,1,32,1]{3,2,1,0} reshape(%sub.391), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/vmap()/sub" stack_frame_id=83} + %sub.393 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%sub.392), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/vmap()/sub" stack_frame_id=83} + %sub.394 = f32[1,1,32,1,41]{4,3,2,1,0} subtract(%vmap_jit__where__.93, %sub.393), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/vmap()/sub" stack_frame_id=83} + %exp.98 = f32[1,1,32,1,41]{4,3,2,1,0} exponential(%sub.394), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/vmap()/exp" stack_frame_id=83} + %reduce_sum.1442 = f32[1,1,32,1]{3,2,1,0} reduce(%exp.98, %constant.123), dimensions={4}, to_apply=%region_252.260, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/vmap()/reduce_sum" stack_frame_id=83} + %broadcast_in_dim.1468 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%reduce_sum.1442), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %div.948 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1468), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/vmap()/div" stack_frame_id=83} + %div.949 = f32[1,1,32,1]{3,2,1,0} reshape(%div.948), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/vmap()/div" stack_frame_id=83} + %div.950 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%div.949), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/vmap()/div" stack_frame_id=83} + %div.951 = f32[1,1,32,1,41]{4,3,2,1,0} divide(%exp.98, %div.950), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/vmap()/div" stack_frame_id=83} + %convert_element_type.1700 = bf16[1,1,32,1,41]{4,3,2,1,0} convert(%div.951), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/vmap()/convert_element_type" stack_frame_id=83} + %dot_general.1033 = bf16[1,32,128,1,1]{4,3,2,1,0} dot(%reshape.689, %convert_element_type.1700), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=83} + %transpose.97 = bf16[1,1,32,128]{3,2,1,0} reshape(%dot_general.1033), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/reshape;jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/vmap()/transpose;jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %dot_general.1034 = bf16[1,1,4096]{2,1,0} dot(%transpose.97, %arg_tuple.5#282), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=354} + %add.1099 = bf16[1,1,4096]{2,1,0} add(%dot_general.1034, %add.1088), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/add" stack_frame_id=355} + %convert_element_type.1701 = f32[1,1,4096]{2,1,0} convert(%add.1099), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/feedforward_layernorm/convert_element_type" stack_frame_id=359} + %pow.378 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1701, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/feedforward_layernorm/pow" stack_frame_id=362} + %reduce_sum.1443 = f32[1,1]{1,0} reduce(%pow.378, %constant.123), dimensions={2}, to_apply=%region_253.261, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/feedforward_layernorm/reduce_sum" stack_frame_id=92} + %broadcast_in_dim.1469 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1443), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/feedforward_layernorm/broadcast_in_dim" stack_frame_id=92} + %div.952 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1469, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/feedforward_layernorm/div" stack_frame_id=92} + %add.1100 = f32[1,1,1]{2,1,0} add(%div.952, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/feedforward_layernorm/add" stack_frame_id=363} + %rsqrt.189 = f32[1,1,1]{2,1,0} rsqrt(%add.1100), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/feedforward_layernorm/rsqrt" stack_frame_id=366} + %mul.3034 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.189), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/feedforward_layernorm/mul" stack_frame_id=367} + %mul.3035 = f32[1,1]{1,0} reshape(%mul.3034), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/feedforward_layernorm/mul" stack_frame_id=367} + %mul.3036 = f32[1,1,4096]{2,1,0} broadcast(%mul.3035), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/feedforward_layernorm/mul" stack_frame_id=367} + %mul.3037 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1701, %mul.3036), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/feedforward_layernorm/mul" stack_frame_id=367} + %convert_element_type.1702 = f32[4096]{0} convert(%arg_tuple.5#283), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/feedforward_layernorm/convert_element_type" stack_frame_id=368} + %broadcast_in_dim.1470 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1702), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/feedforward_layernorm/broadcast_in_dim" stack_frame_id=368} + %mul.3038 = f32[1,1,4096]{2,1,0} multiply(%mul.3037, %broadcast_in_dim.1470), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/feedforward_layernorm/mul" stack_frame_id=368} + %convert_element_type.1703 = bf16[1,1,4096]{2,1,0} convert(%mul.3038), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/feedforward_layernorm/convert_element_type" stack_frame_id=372} + %dot_general.1036 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1703, %arg_tuple.5#285), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/feedforward_intermediate_dense/dot_general" stack_frame_id=401} + %dot_general.1035 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1703, %arg_tuple.5#284), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/feedforward_gate_dense/dot_general" stack_frame_id=381} + %convert_element_type.1704 = f32[1,1,14336]{2,1,0} convert(%dot_general.1035), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/convert_element_type" stack_frame_id=385} + %jit_silu_.93 = f32[1,1,14336]{2,1,0} call(%convert_element_type.1704), to_apply=%silu_4.141, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/jit(silu)" stack_frame_id=388} + %convert_element_type.1705 = bf16[1,1,14336]{2,1,0} convert(%jit_silu_.93), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/convert_element_type" stack_frame_id=392} + %mul.3039 = bf16[1,1,14336]{2,1,0} multiply(%dot_general.1036, %convert_element_type.1705), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/mul" stack_frame_id=404} + %dot_general.1037 = bf16[1,1,4096]{2,1,0} dot(%mul.3039, %arg_tuple.5#286), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/feedforward_output_dense/dot_general" stack_frame_id=413} + %add.1101 = bf16[1,1,4096]{2,1,0} add(%dot_general.1037, %add.1099), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/add" stack_frame_id=414} + %convert_element_type.1709 = f32[1,1,4096]{2,1,0} convert(%add.1101), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention_layernorm/convert_element_type" stack_frame_id=164} + %pow.379 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1709, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention_layernorm/pow" stack_frame_id=167} + %reduce_sum.1444 = f32[1,1]{1,0} reduce(%pow.379, %constant.123), dimensions={2}, to_apply=%region_254.262, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention_layernorm/reduce_sum" stack_frame_id=73} + %broadcast_in_dim.1475 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1444), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention_layernorm/broadcast_in_dim" stack_frame_id=73} + %div.953 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1475, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention_layernorm/div" stack_frame_id=73} + %add.1105 = f32[1,1,1]{2,1,0} add(%div.953, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention_layernorm/add" stack_frame_id=168} + %rsqrt.190 = f32[1,1,1]{2,1,0} rsqrt(%add.1105), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention_layernorm/rsqrt" stack_frame_id=171} + %mul.3040 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.190), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention_layernorm/mul" stack_frame_id=172} + %mul.3041 = f32[1,1]{1,0} reshape(%mul.3040), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention_layernorm/mul" stack_frame_id=172} + %mul.3042 = f32[1,1,4096]{2,1,0} broadcast(%mul.3041), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention_layernorm/mul" stack_frame_id=172} + %mul.3043 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1709, %mul.3042), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention_layernorm/mul" stack_frame_id=172} + %convert_element_type.1710 = f32[4096]{0} convert(%arg_tuple.5#287), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention_layernorm/convert_element_type" stack_frame_id=173} + %broadcast_in_dim.1476 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1710), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention_layernorm/broadcast_in_dim" stack_frame_id=173} + %mul.3044 = f32[1,1,4096]{2,1,0} multiply(%mul.3043, %broadcast_in_dim.1476), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention_layernorm/mul" stack_frame_id=173} + %convert_element_type.1711 = bf16[1,1,4096]{2,1,0} convert(%mul.3044), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention_layernorm/convert_element_type" stack_frame_id=177} + %dot_general.1041 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1711, %arg_tuple.5#290), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=265} + %lt.617 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/lt" stack_frame_id=329} + %add.1111 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/add" stack_frame_id=329} + %select_n.154 = s32[] select(%lt.617, %add.1111, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/select_n" stack_frame_id=329} + %dynamic_update_slice.128 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.195, %dot_general.1041, %constant.127, %select_n.154, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/dynamic_update_slice" stack_frame_id=329} + %broadcast_in_dim.1478 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.128), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/broadcast_in_dim" stack_frame_id=338} + %reshape.696 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1478), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/reshape" stack_frame_id=338} + %convert_element_type.1706 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/convert_element_type" stack_frame_id=139} + %add.1102 = f32[1,1]{1,0} reshape(%convert_element_type.1706), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_31/add"} + %ge.477 = f32[1,1]{1,0} broadcast(%add.1102), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/ge" stack_frame_id=143} + %ge.478 = f32[1]{0} reshape(%ge.477), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/ge" stack_frame_id=143} + %ge.479 = f32[1,41]{1,0} broadcast(%ge.478), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/ge" stack_frame_id=143} + %iota.506 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/iota" stack_frame_id=142} + %broadcast_in_dim.1471 = f32[1,41]{1,0} reshape(%iota.506), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/broadcast_in_dim" stack_frame_id=143} + %ge.480 = pred[1,41]{1,0} compare(%ge.479, %broadcast_in_dim.1471), direction=GE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/ge" stack_frame_id=143} + %broadcast_in_dim.1472 = pred[1,1,41]{2,1,0} reshape(%ge.480), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/broadcast_in_dim" stack_frame_id=146} + %convert_element_type.1708 = s32[1,1,41]{2,1,0} convert(%broadcast_in_dim.1472), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/convert_element_type" stack_frame_id=160} + %add.1103 = s32[1,1]{1,0} reshape(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/add" stack_frame_id=147} + %lt.612 = s32[1,1]{1,0} broadcast(%add.1103), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/lt" stack_frame_id=152} + %lt.613 = s32[1]{0} reshape(%lt.612), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/lt" stack_frame_id=152} + %lt.614 = s32[1,41]{1,0} broadcast(%lt.613), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/lt" stack_frame_id=152} + %iota.507 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/iota" stack_frame_id=150} + %broadcast_in_dim.1473 = s32[1,41]{1,0} reshape(%iota.507), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/broadcast_in_dim" stack_frame_id=148} + %add.1104 = s32[1,41]{1,0} add(%broadcast_in_dim.1473, %broadcast.43), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/add" stack_frame_id=151} + %lt.615 = pred[1,41]{1,0} compare(%lt.614, %add.1104), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/lt" stack_frame_id=152} + %convert_element_type.1707 = s32[1,41]{1,0} convert(%lt.615), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/convert_element_type" stack_frame_id=156} + %broadcast_in_dim.1474 = s32[1,1,41]{2,1,0} reshape(%convert_element_type.1707), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/broadcast_in_dim" stack_frame_id=160} + %min.94 = s32[1,1,41]{2,1,0} minimum(%convert_element_type.1708, %broadcast_in_dim.1474), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/min" stack_frame_id=160} + %broadcast_in_dim.1479 = s32[1,1,1,41]{3,2,1,0} reshape(%min.94), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/broadcast_in_dim" stack_frame_id=341} + %convert_element_type.1718 = pred[1,1,1,41]{3,2,1,0} compare(%broadcast_in_dim.1479, %broadcast.37), direction=NE, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/convert_element_type" stack_frame_id=345} + %broadcast_in_dim.1480 = pred[1,1,1,1,41]{4,3,2,1,0} reshape(%convert_element_type.1718), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/broadcast_in_dim" stack_frame_id=83} + %and.290 = pred[1,1,1,1,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1480), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/vmap()/and" stack_frame_id=83} + %and.291 = pred[1,1,1,41]{3,2,1,0} reshape(%and.290), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/vmap()/and" stack_frame_id=83} + %and.292 = pred[1,1,32,1,41]{4,3,2,1,0} broadcast(%and.291), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/vmap()/and" stack_frame_id=83} + %dot_general.1038 = bf16[1,1,32,128]{3,2,1,0} dot(%convert_element_type.1711, %arg_tuple.5#288), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=186} + %convert_element_type.1712 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/convert_element_type" stack_frame_id=208} + %add.1106 = f32[1,1]{1,0} reshape(%convert_element_type.1712), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/add"} + %iota.508 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/iota" stack_frame_id=197} + %mul.3045 = f32[64]{0} multiply(%iota.508, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/mul" stack_frame_id=197} + %div.954 = f32[64]{0} divide(%mul.3045, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/div" stack_frame_id=198} + %neg.386 = f32[64]{0} negate(%div.954), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/neg" stack_frame_id=199} + %pow.380 = f32[64]{0} power(%broadcast.39, %neg.386), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/pow" stack_frame_id=202} + %div.955 = f32[64]{0} divide(%pow.380, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/div" stack_frame_id=203} + %dot_general.1039 = f32[1,1,64]{2,1,0} dot(%add.1106, %div.955), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/bi,j->bij/dot_general" stack_frame_id=211} + %stack.1485 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.1039), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/stack" stack_frame_id=214} + %stack.1486 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.1039), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/stack" stack_frame_id=214} + %stack.1487 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1485, %stack.1486), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/stack" stack_frame_id=214} + %reshape.691 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1487), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/reshape"} + %cos.189 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.691), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/cos" stack_frame_id=218} + %convert_element_type.1713 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.189), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/convert_element_type" stack_frame_id=222} + %mul.3046 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1713), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/mul" stack_frame_id=242} + %mul.3047 = bf16[1,1,128]{2,1,0} reshape(%mul.3046), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/mul" stack_frame_id=242} + %mul.3048 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.3047), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/mul" stack_frame_id=242} + %mul.3049 = bf16[1,1,32,128]{3,2,1,0} multiply(%dot_general.1038, %mul.3048), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/mul" stack_frame_id=242} + %split.379 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.1038), slice={[0:1], [0:1], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/split" stack_frame_id=234} + %neg.387 = bf16[1,1,32,64]{3,2,1,0} negate(%split.379), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/neg" stack_frame_id=235} + %stack.1488 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%neg.387), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/stack" stack_frame_id=238} + %split.378 = bf16[1,1,32,64]{3,2,1,0} slice(%dot_general.1038), slice={[0:1], [0:1], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/split" stack_frame_id=234} + %stack.1489 = bf16[1,1,32,1,64]{4,3,2,1,0} reshape(%split.378), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/stack" stack_frame_id=238} + %stack.1490 = bf16[1,1,32,2,64]{4,3,2,1,0} concatenate(%stack.1488, %stack.1489), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/stack" stack_frame_id=238} + %reshape.692 = bf16[1,1,32,128]{3,2,1,0} reshape(%stack.1490), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/reshape" stack_frame_id=241} + %sin.189 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.691), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/sin" stack_frame_id=226} + %convert_element_type.1714 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.189), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/convert_element_type" stack_frame_id=230} + %mul.3050 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1714), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/mul" stack_frame_id=243} + %mul.3051 = bf16[1,1,128]{2,1,0} reshape(%mul.3050), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/mul" stack_frame_id=243} + %mul.3052 = bf16[1,1,32,128]{3,2,1,0} broadcast(%mul.3051), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/mul" stack_frame_id=243} + %mul.3053 = bf16[1,1,32,128]{3,2,1,0} multiply(%reshape.692, %mul.3052), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/mul" stack_frame_id=243} + %add.1107 = bf16[1,1,32,128]{3,2,1,0} add(%mul.3049, %mul.3053), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/add" stack_frame_id=244} + %reshape.697 = bf16[1,1,32,1,128]{4,3,2,1,0} reshape(%add.1107), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/reshape" stack_frame_id=83} + %slice.190 = bf16[1,1,41,8,128]{4,3,2,1,0} slice(%squeeze.193), slice={[0:1], [0:1], [0:41], [0:8], [0:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/slice" stack_frame_id=245} + %squeeze.194 = bf16[1,41,8,128]{3,2,1,0} reshape(%slice.190), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/squeeze" stack_frame_id=245} + %dot_general.1040 = bf16[1,1,8,128]{3,2,1,0} dot(%convert_element_type.1711, %arg_tuple.5#289), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=256} + %convert_element_type.1715 = f32[] convert(%sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/convert_element_type" stack_frame_id=287} + %add.1108 = f32[1,1]{1,0} reshape(%convert_element_type.1715), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/add"} + %iota.509 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/iota" stack_frame_id=276} + %mul.3054 = f32[64]{0} multiply(%iota.509, %broadcast.41), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/mul" stack_frame_id=276} + %div.956 = f32[64]{0} divide(%mul.3054, %broadcast.40), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/div" stack_frame_id=277} + %neg.388 = f32[64]{0} negate(%div.956), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/neg" stack_frame_id=278} + %pow.381 = f32[64]{0} power(%broadcast.39, %neg.388), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/pow" stack_frame_id=281} + %div.957 = f32[64]{0} divide(%pow.381, %broadcast.38), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/div" stack_frame_id=282} + %dot_general.1042 = f32[1,1,64]{2,1,0} dot(%add.1108, %div.957), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/bi,j->bij/dot_general" stack_frame_id=290} + %stack.1491 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.1042), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/stack" stack_frame_id=293} + %stack.1492 = f32[1,1,1,64]{3,2,1,0} reshape(%dot_general.1042), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/stack" stack_frame_id=293} + %stack.1493 = f32[1,1,2,64]{3,2,1,0} concatenate(%stack.1491, %stack.1492), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/stack" stack_frame_id=293} + %reshape.693 = f32[1,1,1,128]{3,2,1,0} reshape(%stack.1493), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/broadcast_in_dim;jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/reshape"} + %cos.190 = f32[1,1,1,128]{3,2,1,0} cosine(%reshape.693), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/cos" stack_frame_id=297} + %convert_element_type.1716 = bf16[1,1,1,128]{3,2,1,0} convert(%cos.190), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/convert_element_type" stack_frame_id=301} + %mul.3055 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1716), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/mul" stack_frame_id=321} + %mul.3056 = bf16[1,1,128]{2,1,0} reshape(%mul.3055), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/mul" stack_frame_id=321} + %mul.3057 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.3056), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/mul" stack_frame_id=321} + %mul.3058 = bf16[1,1,8,128]{3,2,1,0} multiply(%dot_general.1040, %mul.3057), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/mul" stack_frame_id=321} + %split.381 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.1040), slice={[0:1], [0:1], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/split" stack_frame_id=313} + %neg.389 = bf16[1,1,8,64]{3,2,1,0} negate(%split.381), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/neg" stack_frame_id=314} + %stack.1494 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%neg.389), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/stack" stack_frame_id=317} + %split.380 = bf16[1,1,8,64]{3,2,1,0} slice(%dot_general.1040), slice={[0:1], [0:1], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/split" stack_frame_id=313} + %stack.1495 = bf16[1,1,8,1,64]{4,3,2,1,0} reshape(%split.380), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/stack" stack_frame_id=317} + %stack.1496 = bf16[1,1,8,2,64]{4,3,2,1,0} concatenate(%stack.1494, %stack.1495), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/stack" stack_frame_id=317} + %reshape.694 = bf16[1,1,8,128]{3,2,1,0} reshape(%stack.1496), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/reshape" stack_frame_id=320} + %sin.190 = f32[1,1,1,128]{3,2,1,0} sine(%reshape.693), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/sin" stack_frame_id=305} + %convert_element_type.1717 = bf16[1,1,1,128]{3,2,1,0} convert(%sin.190), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/convert_element_type" stack_frame_id=309} + %mul.3059 = bf16[1,1,1,128]{3,2,1,0} broadcast(%convert_element_type.1717), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/mul" stack_frame_id=322} + %mul.3060 = bf16[1,1,128]{2,1,0} reshape(%mul.3059), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/mul" stack_frame_id=322} + %mul.3061 = bf16[1,1,8,128]{3,2,1,0} broadcast(%mul.3060), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/mul" stack_frame_id=322} + %mul.3062 = bf16[1,1,8,128]{3,2,1,0} multiply(%reshape.694, %mul.3061), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/mul" stack_frame_id=322} + %add.1109 = bf16[1,1,8,128]{3,2,1,0} add(%mul.3058, %mul.3062), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/rotary_embedding_31/add" stack_frame_id=323} + %lt.616 = pred[] compare(%sub.270, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/lt" stack_frame_id=326} + %add.1110 = s32[] add(%sub.270, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/add" stack_frame_id=326} + %select_n.153 = s32[] select(%lt.616, %add.1110, %sub.270), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/select_n" stack_frame_id=326} + %dynamic_update_slice.127 = bf16[1,41,8,128]{3,2,1,0} dynamic-update-slice(%squeeze.194, %add.1109, %constant.127, %select_n.153, %constant.127, /*index=5*/%constant.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/dynamic_update_slice" stack_frame_id=326} + %broadcast_in_dim.1477 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dynamic_update_slice.127), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/broadcast_in_dim" stack_frame_id=335} + %reshape.695 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1477), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/reshape" stack_frame_id=335} + %dot_general.1043 = f32[1,32,1,1,41]{4,3,2,1,0} dot(%reshape.697, %reshape.695), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=83} + %mul.3063 = f32[1,32,1,1,41]{4,3,2,1,0} multiply(%dot_general.1043, %broadcast.36), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/vmap()/mul" stack_frame_id=83} + %vmap_jit__where__.94 = f32[1,1,32,1,41]{4,3,2,1,0} call(%and.292, %mul.3063, %constant.114), to_apply=%_where_3.137, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/vmap(jit(_where))" stack_frame_id=83} + %reduce_max.479 = f32[1,1,32,1]{3,2,1,0} reduce(%vmap_jit__where__.94, %constant.113), dimensions={4}, to_apply=%region_255.263, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/vmap()/reduce_max" stack_frame_id=83} + %max.103 = f32[1,1,32,1]{3,2,1,0} maximum(%reduce_max.479, %broadcast.35), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/vmap()/max" stack_frame_id=83} + %broadcast_in_dim.1481 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%max.103), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %sub.395 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1481), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/vmap()/sub" stack_frame_id=83} + %sub.396 = f32[1,1,32,1]{3,2,1,0} reshape(%sub.395), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/vmap()/sub" stack_frame_id=83} + %sub.397 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%sub.396), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/vmap()/sub" stack_frame_id=83} + %sub.398 = f32[1,1,32,1,41]{4,3,2,1,0} subtract(%vmap_jit__where__.94, %sub.397), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/vmap()/sub" stack_frame_id=83} + %exp.99 = f32[1,1,32,1,41]{4,3,2,1,0} exponential(%sub.398), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/vmap()/exp" stack_frame_id=83} + %reduce_sum.1445 = f32[1,1,32,1]{3,2,1,0} reduce(%exp.99, %constant.123), dimensions={4}, to_apply=%region_256.264, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/vmap()/reduce_sum" stack_frame_id=83} + %broadcast_in_dim.1482 = f32[1,1,32,1,1]{4,3,2,1,0} reshape(%reduce_sum.1445), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/vmap()/broadcast_in_dim" stack_frame_id=83} + %div.958 = f32[1,1,32,1,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1482), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/vmap()/div" stack_frame_id=83} + %div.959 = f32[1,1,32,1]{3,2,1,0} reshape(%div.958), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/vmap()/div" stack_frame_id=83} + %div.960 = f32[1,1,32,1,41]{4,3,2,1,0} broadcast(%div.959), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/vmap()/div" stack_frame_id=83} + %div.961 = f32[1,1,32,1,41]{4,3,2,1,0} divide(%exp.99, %div.960), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/vmap()/div" stack_frame_id=83} + %convert_element_type.1719 = bf16[1,1,32,1,41]{4,3,2,1,0} convert(%div.961), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/vmap()/convert_element_type" stack_frame_id=83} + %dot_general.1044 = bf16[1,32,128,1,1]{4,3,2,1,0} dot(%reshape.696, %convert_element_type.1719), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=83} + %transpose.98 = bf16[1,1,32,128]{3,2,1,0} reshape(%dot_general.1044), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/reshape;jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/vmap()/transpose;jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %dot_general.1045 = bf16[1,1,4096]{2,1,0} dot(%transpose.98, %arg_tuple.5#291), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=354} + %add.1112 = bf16[1,1,4096]{2,1,0} add(%dot_general.1045, %add.1101), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/add" stack_frame_id=355} + %convert_element_type.1720 = f32[1,1,4096]{2,1,0} convert(%add.1112), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/feedforward_layernorm/convert_element_type" stack_frame_id=359} + %pow.382 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1720, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/feedforward_layernorm/pow" stack_frame_id=362} + %reduce_sum.1446 = f32[1,1]{1,0} reduce(%pow.382, %constant.123), dimensions={2}, to_apply=%region_257.265, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/feedforward_layernorm/reduce_sum" stack_frame_id=92} + %broadcast_in_dim.1483 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1446), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/feedforward_layernorm/broadcast_in_dim" stack_frame_id=92} + %div.962 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1483, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/feedforward_layernorm/div" stack_frame_id=92} + %add.1113 = f32[1,1,1]{2,1,0} add(%div.962, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/feedforward_layernorm/add" stack_frame_id=363} + %rsqrt.191 = f32[1,1,1]{2,1,0} rsqrt(%add.1113), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/feedforward_layernorm/rsqrt" stack_frame_id=366} + %mul.3064 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.191), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/feedforward_layernorm/mul" stack_frame_id=367} + %mul.3065 = f32[1,1]{1,0} reshape(%mul.3064), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/feedforward_layernorm/mul" stack_frame_id=367} + %mul.3066 = f32[1,1,4096]{2,1,0} broadcast(%mul.3065), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/feedforward_layernorm/mul" stack_frame_id=367} + %mul.3067 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1720, %mul.3066), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/feedforward_layernorm/mul" stack_frame_id=367} + %convert_element_type.1721 = f32[4096]{0} convert(%arg_tuple.5#292), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/feedforward_layernorm/convert_element_type" stack_frame_id=368} + %broadcast_in_dim.1484 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1721), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/feedforward_layernorm/broadcast_in_dim" stack_frame_id=368} + %mul.3068 = f32[1,1,4096]{2,1,0} multiply(%mul.3067, %broadcast_in_dim.1484), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/feedforward_layernorm/mul" stack_frame_id=368} + %convert_element_type.1722 = bf16[1,1,4096]{2,1,0} convert(%mul.3068), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/feedforward_layernorm/convert_element_type" stack_frame_id=372} + %dot_general.1047 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1722, %arg_tuple.5#294), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/feedforward_intermediate_dense/dot_general" stack_frame_id=401} + %dot_general.1046 = bf16[1,1,14336]{2,1,0} dot(%convert_element_type.1722, %arg_tuple.5#293), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/feedforward_gate_dense/dot_general" stack_frame_id=381} + %convert_element_type.1723 = f32[1,1,14336]{2,1,0} convert(%dot_general.1046), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/convert_element_type" stack_frame_id=385} + %jit_silu_.94 = f32[1,1,14336]{2,1,0} call(%convert_element_type.1723), to_apply=%silu_4.141, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/jit(silu)" stack_frame_id=388} + %convert_element_type.1724 = bf16[1,1,14336]{2,1,0} convert(%jit_silu_.94), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/convert_element_type" stack_frame_id=392} + %mul.3069 = bf16[1,1,14336]{2,1,0} multiply(%dot_general.1047, %convert_element_type.1724), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/mul" stack_frame_id=404} + %dot_general.1048 = bf16[1,1,4096]{2,1,0} dot(%mul.3069, %arg_tuple.5#295), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/feedforward_output_dense/dot_general" stack_frame_id=413} + %add.1114 = bf16[1,1,4096]{2,1,0} add(%dot_general.1048, %add.1112), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/add" stack_frame_id=414} + %convert_element_type.1725 = f32[1,1,4096]{2,1,0} convert(%add.1114), metadata={op_name="jit(compiled_generate_function)/while/body/sequence_output_layernorm/convert_element_type" stack_frame_id=421} + %pow.383 = f32[1,1,4096]{2,1,0} power(%convert_element_type.1725, %broadcast.42), metadata={op_name="jit(compiled_generate_function)/while/body/sequence_output_layernorm/pow" stack_frame_id=424} + %reduce_sum.1447 = f32[1,1]{1,0} reduce(%pow.383, %constant.123), dimensions={2}, to_apply=%region_258.266, metadata={op_name="jit(compiled_generate_function)/while/body/sequence_output_layernorm/reduce_sum" stack_frame_id=101} + %broadcast_in_dim.1485 = f32[1,1,1]{2,1,0} reshape(%reduce_sum.1447), metadata={op_name="jit(compiled_generate_function)/while/body/sequence_output_layernorm/broadcast_in_dim" stack_frame_id=101} + %div.963 = f32[1,1,1]{2,1,0} divide(%broadcast_in_dim.1485, %constant.122), metadata={op_name="jit(compiled_generate_function)/while/body/sequence_output_layernorm/div" stack_frame_id=101} + %add.1115 = f32[1,1,1]{2,1,0} add(%div.963, %constant.121), metadata={op_name="jit(compiled_generate_function)/while/body/sequence_output_layernorm/add" stack_frame_id=425} + %rsqrt.192 = f32[1,1,1]{2,1,0} rsqrt(%add.1115), metadata={op_name="jit(compiled_generate_function)/while/body/sequence_output_layernorm/rsqrt" stack_frame_id=428} + %mul.3070 = f32[1,1,1]{2,1,0} broadcast(%rsqrt.192), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/while/body/sequence_output_layernorm/mul" stack_frame_id=429} + %mul.3071 = f32[1,1]{1,0} reshape(%mul.3070), metadata={op_name="jit(compiled_generate_function)/while/body/sequence_output_layernorm/mul" stack_frame_id=429} + %mul.3072 = f32[1,1,4096]{2,1,0} broadcast(%mul.3071), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/sequence_output_layernorm/mul" stack_frame_id=429} + %mul.3073 = f32[1,1,4096]{2,1,0} multiply(%convert_element_type.1725, %mul.3072), metadata={op_name="jit(compiled_generate_function)/while/body/sequence_output_layernorm/mul" stack_frame_id=429} + %convert_element_type.1726 = f32[4096]{0} convert(%arg_tuple.5#296), metadata={op_name="jit(compiled_generate_function)/while/body/sequence_output_layernorm/convert_element_type" stack_frame_id=430} + %broadcast_in_dim.1486 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1726), metadata={op_name="jit(compiled_generate_function)/while/body/sequence_output_layernorm/broadcast_in_dim" stack_frame_id=430} + %mul.3074 = f32[1,1,4096]{2,1,0} multiply(%mul.3073, %broadcast_in_dim.1486), metadata={op_name="jit(compiled_generate_function)/while/body/sequence_output_layernorm/mul" stack_frame_id=430} + %convert_element_type.1727 = bf16[1,1,4096]{2,1,0} convert(%mul.3074), metadata={op_name="jit(compiled_generate_function)/while/body/sequence_output_layernorm/convert_element_type" stack_frame_id=434} + %dot_general.1049 = bf16[1,1,32000]{2,1,0} dot(%convert_element_type.1727, %arg_tuple.5#297), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/while/body/token_embedding/dot_general" stack_frame_id=443} + %squeeze.196 = bf16[1,32000]{1,0} reshape(%dot_general.1049), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=446} + %convert_element_type.1728 = f32[1,32000]{1,0} convert(%squeeze.196), metadata={op_name="jit(compiled_generate_function)/while/body/convert_element_type" stack_frame_id=450} + %constant.111 = f32[] constant(1) + %broadcast.34 = f32[1,32000]{1,0} broadcast(%constant.111), dimensions={} + %div.964 = f32[1,32000]{1,0} divide(%convert_element_type.1728, %broadcast.34), metadata={op_name="jit(compiled_generate_function)/while/body/div" stack_frame_id=451} + %reduce_max.480 = f32[1]{0} reduce(%div.964, %constant.113), dimensions={1}, to_apply=%region_259.267, metadata={op_name="jit(compiled_generate_function)/while/body/reduce_max" stack_frame_id=106} + %constant.110 = f32[1]{0} constant({-inf}) + %max.104 = f32[1]{0} maximum(%reduce_max.480, %constant.110), metadata={op_name="jit(compiled_generate_function)/while/body/max" stack_frame_id=106} + %broadcast_in_dim.1487 = f32[1,1]{1,0} reshape(%max.104), metadata={op_name="jit(compiled_generate_function)/while/body/broadcast_in_dim" stack_frame_id=106} + %sub.399 = f32[1,1]{1,0} broadcast(%broadcast_in_dim.1487), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/sub" stack_frame_id=106} + %sub.400 = f32[1]{0} reshape(%sub.399), metadata={op_name="jit(compiled_generate_function)/while/body/sub" stack_frame_id=106} + %sub.401 = f32[1,32000]{1,0} broadcast(%sub.400), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/sub" stack_frame_id=106} + %sub.402 = f32[1,32000]{1,0} subtract(%div.964, %sub.401), metadata={op_name="jit(compiled_generate_function)/while/body/sub" stack_frame_id=106} + %exp.100 = f32[1,32000]{1,0} exponential(%sub.402), metadata={op_name="jit(compiled_generate_function)/while/body/exp" stack_frame_id=106} + %reduce_sum.1448 = f32[1]{0} reduce(%exp.100, %constant.123), dimensions={1}, to_apply=%region_260.268, metadata={op_name="jit(compiled_generate_function)/while/body/reduce_sum" stack_frame_id=106} + %broadcast_in_dim.1488 = f32[1,1]{1,0} reshape(%reduce_sum.1448), metadata={op_name="jit(compiled_generate_function)/while/body/broadcast_in_dim" stack_frame_id=106} + %div.965 = f32[1,1]{1,0} broadcast(%broadcast_in_dim.1488), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/while/body/div" stack_frame_id=106} + %div.966 = f32[1]{0} reshape(%div.965), metadata={op_name="jit(compiled_generate_function)/while/body/div" stack_frame_id=106} + %div.967 = f32[1,32000]{1,0} broadcast(%div.966), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/while/body/div" stack_frame_id=106} + %div.968 = f32[1,32000]{1,0} divide(%exp.100, %div.967), metadata={op_name="jit(compiled_generate_function)/while/body/div" stack_frame_id=106} + %top_k.3 = (f32[1,5]{1,0}, s32[1,5]{1,0}) topk(%div.968), k=5, largest=true, is_stable=true, metadata={op_name="jit(compiled_generate_function)/while/body/top_k" stack_frame_id=455} + %jit__gumbel_.1 = f32[1,5,1]{2,1,0} call(%arg_tuple.5#0), to_apply=%_gumbel.274, metadata={op_name="jit(compiled_generate_function)/while/body/jit(_gumbel)" stack_frame_id=469} + %log.5 = f32[1,5]{1,0} log(%top_k.3#0), metadata={op_name="jit(compiled_generate_function)/while/body/log" stack_frame_id=459} + %broadcast_in_dim.1489 = f32[1,5,1]{2,1,0} reshape(%log.5), metadata={op_name="jit(compiled_generate_function)/while/body/broadcast_in_dim" stack_frame_id=468} + %add.1117 = f32[1,5,1]{2,1,0} add(%jit__gumbel_.1, %broadcast_in_dim.1489), metadata={op_name="jit(compiled_generate_function)/while/body/add" stack_frame_id=469} + %body.3 = s32[1,1]{1,0} call(%add.1117), to_apply=%argmax.276, metadata={op_name="jit(compiled_generate_function)/while/body" stack_frame_id=469} + %jit_take_along_axis_.1 = s32[1,1]{1,0} call(%top_k.3#1, %body.3), to_apply=%take_along_axis.278, metadata={op_name="jit(compiled_generate_function)/while/body/jit(take_along_axis)" stack_frame_id=472} + %squeeze.197 = s32[1]{0} reshape(%jit_take_along_axis_.1), metadata={op_name="jit(compiled_generate_function)/while/body/squeeze" stack_frame_id=475} + %jit__where_.5 = s32[1]{0} call(%squeeze.198, %squeeze.199, %squeeze.197), to_apply=%_where_5.279, metadata={op_name="jit(compiled_generate_function)/while/body/jit(_where)" stack_frame_id=480} + %broadcast_in_dim.1490 = s32[1,1]{1,0} reshape(%jit__where_.5), metadata={op_name="jit(compiled_generate_function)/while/body/broadcast_in_dim" stack_frame_id=481} + %lt.620 = pred[] compare(%arg_tuple.5#3, %constant.127), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/body/lt" stack_frame_id=484} + %add.1120 = s32[] add(%arg_tuple.5#3, %constant.126), metadata={op_name="jit(compiled_generate_function)/while/body/add" stack_frame_id=484} + %select_n.157 = s32[] select(%lt.620, %add.1120, %arg_tuple.5#3), metadata={op_name="jit(compiled_generate_function)/while/body/select_n" stack_frame_id=484} + %dynamic_update_slice.129 = s32[1,41]{1,0} dynamic-update-slice(%arg_tuple.5#1, %broadcast_in_dim.1490, %constant.127, %select_n.157), metadata={op_name="jit(compiled_generate_function)/while/body/dynamic_update_slice" stack_frame_id=484} + %stack.1032 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.65), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/stack" stack_frame_id=332} + %stack.1033 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.66), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/stack" stack_frame_id=332} + %stack.1034 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1032, %stack.1033), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_0/self_attention/stack" stack_frame_id=332} + %stack.1500 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1034), metadata={op_name="jit(compiled_generate_function)/while/body/stack" stack_frame_id=417} + %stack.1047 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.67), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/stack" stack_frame_id=332} + %stack.1048 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.68), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/stack" stack_frame_id=332} + %stack.1049 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1047, %stack.1048), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_1/self_attention/stack" stack_frame_id=332} + %stack.1501 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1049), metadata={op_name="jit(compiled_generate_function)/while/body/stack" stack_frame_id=417} + %stack.1062 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.69), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/stack" stack_frame_id=332} + %stack.1063 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.70), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/stack" stack_frame_id=332} + %stack.1064 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1062, %stack.1063), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_2/self_attention/stack" stack_frame_id=332} + %stack.1502 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1064), metadata={op_name="jit(compiled_generate_function)/while/body/stack" stack_frame_id=417} + %stack.1077 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.71), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/stack" stack_frame_id=332} + %stack.1078 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.72), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/stack" stack_frame_id=332} + %stack.1079 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1077, %stack.1078), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_3/self_attention/stack" stack_frame_id=332} + %stack.1503 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1079), metadata={op_name="jit(compiled_generate_function)/while/body/stack" stack_frame_id=417} + %stack.1092 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.73), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/stack" stack_frame_id=332} + %stack.1093 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.74), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/stack" stack_frame_id=332} + %stack.1094 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1092, %stack.1093), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_4/self_attention/stack" stack_frame_id=332} + %stack.1504 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1094), metadata={op_name="jit(compiled_generate_function)/while/body/stack" stack_frame_id=417} + %stack.1107 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.75), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/stack" stack_frame_id=332} + %stack.1108 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.76), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/stack" stack_frame_id=332} + %stack.1109 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1107, %stack.1108), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_5/self_attention/stack" stack_frame_id=332} + %stack.1505 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1109), metadata={op_name="jit(compiled_generate_function)/while/body/stack" stack_frame_id=417} + %stack.1122 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.77), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/stack" stack_frame_id=332} + %stack.1123 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.78), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/stack" stack_frame_id=332} + %stack.1124 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1122, %stack.1123), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_6/self_attention/stack" stack_frame_id=332} + %stack.1506 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1124), metadata={op_name="jit(compiled_generate_function)/while/body/stack" stack_frame_id=417} + %stack.1137 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.79), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/stack" stack_frame_id=332} + %stack.1138 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.80), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/stack" stack_frame_id=332} + %stack.1139 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1137, %stack.1138), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_7/self_attention/stack" stack_frame_id=332} + %stack.1507 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1139), metadata={op_name="jit(compiled_generate_function)/while/body/stack" stack_frame_id=417} + %stack.1152 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.81), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/stack" stack_frame_id=332} + %stack.1153 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.82), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/stack" stack_frame_id=332} + %stack.1154 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1152, %stack.1153), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_8/self_attention/stack" stack_frame_id=332} + %stack.1508 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1154), metadata={op_name="jit(compiled_generate_function)/while/body/stack" stack_frame_id=417} + %stack.1167 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.83), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/stack" stack_frame_id=332} + %stack.1168 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.84), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/stack" stack_frame_id=332} + %stack.1169 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1167, %stack.1168), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_9/self_attention/stack" stack_frame_id=332} + %stack.1509 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1169), metadata={op_name="jit(compiled_generate_function)/while/body/stack" stack_frame_id=417} + %stack.1182 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.85), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/stack" stack_frame_id=332} + %stack.1183 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.86), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/stack" stack_frame_id=332} + %stack.1184 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1182, %stack.1183), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_10/self_attention/stack" stack_frame_id=332} + %stack.1510 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1184), metadata={op_name="jit(compiled_generate_function)/while/body/stack" stack_frame_id=417} + %stack.1197 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.87), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/stack" stack_frame_id=332} + %stack.1198 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.88), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/stack" stack_frame_id=332} + %stack.1199 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1197, %stack.1198), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_11/self_attention/stack" stack_frame_id=332} + %stack.1511 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1199), metadata={op_name="jit(compiled_generate_function)/while/body/stack" stack_frame_id=417} + %stack.1212 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.89), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/stack" stack_frame_id=332} + %stack.1213 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.90), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/stack" stack_frame_id=332} + %stack.1214 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1212, %stack.1213), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_12/self_attention/stack" stack_frame_id=332} + %stack.1512 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1214), metadata={op_name="jit(compiled_generate_function)/while/body/stack" stack_frame_id=417} + %stack.1227 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.91), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/stack" stack_frame_id=332} + %stack.1228 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.92), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/stack" stack_frame_id=332} + %stack.1229 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1227, %stack.1228), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_13/self_attention/stack" stack_frame_id=332} + %stack.1513 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1229), metadata={op_name="jit(compiled_generate_function)/while/body/stack" stack_frame_id=417} + %stack.1242 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.93), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/stack" stack_frame_id=332} + %stack.1243 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.94), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/stack" stack_frame_id=332} + %stack.1244 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1242, %stack.1243), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_14/self_attention/stack" stack_frame_id=332} + %stack.1514 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1244), metadata={op_name="jit(compiled_generate_function)/while/body/stack" stack_frame_id=417} + %stack.1257 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.95), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/stack" stack_frame_id=332} + %stack.1258 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.96), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/stack" stack_frame_id=332} + %stack.1259 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1257, %stack.1258), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_15/self_attention/stack" stack_frame_id=332} + %stack.1515 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1259), metadata={op_name="jit(compiled_generate_function)/while/body/stack" stack_frame_id=417} + %stack.1272 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.97), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/stack" stack_frame_id=332} + %stack.1273 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.98), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/stack" stack_frame_id=332} + %stack.1274 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1272, %stack.1273), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_16/self_attention/stack" stack_frame_id=332} + %stack.1516 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1274), metadata={op_name="jit(compiled_generate_function)/while/body/stack" stack_frame_id=417} + %stack.1287 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.99), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/stack" stack_frame_id=332} + %stack.1288 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.100), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/stack" stack_frame_id=332} + %stack.1289 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1287, %stack.1288), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_17/self_attention/stack" stack_frame_id=332} + %stack.1517 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1289), metadata={op_name="jit(compiled_generate_function)/while/body/stack" stack_frame_id=417} + %stack.1302 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.101), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/stack" stack_frame_id=332} + %stack.1303 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.102), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/stack" stack_frame_id=332} + %stack.1304 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1302, %stack.1303), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_18/self_attention/stack" stack_frame_id=332} + %stack.1518 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1304), metadata={op_name="jit(compiled_generate_function)/while/body/stack" stack_frame_id=417} + %stack.1317 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.103), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/stack" stack_frame_id=332} + %stack.1318 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.104), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/stack" stack_frame_id=332} + %stack.1319 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1317, %stack.1318), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_19/self_attention/stack" stack_frame_id=332} + %stack.1519 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1319), metadata={op_name="jit(compiled_generate_function)/while/body/stack" stack_frame_id=417} + %stack.1332 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.105), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/stack" stack_frame_id=332} + %stack.1333 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.106), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/stack" stack_frame_id=332} + %stack.1334 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1332, %stack.1333), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_20/self_attention/stack" stack_frame_id=332} + %stack.1520 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1334), metadata={op_name="jit(compiled_generate_function)/while/body/stack" stack_frame_id=417} + %stack.1347 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.107), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/stack" stack_frame_id=332} + %stack.1348 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.108), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/stack" stack_frame_id=332} + %stack.1349 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1347, %stack.1348), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_21/self_attention/stack" stack_frame_id=332} + %stack.1521 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1349), metadata={op_name="jit(compiled_generate_function)/while/body/stack" stack_frame_id=417} + %stack.1362 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.109), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/stack" stack_frame_id=332} + %stack.1363 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.110), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/stack" stack_frame_id=332} + %stack.1364 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1362, %stack.1363), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_22/self_attention/stack" stack_frame_id=332} + %stack.1522 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1364), metadata={op_name="jit(compiled_generate_function)/while/body/stack" stack_frame_id=417} + %stack.1377 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.111), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/stack" stack_frame_id=332} + %stack.1378 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.112), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/stack" stack_frame_id=332} + %stack.1379 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1377, %stack.1378), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_23/self_attention/stack" stack_frame_id=332} + %stack.1523 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1379), metadata={op_name="jit(compiled_generate_function)/while/body/stack" stack_frame_id=417} + %stack.1392 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.113), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/stack" stack_frame_id=332} + %stack.1393 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.114), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/stack" stack_frame_id=332} + %stack.1394 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1392, %stack.1393), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_24/self_attention/stack" stack_frame_id=332} + %stack.1524 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1394), metadata={op_name="jit(compiled_generate_function)/while/body/stack" stack_frame_id=417} + %stack.1407 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.115), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/stack" stack_frame_id=332} + %stack.1408 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.116), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/stack" stack_frame_id=332} + %stack.1409 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1407, %stack.1408), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_25/self_attention/stack" stack_frame_id=332} + %stack.1525 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1409), metadata={op_name="jit(compiled_generate_function)/while/body/stack" stack_frame_id=417} + %stack.1422 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.117), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/stack" stack_frame_id=332} + %stack.1423 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.118), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/stack" stack_frame_id=332} + %stack.1424 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1422, %stack.1423), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_26/self_attention/stack" stack_frame_id=332} + %stack.1526 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1424), metadata={op_name="jit(compiled_generate_function)/while/body/stack" stack_frame_id=417} + %stack.1437 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.119), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/stack" stack_frame_id=332} + %stack.1438 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.120), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/stack" stack_frame_id=332} + %stack.1439 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1437, %stack.1438), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_27/self_attention/stack" stack_frame_id=332} + %stack.1527 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1439), metadata={op_name="jit(compiled_generate_function)/while/body/stack" stack_frame_id=417} + %stack.1452 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.121), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/stack" stack_frame_id=332} + %stack.1453 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.122), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/stack" stack_frame_id=332} + %stack.1454 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1452, %stack.1453), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_28/self_attention/stack" stack_frame_id=332} + %stack.1528 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1454), metadata={op_name="jit(compiled_generate_function)/while/body/stack" stack_frame_id=417} + %stack.1467 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.123), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/stack" stack_frame_id=332} + %stack.1468 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.124), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/stack" stack_frame_id=332} + %stack.1469 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1467, %stack.1468), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_29/self_attention/stack" stack_frame_id=332} + %stack.1529 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1469), metadata={op_name="jit(compiled_generate_function)/while/body/stack" stack_frame_id=417} + %stack.1482 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.125), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/stack" stack_frame_id=332} + %stack.1483 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.126), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/stack" stack_frame_id=332} + %stack.1484 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1482, %stack.1483), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_30/self_attention/stack" stack_frame_id=332} + %stack.1530 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1484), metadata={op_name="jit(compiled_generate_function)/while/body/stack" stack_frame_id=417} + %stack.1497 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.127), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/stack" stack_frame_id=332} + %stack.1498 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dynamic_update_slice.128), metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/stack" stack_frame_id=332} + %stack.1499 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1497, %stack.1498), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/while/body/transformer_layer_31/self_attention/stack" stack_frame_id=332} + %stack.1531 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1499), metadata={op_name="jit(compiled_generate_function)/while/body/stack" stack_frame_id=417} + %stack.1532 = bf16[1,32,2,41,8,128]{5,4,3,2,1,0} concatenate(%stack.1500, %stack.1501, %stack.1502, %stack.1503, %stack.1504, /*index=5*/%stack.1505, %stack.1506, %stack.1507, %stack.1508, %stack.1509, /*index=10*/%stack.1510, %stack.1511, %stack.1512, %stack.1513, %stack.1514, /*index=15*/%stack.1515, %stack.1516, %stack.1517, %stack.1518, %stack.1519, /*index=20*/%stack.1520, %stack.1521, %stack.1522, %stack.1523, %stack.1524, /*index=25*/%stack.1525, %stack.1526, %stack.1527, %stack.1528, %stack.1529, /*index=30*/%stack.1530, %stack.1531), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/while/body/stack" stack_frame_id=417} + %add.1121 = s32[] add(%arg_tuple.5#3, %constant.128), metadata={op_name="jit(compiled_generate_function)/while/body/add" stack_frame_id=485} + %add.1122 = s32[] add(%arg_tuple.5#4, %constant.128), metadata={op_name="jit(compiled_generate_function)/while/body/add" stack_frame_id=486} + ROOT %tuple.9 = (u32[2]{0}, s32[1,41]{1,0}, bf16[1,32,2,41,8,128]{5,4,3,2,1,0}, s32[], s32[], /*index=5*/pred[1,41]{1,0}, s32[], bf16[32000,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=10*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=15*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=20*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=25*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=30*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=35*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=40*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=45*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=50*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=55*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=60*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=65*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=70*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=75*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=80*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=85*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=90*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=95*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=100*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=105*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=110*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=115*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=120*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=125*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=130*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=135*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=140*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=145*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=150*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=155*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=160*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=165*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=170*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=175*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=180*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=185*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=190*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=195*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=200*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=205*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=210*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=215*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=220*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=225*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=230*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=235*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=240*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=245*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=250*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=255*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=260*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=265*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=270*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=275*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=280*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=285*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=290*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=295*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32000]{1,0}) tuple(%add.1116, %dynamic_update_slice.129, %stack.1532, %add.1121, %add.1122, /*index=5*/%arg_tuple.5#5, %arg_tuple.5#6, %arg_tuple.5#7, %arg_tuple.5#8, %arg_tuple.5#9, /*index=10*/%arg_tuple.5#10, %arg_tuple.5#11, %arg_tuple.5#12, %arg_tuple.5#13, %arg_tuple.5#14, /*index=15*/%arg_tuple.5#15, %arg_tuple.5#16, %arg_tuple.5#17, %arg_tuple.5#18, %arg_tuple.5#19, /*index=20*/%arg_tuple.5#20, %arg_tuple.5#21, %arg_tuple.5#22, %arg_tuple.5#23, %arg_tuple.5#24, /*index=25*/%arg_tuple.5#25, %arg_tuple.5#26, %arg_tuple.5#27, %arg_tuple.5#28, %arg_tuple.5#29, /*index=30*/%arg_tuple.5#30, %arg_tuple.5#31, %arg_tuple.5#32, %arg_tuple.5#33, %arg_tuple.5#34, /*index=35*/%arg_tuple.5#35, %arg_tuple.5#36, %arg_tuple.5#37, %arg_tuple.5#38, %arg_tuple.5#39, /*index=40*/%arg_tuple.5#40, %arg_tuple.5#41, %arg_tuple.5#42, %arg_tuple.5#43, %arg_tuple.5#44, /*index=45*/%arg_tuple.5#45, %arg_tuple.5#46, %arg_tuple.5#47, %arg_tuple.5#48, %arg_tuple.5#49, /*index=50*/%arg_tuple.5#50, %arg_tuple.5#51, %arg_tuple.5#52, %arg_tuple.5#53, %arg_tuple.5#54, /*index=55*/%arg_tuple.5#55, %arg_tuple.5#56, %arg_tuple.5#57, %arg_tuple.5#58, %arg_tuple.5#59, /*index=60*/%arg_tuple.5#60, %arg_tuple.5#61, %arg_tuple.5#62, %arg_tuple.5#63, %arg_tuple.5#64, /*index=65*/%arg_tuple.5#65, %arg_tuple.5#66, %arg_tuple.5#67, %arg_tuple.5#68, %arg_tuple.5#69, /*index=70*/%arg_tuple.5#70, %arg_tuple.5#71, %arg_tuple.5#72, %arg_tuple.5#73, %arg_tuple.5#74, /*index=75*/%arg_tuple.5#75, %arg_tuple.5#76, %arg_tuple.5#77, %arg_tuple.5#78, %arg_tuple.5#79, /*index=80*/%arg_tuple.5#80, %arg_tuple.5#81, %arg_tuple.5#82, %arg_tuple.5#83, %arg_tuple.5#84, /*index=85*/%arg_tuple.5#85, %arg_tuple.5#86, %arg_tuple.5#87, %arg_tuple.5#88, %arg_tuple.5#89, /*index=90*/%arg_tuple.5#90, %arg_tuple.5#91, %arg_tuple.5#92, %arg_tuple.5#93, %arg_tuple.5#94, /*index=95*/%arg_tuple.5#95, %arg_tuple.5#96, %arg_tuple.5#97, %arg_tuple.5#98, %arg_tuple.5#99, /*index=100*/%arg_tuple.5#100, %arg_tuple.5#101, %arg_tuple.5#102, %arg_tuple.5#103, %arg_tuple.5#104, /*index=105*/%arg_tuple.5#105, %arg_tuple.5#106, %arg_tuple.5#107, %arg_tuple.5#108, %arg_tuple.5#109, /*index=110*/%arg_tuple.5#110, %arg_tuple.5#111, %arg_tuple.5#112, %arg_tuple.5#113, %arg_tuple.5#114, /*index=115*/%arg_tuple.5#115, %arg_tuple.5#116, %arg_tuple.5#117, %arg_tuple.5#118, %arg_tuple.5#119, /*index=120*/%arg_tuple.5#120, %arg_tuple.5#121, %arg_tuple.5#122, %arg_tuple.5#123, %arg_tuple.5#124, /*index=125*/%arg_tuple.5#125, %arg_tuple.5#126, %arg_tuple.5#127, %arg_tuple.5#128, %arg_tuple.5#129, /*index=130*/%arg_tuple.5#130, %arg_tuple.5#131, %arg_tuple.5#132, %arg_tuple.5#133, %arg_tuple.5#134, /*index=135*/%arg_tuple.5#135, %arg_tuple.5#136, %arg_tuple.5#137, %arg_tuple.5#138, %arg_tuple.5#139, /*index=140*/%arg_tuple.5#140, %arg_tuple.5#141, %arg_tuple.5#142, %arg_tuple.5#143, %arg_tuple.5#144, /*index=145*/%arg_tuple.5#145, %arg_tuple.5#146, %arg_tuple.5#147, %arg_tuple.5#148, %arg_tuple.5#149, /*index=150*/%arg_tuple.5#150, %arg_tuple.5#151, %arg_tuple.5#152, %arg_tuple.5#153, %arg_tuple.5#154, /*index=155*/%arg_tuple.5#155, %arg_tuple.5#156, %arg_tuple.5#157, %arg_tuple.5#158, %arg_tuple.5#159, /*index=160*/%arg_tuple.5#160, %arg_tuple.5#161, %arg_tuple.5#162, %arg_tuple.5#163, %arg_tuple.5#164, /*index=165*/%arg_tuple.5#165, %arg_tuple.5#166, %arg_tuple.5#167, %arg_tuple.5#168, %arg_tuple.5#169, /*index=170*/%arg_tuple.5#170, %arg_tuple.5#171, %arg_tuple.5#172, %arg_tuple.5#173, %arg_tuple.5#174, /*index=175*/%arg_tuple.5#175, %arg_tuple.5#176, %arg_tuple.5#177, %arg_tuple.5#178, %arg_tuple.5#179, /*index=180*/%arg_tuple.5#180, %arg_tuple.5#181, %arg_tuple.5#182, %arg_tuple.5#183, %arg_tuple.5#184, /*index=185*/%arg_tuple.5#185, %arg_tuple.5#186, %arg_tuple.5#187, %arg_tuple.5#188, %arg_tuple.5#189, /*index=190*/%arg_tuple.5#190, %arg_tuple.5#191, %arg_tuple.5#192, %arg_tuple.5#193, %arg_tuple.5#194, /*index=195*/%arg_tuple.5#195, %arg_tuple.5#196, %arg_tuple.5#197, %arg_tuple.5#198, %arg_tuple.5#199, /*index=200*/%arg_tuple.5#200, %arg_tuple.5#201, %arg_tuple.5#202, %arg_tuple.5#203, %arg_tuple.5#204, /*index=205*/%arg_tuple.5#205, %arg_tuple.5#206, %arg_tuple.5#207, %arg_tuple.5#208, %arg_tuple.5#209, /*index=210*/%arg_tuple.5#210, %arg_tuple.5#211, %arg_tuple.5#212, %arg_tuple.5#213, %arg_tuple.5#214, /*index=215*/%arg_tuple.5#215, %arg_tuple.5#216, %arg_tuple.5#217, %arg_tuple.5#218, %arg_tuple.5#219, /*index=220*/%arg_tuple.5#220, %arg_tuple.5#221, %arg_tuple.5#222, %arg_tuple.5#223, %arg_tuple.5#224, /*index=225*/%arg_tuple.5#225, %arg_tuple.5#226, %arg_tuple.5#227, %arg_tuple.5#228, %arg_tuple.5#229, /*index=230*/%arg_tuple.5#230, %arg_tuple.5#231, %arg_tuple.5#232, %arg_tuple.5#233, %arg_tuple.5#234, /*index=235*/%arg_tuple.5#235, %arg_tuple.5#236, %arg_tuple.5#237, %arg_tuple.5#238, %arg_tuple.5#239, /*index=240*/%arg_tuple.5#240, %arg_tuple.5#241, %arg_tuple.5#242, %arg_tuple.5#243, %arg_tuple.5#244, /*index=245*/%arg_tuple.5#245, %arg_tuple.5#246, %arg_tuple.5#247, %arg_tuple.5#248, %arg_tuple.5#249, /*index=250*/%arg_tuple.5#250, %arg_tuple.5#251, %arg_tuple.5#252, %arg_tuple.5#253, %arg_tuple.5#254, /*index=255*/%arg_tuple.5#255, %arg_tuple.5#256, %arg_tuple.5#257, %arg_tuple.5#258, %arg_tuple.5#259, /*index=260*/%arg_tuple.5#260, %arg_tuple.5#261, %arg_tuple.5#262, %arg_tuple.5#263, %arg_tuple.5#264, /*index=265*/%arg_tuple.5#265, %arg_tuple.5#266, %arg_tuple.5#267, %arg_tuple.5#268, %arg_tuple.5#269, /*index=270*/%arg_tuple.5#270, %arg_tuple.5#271, %arg_tuple.5#272, %arg_tuple.5#273, %arg_tuple.5#274, /*index=275*/%arg_tuple.5#275, %arg_tuple.5#276, %arg_tuple.5#277, %arg_tuple.5#278, %arg_tuple.5#279, /*index=280*/%arg_tuple.5#280, %arg_tuple.5#281, %arg_tuple.5#282, %arg_tuple.5#283, %arg_tuple.5#284, /*index=285*/%arg_tuple.5#285, %arg_tuple.5#286, %arg_tuple.5#287, %arg_tuple.5#288, %arg_tuple.5#289, /*index=290*/%arg_tuple.5#290, %arg_tuple.5#291, %arg_tuple.5#292, %arg_tuple.5#293, %arg_tuple.5#294, /*index=295*/%arg_tuple.5#295, %arg_tuple.5#296, %arg_tuple.5#297) +} + +%region_266.281 (reduce_or.3: pred[], reduce_or.4: pred[]) -> pred[] { + %reduce_or.3 = pred[] parameter(0), metadata={op_name="reduce_or"} + %reduce_or.4 = pred[] parameter(1), metadata={op_name="reduce_or"} + ROOT %reduce_or.5 = pred[] or(%reduce_or.3, %reduce_or.4), metadata={op_name="jit(compiled_generate_function)/while/cond/reduce_or" stack_frame_id=491} +} + +%region_267.282 (reduce_and.27: pred[], reduce_and.28: pred[]) -> pred[] { + %reduce_and.27 = pred[] parameter(0), metadata={op_name="reduce_and"} + %reduce_and.28 = pred[] parameter(1), metadata={op_name="reduce_and"} + ROOT %reduce_and.29 = pred[] and(%reduce_and.27, %reduce_and.28), metadata={op_name="jit(compiled_generate_function)/while/cond/reduce_and" stack_frame_id=494} +} + +%region_265.283 (arg_tuple.7: (u32[2], s32[1,41], bf16[1,32,2,41,8,128], s32[], s32[], /*index=5*/pred[1,41], s32[], bf16[32000,4096], bf16[4096], bf16[4096,32,128], /*index=10*/bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], /*index=15*/bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], /*index=20*/bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], /*index=25*/bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], /*index=30*/bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], /*index=35*/bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], /*index=40*/bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], /*index=45*/bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], /*index=50*/bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], /*index=55*/bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], /*index=60*/bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], /*index=65*/bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], /*index=70*/bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], /*index=75*/bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], /*index=80*/bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], /*index=85*/bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], /*index=90*/bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], /*index=95*/bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], /*index=100*/bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], /*index=105*/bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], /*index=110*/bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], /*index=115*/bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], /*index=120*/bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], /*index=125*/bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], /*index=130*/bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], /*index=135*/bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], /*index=140*/bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], /*index=145*/bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], /*index=150*/bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], /*index=155*/bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], /*index=160*/bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], /*index=165*/bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], /*index=170*/bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], /*index=175*/bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], /*index=180*/bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], /*index=185*/bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], /*index=190*/bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], /*index=195*/bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], /*index=200*/bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], /*index=205*/bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], /*index=210*/bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], /*index=215*/bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], /*index=220*/bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], /*index=225*/bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], /*index=230*/bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], /*index=235*/bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], /*index=240*/bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], /*index=245*/bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], /*index=250*/bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], /*index=255*/bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], /*index=260*/bf16[4096], bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], /*index=265*/bf16[4096], bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], /*index=270*/bf16[4096,32,128], bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], /*index=275*/bf16[4096,14336], bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], /*index=280*/bf16[4096,8,128], bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], /*index=285*/bf16[4096,14336], bf16[14336,4096], bf16[4096], bf16[4096,32,128], bf16[4096,8,128], /*index=290*/bf16[4096,8,128], bf16[32,128,4096], bf16[4096], bf16[4096,14336], bf16[4096,14336], /*index=295*/bf16[14336,4096], bf16[4096], bf16[4096,32000])) -> pred[] { + %arg_tuple.7 = (u32[2]{0}, s32[1,41]{1,0}, bf16[1,32,2,41,8,128]{5,4,3,2,1,0}, s32[], s32[], /*index=5*/pred[1,41]{1,0}, s32[], bf16[32000,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=10*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=15*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=20*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=25*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=30*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=35*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=40*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=45*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=50*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=55*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=60*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=65*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=70*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=75*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=80*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=85*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=90*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=95*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=100*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=105*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=110*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=115*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=120*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=125*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=130*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=135*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=140*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=145*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=150*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=155*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=160*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=165*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=170*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=175*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=180*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=185*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=190*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=195*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=200*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=205*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=210*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=215*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=220*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=225*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=230*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=235*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=240*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=245*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=250*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=255*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=260*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=265*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=270*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=275*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=280*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=285*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=290*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=295*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32000]{1,0}) parameter(0) + %constant.134 = s32[] constant(2) + %broadcast.45 = s32[1,41]{1,0} broadcast(%constant.134), dimensions={} + %eq.3 = pred[1,41]{1,0} compare(%arg_tuple.7#1, %broadcast.45), direction=EQ, metadata={op_name="jit(compiled_generate_function)/while/cond/eq" stack_frame_id=499} + %not.2 = pred[1,41]{1,0} not(%arg_tuple.7#5), metadata={op_name="jit(compiled_generate_function)/while/cond/not" stack_frame_id=495} + %and.295 = pred[1,41]{1,0} and(%eq.3, %not.2), metadata={op_name="jit(compiled_generate_function)/while/cond/and" stack_frame_id=502} + %constant.133 = pred[] constant(false) + %reduce_or.7 = pred[1]{0} reduce(%and.295, %constant.133), dimensions={1}, to_apply=%region_266.281, metadata={op_name="jit(compiled_generate_function)/while/cond/reduce_or" stack_frame_id=491} + %constant.132 = pred[] constant(true) + %reduce_and.31 = pred[] reduce(%reduce_or.7, %constant.132), dimensions={0}, to_apply=%region_267.282, metadata={op_name="jit(compiled_generate_function)/while/cond/reduce_and" stack_frame_id=494} + %not.3 = pred[] not(%reduce_and.31), metadata={op_name="jit(compiled_generate_function)/while/cond/not" stack_frame_id=505} + %lt.622 = pred[] compare(%arg_tuple.7#4, %arg_tuple.7#6), direction=LT, metadata={op_name="jit(compiled_generate_function)/while/cond/lt" stack_frame_id=506} + ROOT %and.296 = pred[] and(%not.3, %lt.622), metadata={op_name="jit(compiled_generate_function)/while/cond/and" stack_frame_id=487} +} + +%region_268.284 (reduce_window_sum.3: s32[], reduce_window_sum.4: s32[]) -> s32[] { + %reduce_window_sum.3 = s32[] parameter(0), metadata={op_name="reduce_window_sum"} + %reduce_window_sum.4 = s32[] parameter(1), metadata={op_name="reduce_window_sum"} + ROOT %reduce_window_sum.5 = s32[] add(%reduce_window_sum.3, %reduce_window_sum.4), metadata={op_name="reduce_window_sum"} +} + +%cumsum_6.285 (Arg_0.30: s32[1,41]) -> s32[1,41] { + %Arg_0.30 = s32[1,41]{1,0} parameter(0) + %constant.136 = s32[] constant(0) + ROOT %reduce_window_sum.7 = s32[1,41]{1,0} reduce-window(%Arg_0.30, %constant.136), window={size=1x41 pad=0_0x40_0}, to_apply=%region_268.284, metadata={op_name="reduce_window_sum"} +} + +%cumsum.286 (Arg_0.31: s32[1,41]) -> s32[1,41] { + %Arg_0.31 = s32[1,41]{1,0} parameter(0) + ROOT %call.3 = s32[1,41]{1,0} call(%Arg_0.31), to_apply=%cumsum_6.285 +} + +ENTRY %main.287 (inputs__padding_mask__.1: pred[1,41], inputs__token_ids__.1: s32[1,41], state_0__0_.1: u32[2], state_1__0_.1: bf16[32000,4096], state_1__1_.1: bf16[4096,32000], state_1__2_.1: bf16[4096,32,128], state_1__3_.1: bf16[4096,8,128], state_1__4_.1: bf16[4096,8,128], state_1__5_.1: bf16[32,128,4096], state_1__6_.1: bf16[4096], state_1__7_.1: bf16[4096,14336], state_1__8_.1: bf16[4096,14336], state_1__9_.1: bf16[14336,4096], state_1__10_.1: bf16[4096], state_1__11_.1: bf16[4096,32,128], state_1__12_.1: bf16[4096,8,128], state_1__13_.1: bf16[4096,8,128], state_1__14_.1: bf16[32,128,4096], state_1__15_.1: bf16[4096], state_1__16_.1: bf16[4096,14336], state_1__17_.1: bf16[4096,14336], state_1__18_.1: bf16[14336,4096], state_1__19_.1: bf16[4096], state_1__20_.1: bf16[4096,32,128], state_1__21_.1: bf16[4096,8,128], state_1__22_.1: bf16[4096,8,128], state_1__23_.1: bf16[32,128,4096], state_1__24_.1: bf16[4096], state_1__25_.1: bf16[4096,14336], state_1__26_.1: bf16[4096,14336], state_1__27_.1: bf16[14336,4096], state_1__28_.1: bf16[4096], state_1__29_.1: bf16[4096,32,128], state_1__30_.1: bf16[4096,8,128], state_1__31_.1: bf16[4096,8,128], state_1__32_.1: bf16[32,128,4096], state_1__33_.1: bf16[4096], state_1__34_.1: bf16[4096,14336], state_1__35_.1: bf16[4096,14336], state_1__36_.1: bf16[14336,4096], state_1__37_.1: bf16[4096], state_1__38_.1: bf16[4096,32,128], state_1__39_.1: bf16[4096,8,128], state_1__40_.1: bf16[4096,8,128], state_1__41_.1: bf16[32,128,4096], state_1__42_.1: bf16[4096], state_1__43_.1: bf16[4096,14336], state_1__44_.1: bf16[4096,14336], state_1__45_.1: bf16[14336,4096], state_1__46_.1: bf16[4096], state_1__47_.1: bf16[4096,32,128], state_1__48_.1: bf16[4096,8,128], state_1__49_.1: bf16[4096,8,128], state_1__50_.1: bf16[32,128,4096], state_1__51_.1: bf16[4096], state_1__52_.1: bf16[4096,14336], state_1__53_.1: bf16[4096,14336], state_1__54_.1: bf16[14336,4096], state_1__55_.1: bf16[4096], state_1__56_.1: bf16[4096,32,128], state_1__57_.1: bf16[4096,8,128], state_1__58_.1: bf16[4096,8,128], state_1__59_.1: bf16[32,128,4096], state_1__60_.1: bf16[4096], state_1__61_.1: bf16[4096,14336], state_1__62_.1: bf16[4096,14336], state_1__63_.1: bf16[14336,4096], state_1__64_.1: bf16[4096], state_1__65_.1: bf16[4096,32,128], state_1__66_.1: bf16[4096,8,128], state_1__67_.1: bf16[4096,8,128], state_1__68_.1: bf16[32,128,4096], state_1__69_.1: bf16[4096], state_1__70_.1: bf16[4096,14336], state_1__71_.1: bf16[4096,14336], state_1__72_.1: bf16[14336,4096], state_1__73_.1: bf16[4096], state_1__74_.1: bf16[4096,32,128], state_1__75_.1: bf16[4096,8,128], state_1__76_.1: bf16[4096,8,128], state_1__77_.1: bf16[32,128,4096], state_1__78_.1: bf16[4096], state_1__79_.1: bf16[4096,14336], state_1__80_.1: bf16[4096,14336], state_1__81_.1: bf16[14336,4096], state_1__82_.1: bf16[4096], state_1__83_.1: bf16[4096,32,128], state_1__84_.1: bf16[4096,8,128], state_1__85_.1: bf16[4096,8,128], state_1__86_.1: bf16[32,128,4096], state_1__87_.1: bf16[4096], state_1__88_.1: bf16[4096,14336], state_1__89_.1: bf16[4096,14336], state_1__90_.1: bf16[14336,4096], state_1__91_.1: bf16[4096], state_1__92_.1: bf16[4096,32,128], state_1__93_.1: bf16[4096,8,128], state_1__94_.1: bf16[4096,8,128], state_1__95_.1: bf16[32,128,4096], state_1__96_.1: bf16[4096], state_1__97_.1: bf16[4096,14336], state_1__98_.1: bf16[4096,14336], state_1__99_.1: bf16[14336,4096], state_1__100_.1: bf16[4096], state_1__101_.1: bf16[4096,32,128], state_1__102_.1: bf16[4096,8,128], state_1__103_.1: bf16[4096,8,128], state_1__104_.1: bf16[32,128,4096], state_1__105_.1: bf16[4096], state_1__106_.1: bf16[4096,14336], state_1__107_.1: bf16[4096,14336], state_1__108_.1: bf16[14336,4096], state_1__109_.1: bf16[4096], state_1__110_.1: bf16[4096,32,128], state_1__111_.1: bf16[4096,8,128], state_1__112_.1: bf16[4096,8,128], state_1__113_.1: bf16[32,128,4096], state_1__114_.1: bf16[4096], state_1__115_.1: bf16[4096,14336], state_1__116_.1: bf16[4096,14336], state_1__117_.1: bf16[14336,4096], state_1__118_.1: bf16[4096], state_1__119_.1: bf16[4096,32,128], state_1__120_.1: bf16[4096,8,128], state_1__121_.1: bf16[4096,8,128], state_1__122_.1: bf16[32,128,4096], state_1__123_.1: bf16[4096], state_1__124_.1: bf16[4096,14336], state_1__125_.1: bf16[4096,14336], state_1__126_.1: bf16[14336,4096], state_1__127_.1: bf16[4096], state_1__128_.1: bf16[4096,32,128], state_1__129_.1: bf16[4096,8,128], state_1__130_.1: bf16[4096,8,128], state_1__131_.1: bf16[32,128,4096], state_1__132_.1: bf16[4096], state_1__133_.1: bf16[4096,14336], state_1__134_.1: bf16[4096,14336], state_1__135_.1: bf16[14336,4096], state_1__136_.1: bf16[4096], state_1__137_.1: bf16[4096,32,128], state_1__138_.1: bf16[4096,8,128], state_1__139_.1: bf16[4096,8,128], state_1__140_.1: bf16[32,128,4096], state_1__141_.1: bf16[4096], state_1__142_.1: bf16[4096,14336], state_1__143_.1: bf16[4096,14336], state_1__144_.1: bf16[14336,4096], state_1__145_.1: bf16[4096], state_1__146_.1: bf16[4096,32,128], state_1__147_.1: bf16[4096,8,128], state_1__148_.1: bf16[4096,8,128], state_1__149_.1: bf16[32,128,4096], state_1__150_.1: bf16[4096], state_1__151_.1: bf16[4096,14336], state_1__152_.1: bf16[4096,14336], state_1__153_.1: bf16[14336,4096], state_1__154_.1: bf16[4096], state_1__155_.1: bf16[4096,32,128], state_1__156_.1: bf16[4096,8,128], state_1__157_.1: bf16[4096,8,128], state_1__158_.1: bf16[32,128,4096], state_1__159_.1: bf16[4096], state_1__160_.1: bf16[4096,14336], state_1__161_.1: bf16[4096,14336], state_1__162_.1: bf16[14336,4096], state_1__163_.1: bf16[4096], state_1__164_.1: bf16[4096,32,128], state_1__165_.1: bf16[4096,8,128], state_1__166_.1: bf16[4096,8,128], state_1__167_.1: bf16[32,128,4096], state_1__168_.1: bf16[4096], state_1__169_.1: bf16[4096,14336], state_1__170_.1: bf16[4096,14336], state_1__171_.1: bf16[14336,4096], state_1__172_.1: bf16[4096], state_1__173_.1: bf16[4096,32,128], state_1__174_.1: bf16[4096,8,128], state_1__175_.1: bf16[4096,8,128], state_1__176_.1: bf16[32,128,4096], state_1__177_.1: bf16[4096], state_1__178_.1: bf16[4096,14336], state_1__179_.1: bf16[4096,14336], state_1__180_.1: bf16[14336,4096], state_1__181_.1: bf16[4096], state_1__182_.1: bf16[4096,32,128], state_1__183_.1: bf16[4096,8,128], state_1__184_.1: bf16[4096,8,128], state_1__185_.1: bf16[32,128,4096], state_1__186_.1: bf16[4096], state_1__187_.1: bf16[4096,14336], state_1__188_.1: bf16[4096,14336], state_1__189_.1: bf16[14336,4096], state_1__190_.1: bf16[4096], state_1__191_.1: bf16[4096,32,128], state_1__192_.1: bf16[4096,8,128], state_1__193_.1: bf16[4096,8,128], state_1__194_.1: bf16[32,128,4096], state_1__195_.1: bf16[4096], state_1__196_.1: bf16[4096,14336], state_1__197_.1: bf16[4096,14336], state_1__198_.1: bf16[14336,4096], state_1__199_.1: bf16[4096], state_1__200_.1: bf16[4096,32,128], state_1__201_.1: bf16[4096,8,128], state_1__202_.1: bf16[4096,8,128], state_1__203_.1: bf16[32,128,4096], state_1__204_.1: bf16[4096], state_1__205_.1: bf16[4096,14336], state_1__206_.1: bf16[4096,14336], state_1__207_.1: bf16[14336,4096], state_1__208_.1: bf16[4096], state_1__209_.1: bf16[4096,32,128], state_1__210_.1: bf16[4096,8,128], state_1__211_.1: bf16[4096,8,128], state_1__212_.1: bf16[32,128,4096], state_1__213_.1: bf16[4096], state_1__214_.1: bf16[4096,14336], state_1__215_.1: bf16[4096,14336], state_1__216_.1: bf16[14336,4096], state_1__217_.1: bf16[4096], state_1__218_.1: bf16[4096,32,128], state_1__219_.1: bf16[4096,8,128], state_1__220_.1: bf16[4096,8,128], state_1__221_.1: bf16[32,128,4096], state_1__222_.1: bf16[4096], state_1__223_.1: bf16[4096,14336], state_1__224_.1: bf16[4096,14336], state_1__225_.1: bf16[14336,4096], state_1__226_.1: bf16[4096], state_1__227_.1: bf16[4096,32,128], state_1__228_.1: bf16[4096,8,128], state_1__229_.1: bf16[4096,8,128], state_1__230_.1: bf16[32,128,4096], state_1__231_.1: bf16[4096], state_1__232_.1: bf16[4096,14336], state_1__233_.1: bf16[4096,14336], state_1__234_.1: bf16[14336,4096], state_1__235_.1: bf16[4096], state_1__236_.1: bf16[4096,32,128], state_1__237_.1: bf16[4096,8,128], state_1__238_.1: bf16[4096,8,128], state_1__239_.1: bf16[32,128,4096], state_1__240_.1: bf16[4096], state_1__241_.1: bf16[4096,14336], state_1__242_.1: bf16[4096,14336], state_1__243_.1: bf16[14336,4096], state_1__244_.1: bf16[4096], state_1__245_.1: bf16[4096,32,128], state_1__246_.1: bf16[4096,8,128], state_1__247_.1: bf16[4096,8,128], state_1__248_.1: bf16[32,128,4096], state_1__249_.1: bf16[4096], state_1__250_.1: bf16[4096,14336], state_1__251_.1: bf16[4096,14336], state_1__252_.1: bf16[14336,4096], state_1__253_.1: bf16[4096], state_1__254_.1: bf16[4096,32,128], state_1__255_.1: bf16[4096,8,128], state_1__256_.1: bf16[4096,8,128], state_1__257_.1: bf16[32,128,4096], state_1__258_.1: bf16[4096], state_1__259_.1: bf16[4096,14336], state_1__260_.1: bf16[4096,14336], state_1__261_.1: bf16[14336,4096], state_1__262_.1: bf16[4096], state_1__263_.1: bf16[4096,32,128], state_1__264_.1: bf16[4096,8,128], state_1__265_.1: bf16[4096,8,128], state_1__266_.1: bf16[32,128,4096], state_1__267_.1: bf16[4096], state_1__268_.1: bf16[4096,14336], state_1__269_.1: bf16[4096,14336], state_1__270_.1: bf16[14336,4096], state_1__271_.1: bf16[4096], state_1__272_.1: bf16[4096,32,128], state_1__273_.1: bf16[4096,8,128], state_1__274_.1: bf16[4096,8,128], state_1__275_.1: bf16[32,128,4096], state_1__276_.1: bf16[4096], state_1__277_.1: bf16[4096,14336], state_1__278_.1: bf16[4096,14336], state_1__279_.1: bf16[14336,4096], state_1__280_.1: bf16[4096], state_1__281_.1: bf16[4096,32,128], state_1__282_.1: bf16[4096,8,128], state_1__283_.1: bf16[4096,8,128], state_1__284_.1: bf16[32,128,4096], state_1__285_.1: bf16[4096], state_1__286_.1: bf16[4096,14336], state_1__287_.1: bf16[4096,14336], state_1__288_.1: bf16[14336,4096], state_1__289_.1: bf16[4096], state_1__290_.1: bf16[4096]) -> (pred[1,41], s32[1,41], u32[2]) { + %state_0__0_.1 = u32[2]{0} parameter(2), metadata={op_name="state[0][0]"} + %inputs__token_ids__.1 = s32[1,41]{1,0} parameter(1), metadata={op_name="inputs[\'token_ids\']"} + %state_1__0_.1 = bf16[32000,4096]{1,0} parameter(3), metadata={op_name="state[1][0]"} + %jit__take_.3 = bf16[1,41,4096]{2,1,0} call(%state_1__0_.1, %inputs__token_ids__.1), to_apply=%_take.3, metadata={op_name="jit(compiled_generate_function)/token_embedding/jit(_take)" stack_frame_id=520} + %convert_element_type.1734 = f32[1,41,4096]{2,1,0} convert(%jit__take_.3), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention_layernorm/convert_element_type" stack_frame_id=555} + %constant.148 = f32[] constant(2) + %broadcast.56 = f32[1,41,4096]{2,1,0} broadcast(%constant.148), dimensions={} + %pow.384 = f32[1,41,4096]{2,1,0} power(%convert_element_type.1734, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention_layernorm/pow" stack_frame_id=558} + %constant.155 = f32[] constant(0) + %reduce_sum.1449 = f32[1,41]{1,0} reduce(%pow.384, %constant.155), dimensions={2}, to_apply=%region_1.4, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention_layernorm/reduce_sum" stack_frame_id=24} + %broadcast_in_dim.1497 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1449), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention_layernorm/broadcast_in_dim" stack_frame_id=24} + %constant.147 = f32[] constant(4096) + %broadcast.55 = f32[1,41,1]{2,1,0} broadcast(%constant.147), dimensions={} + %div.969 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1497, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention_layernorm/div" stack_frame_id=24} + %constant.146 = f32[] constant(1e-05) + %broadcast.54 = f32[1,41,1]{2,1,0} broadcast(%constant.146), dimensions={} + %add.1124 = f32[1,41,1]{2,1,0} add(%div.969, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention_layernorm/add" stack_frame_id=559} + %rsqrt.193 = f32[1,41,1]{2,1,0} rsqrt(%add.1124), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention_layernorm/rsqrt" stack_frame_id=562} + %mul.3075 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.193), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3076 = f32[1,41]{1,0} reshape(%mul.3075), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3077 = f32[1,41,4096]{2,1,0} broadcast(%mul.3076), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3078 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.1734, %mul.3077), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention_layernorm/mul" stack_frame_id=563} + %state_1__6_.1 = bf16[4096]{0} parameter(9), metadata={op_name="state[1][6]"} + %convert_element_type.1735 = f32[4096]{0} convert(%state_1__6_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention_layernorm/convert_element_type" stack_frame_id=564} + %broadcast_in_dim.1498 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1735), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention_layernorm/broadcast_in_dim" stack_frame_id=564} + %mul.3079 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1498), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3080 = f32[1,4096]{1,0} reshape(%mul.3079), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3081 = f32[1,41,4096]{2,1,0} broadcast(%mul.3080), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3082 = f32[1,41,4096]{2,1,0} multiply(%mul.3078, %mul.3081), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention_layernorm/mul" stack_frame_id=564} + %convert_element_type.1736 = bf16[1,41,4096]{2,1,0} convert(%mul.3082), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention_layernorm/convert_element_type" stack_frame_id=568} + %state_1__3_.1 = bf16[4096,8,128]{2,1,0} parameter(6), metadata={op_name="state[1][3]"} + %dot_general.1052 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.1736, %state_1__3_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=647} + %iota.517 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/iota" stack_frame_id=677} + %broadcast_in_dim.1500 = f32[1,41]{1,0} reshape(%iota.517), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/broadcast_in_dim" stack_frame_id=680} + %constant.143 = f32[] constant(10000) + %broadcast.51 = f32[64]{0} broadcast(%constant.143), dimensions={} + %iota.516 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/iota" stack_frame_id=667} + %constant.145 = f32[] constant(2) + %broadcast.53 = f32[64]{0} broadcast(%constant.145), dimensions={} + %mul.3092 = f32[64]{0} multiply(%iota.516, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=667} + %constant.144 = f32[] constant(128) + %broadcast.52 = f32[64]{0} broadcast(%constant.144), dimensions={} + %div.972 = f32[64]{0} divide(%mul.3092, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/div" stack_frame_id=668} + %neg.392 = f32[64]{0} negate(%div.972), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/neg" stack_frame_id=669} + %pow.386 = f32[64]{0} power(%broadcast.51, %neg.392), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/pow" stack_frame_id=672} + %constant.142 = f32[] constant(1) + %broadcast.50 = f32[64]{0} broadcast(%constant.142), dimensions={} + %div.973 = f32[64]{0} divide(%pow.386, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/div" stack_frame_id=673} + %dot_general.1054 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1500, %div.973), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/bi,j->bij/dot_general" stack_frame_id=683} + %stack.1539 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1054), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/stack" stack_frame_id=686} + %stack.1540 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1054), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/stack" stack_frame_id=686} + %stack.1541 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1539, %stack.1540), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/stack" stack_frame_id=686} + %reshape.700 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1541), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/reshape"} + %cos.192 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.700), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/cos" stack_frame_id=690} + %convert_element_type.1739 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.192), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/convert_element_type" stack_frame_id=694} + %mul.3093 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1739), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=714} + %mul.3094 = bf16[1,41,128]{2,1,0} reshape(%mul.3093), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=714} + %mul.3095 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3094), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=714} + %mul.3096 = bf16[1,41,8,128]{3,2,1,0} multiply(%dot_general.1052, %mul.3095), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=714} + %split.385 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1052), slice={[0:1], [0:41], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/split" stack_frame_id=706} + %neg.393 = bf16[1,41,8,64]{3,2,1,0} negate(%split.385), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/neg" stack_frame_id=707} + %stack.1542 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%neg.393), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/stack" stack_frame_id=710} + %split.384 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1052), slice={[0:1], [0:41], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/split" stack_frame_id=706} + %stack.1543 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%split.384), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/stack" stack_frame_id=710} + %stack.1544 = bf16[1,41,8,2,64]{4,3,2,1,0} concatenate(%stack.1542, %stack.1543), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/stack" stack_frame_id=710} + %reshape.701 = bf16[1,41,8,128]{3,2,1,0} reshape(%stack.1544), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/reshape" stack_frame_id=713} + %sin.192 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.700), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/sin" stack_frame_id=698} + %convert_element_type.1740 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.192), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/convert_element_type" stack_frame_id=702} + %mul.3097 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1740), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=715} + %mul.3098 = bf16[1,41,128]{2,1,0} reshape(%mul.3097), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=715} + %mul.3099 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3098), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=715} + %mul.3100 = bf16[1,41,8,128]{3,2,1,0} multiply(%reshape.701, %mul.3099), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=715} + %add.1126 = bf16[1,41,8,128]{3,2,1,0} add(%mul.3096, %mul.3100), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/add" stack_frame_id=716} + %stack.1545 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%add.1126), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/stack" stack_frame_id=719} + %state_1__4_.1 = bf16[4096,8,128]{2,1,0} parameter(7), metadata={op_name="state[1][4]"} + %dot_general.1053 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.1736, %state_1__4_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=656} + %stack.1546 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dot_general.1053), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/stack" stack_frame_id=719} + %stack.1547 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1545, %stack.1546), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/stack" stack_frame_id=719} + %stack.2007 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1547), metadata={op_name="jit(compiled_generate_function)/stack" stack_frame_id=804} + %broadcast_in_dim.1502 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dot_general.1053), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/broadcast_in_dim" stack_frame_id=725} + %reshape.703 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1502), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/reshape" stack_frame_id=725} + %iota.510 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/iota" stack_frame_id=525} + %broadcast_in_dim.1491 = f32[41,1]{1,0} reshape(%iota.510), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/broadcast_in_dim" stack_frame_id=528} + %ge.481 = f32[41,1]{1,0} broadcast(%broadcast_in_dim.1491), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/ge" stack_frame_id=532} + %ge.482 = f32[41]{0} reshape(%ge.481), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/ge" stack_frame_id=532} + %ge.483 = f32[41,41]{1,0} broadcast(%ge.482), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/ge" stack_frame_id=532} + %iota.511 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/iota" stack_frame_id=531} + %broadcast_in_dim.1492 = f32[1,41]{1,0} reshape(%iota.511), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/broadcast_in_dim" stack_frame_id=532} + %ge.484 = f32[1,41]{1,0} broadcast(%broadcast_in_dim.1492), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/ge" stack_frame_id=532} + %ge.485 = f32[41]{0} reshape(%ge.484), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/ge" stack_frame_id=532} + %ge.486 = f32[41,41]{1,0} broadcast(%ge.485), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/ge" stack_frame_id=532} + %ge.487 = pred[41,41]{1,0} compare(%ge.483, %ge.486), direction=GE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/ge" stack_frame_id=532} + %broadcast_in_dim.1493 = pred[1,41,41]{2,1,0} reshape(%ge.487), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/broadcast_in_dim" stack_frame_id=535} + %convert_element_type.1733 = s32[1,41,41]{2,1,0} convert(%broadcast_in_dim.1493), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/convert_element_type" stack_frame_id=551} + %iota.512 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/iota" stack_frame_id=538} + %broadcast_in_dim.1494 = s32[41,1]{1,0} reshape(%iota.512), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/broadcast_in_dim" stack_frame_id=536} + %lt.623 = s32[41,1]{1,0} broadcast(%broadcast_in_dim.1494), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/lt" stack_frame_id=543} + %lt.624 = s32[41]{0} reshape(%lt.623), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/lt" stack_frame_id=543} + %lt.625 = s32[41,41]{1,0} broadcast(%lt.624), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/lt" stack_frame_id=543} + %iota.513 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/iota" stack_frame_id=541} + %broadcast_in_dim.1495 = s32[1,41]{1,0} reshape(%iota.513), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/broadcast_in_dim" stack_frame_id=539} + %constant.149 = s32[] constant(4096) + %broadcast.57 = s32[1,41]{1,0} broadcast(%constant.149), dimensions={} + %add.1123 = s32[1,41]{1,0} add(%broadcast_in_dim.1495, %broadcast.57), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/add" stack_frame_id=542} + %lt.626 = s32[1,41]{1,0} broadcast(%add.1123), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/lt" stack_frame_id=543} + %lt.627 = s32[41]{0} reshape(%lt.626), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/lt" stack_frame_id=543} + %lt.628 = s32[41,41]{1,0} broadcast(%lt.627), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/lt" stack_frame_id=543} + %lt.629 = pred[41,41]{1,0} compare(%lt.625, %lt.628), direction=LT, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/lt" stack_frame_id=543} + %convert_element_type.1732 = s32[41,41]{1,0} convert(%lt.629), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/convert_element_type" stack_frame_id=547} + %broadcast_in_dim.1496 = s32[1,41,41]{2,1,0} reshape(%convert_element_type.1732), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/broadcast_in_dim" stack_frame_id=551} + %min.95 = s32[1,41,41]{2,1,0} minimum(%convert_element_type.1733, %broadcast_in_dim.1496), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/min" stack_frame_id=551} + %broadcast_in_dim.1503 = s32[1,1,41,41]{3,2,1,0} reshape(%min.95), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/broadcast_in_dim" stack_frame_id=728} + %constant.141 = s32[] constant(0) + %broadcast.49 = s32[1,1,41,41]{3,2,1,0} broadcast(%constant.141), dimensions={} + %convert_element_type.1741 = pred[1,1,41,41]{3,2,1,0} compare(%broadcast_in_dim.1503, %broadcast.49), direction=NE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/convert_element_type" stack_frame_id=732} + %broadcast_in_dim.1504 = pred[1,1,1,41,41]{4,3,2,1,0} reshape(%convert_element_type.1741), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/broadcast_in_dim" stack_frame_id=34} + %and.298 = pred[1,1,1,41,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1504), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/vmap()/and" stack_frame_id=34} + %and.299 = pred[1,1,41,41]{3,2,1,0} reshape(%and.298), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/vmap()/and" stack_frame_id=34} + %and.300 = pred[1,1,32,41,41]{4,3,2,1,0} broadcast(%and.299), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/vmap()/and" stack_frame_id=34} + %state_1__2_.1 = bf16[4096,32,128]{2,1,0} parameter(5), metadata={op_name="state[1][2]"} + %dot_general.1050 = bf16[1,41,32,128]{3,2,1,0} dot(%convert_element_type.1736, %state_1__2_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=577} + %iota.515 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/iota" stack_frame_id=598} + %broadcast_in_dim.1499 = f32[1,41]{1,0} reshape(%iota.515), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/broadcast_in_dim" stack_frame_id=601} + %iota.514 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/iota" stack_frame_id=588} + %mul.3083 = f32[64]{0} multiply(%iota.514, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=588} + %div.970 = f32[64]{0} divide(%mul.3083, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/div" stack_frame_id=589} + %neg.390 = f32[64]{0} negate(%div.970), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/neg" stack_frame_id=590} + %pow.385 = f32[64]{0} power(%broadcast.51, %neg.390), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/pow" stack_frame_id=593} + %div.971 = f32[64]{0} divide(%pow.385, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/div" stack_frame_id=594} + %dot_general.1051 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1499, %div.971), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/bi,j->bij/dot_general" stack_frame_id=604} + %stack.1533 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1051), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/stack" stack_frame_id=607} + %stack.1534 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1051), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/stack" stack_frame_id=607} + %stack.1535 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1533, %stack.1534), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/stack" stack_frame_id=607} + %reshape.698 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1535), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/reshape"} + %cos.191 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.698), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/cos" stack_frame_id=611} + %convert_element_type.1737 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.191), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/convert_element_type" stack_frame_id=615} + %mul.3084 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1737), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=635} + %mul.3085 = bf16[1,41,128]{2,1,0} reshape(%mul.3084), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=635} + %mul.3086 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3085), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=635} + %mul.3087 = bf16[1,41,32,128]{3,2,1,0} multiply(%dot_general.1050, %mul.3086), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=635} + %split.383 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1050), slice={[0:1], [0:41], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/split" stack_frame_id=627} + %neg.391 = bf16[1,41,32,64]{3,2,1,0} negate(%split.383), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/neg" stack_frame_id=628} + %stack.1536 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%neg.391), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/stack" stack_frame_id=631} + %split.382 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1050), slice={[0:1], [0:41], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/split" stack_frame_id=627} + %stack.1537 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%split.382), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/stack" stack_frame_id=631} + %stack.1538 = bf16[1,41,32,2,64]{4,3,2,1,0} concatenate(%stack.1536, %stack.1537), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/stack" stack_frame_id=631} + %reshape.699 = bf16[1,41,32,128]{3,2,1,0} reshape(%stack.1538), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/reshape" stack_frame_id=634} + %sin.191 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.698), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/sin" stack_frame_id=619} + %convert_element_type.1738 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.191), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/convert_element_type" stack_frame_id=623} + %mul.3088 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1738), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=636} + %mul.3089 = bf16[1,41,128]{2,1,0} reshape(%mul.3088), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=636} + %mul.3090 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3089), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=636} + %mul.3091 = bf16[1,41,32,128]{3,2,1,0} multiply(%reshape.699, %mul.3090), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/mul" stack_frame_id=636} + %add.1125 = bf16[1,41,32,128]{3,2,1,0} add(%mul.3087, %mul.3091), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/rotary_embedding/add" stack_frame_id=637} + %reshape.704 = bf16[1,41,32,1,128]{4,3,2,1,0} reshape(%add.1125), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/reshape" stack_frame_id=34} + %broadcast_in_dim.1501 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%add.1126), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/broadcast_in_dim" stack_frame_id=722} + %reshape.702 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1501), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/reshape" stack_frame_id=722} + %dot_general.1055 = f32[1,32,41,1,41]{4,3,2,1,0} dot(%reshape.704, %reshape.702), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=34} + %constant.140 = f32[] constant(0.0883883461) + %broadcast.48 = f32[1,32,41,1,41]{4,3,2,1,0} broadcast(%constant.140), dimensions={} + %mul.3101 = f32[1,32,41,1,41]{4,3,2,1,0} multiply(%dot_general.1055, %broadcast.48), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/vmap()/mul" stack_frame_id=34} + %constant.153 = f32[] constant(-2.38197633e+38) + %vmap_jit__where__.95 = f32[1,1,32,41,41]{4,3,2,1,0} call(%and.300, %mul.3101, %constant.153), to_apply=%_where_0.5, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/vmap(jit(_where))" stack_frame_id=34} + %constant.152 = f32[] constant(-inf) + %reduce_max.481 = f32[1,1,32,41]{3,2,1,0} reduce(%vmap_jit__where__.95, %constant.152), dimensions={4}, to_apply=%region_2.6, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/vmap()/reduce_max" stack_frame_id=34} + %constant.139 = f32[] constant(-inf) + %broadcast.47 = f32[1,1,32,41]{3,2,1,0} broadcast(%constant.139), dimensions={} + %max.105 = f32[1,1,32,41]{3,2,1,0} maximum(%reduce_max.481, %broadcast.47), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/vmap()/max" stack_frame_id=34} + %broadcast_in_dim.1505 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%max.105), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %sub.404 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1505), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/vmap()/sub" stack_frame_id=34} + %sub.405 = f32[1,1,32,41]{3,2,1,0} reshape(%sub.404), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/vmap()/sub" stack_frame_id=34} + %sub.406 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%sub.405), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/vmap()/sub" stack_frame_id=34} + %sub.407 = f32[1,1,32,41,41]{4,3,2,1,0} subtract(%vmap_jit__where__.95, %sub.406), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/vmap()/sub" stack_frame_id=34} + %exp.101 = f32[1,1,32,41,41]{4,3,2,1,0} exponential(%sub.407), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/vmap()/exp" stack_frame_id=34} + %reduce_sum.1450 = f32[1,1,32,41]{3,2,1,0} reduce(%exp.101, %constant.155), dimensions={4}, to_apply=%region_3.7, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/vmap()/reduce_sum" stack_frame_id=34} + %broadcast_in_dim.1506 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%reduce_sum.1450), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %div.974 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1506), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/vmap()/div" stack_frame_id=34} + %div.975 = f32[1,1,32,41]{3,2,1,0} reshape(%div.974), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/vmap()/div" stack_frame_id=34} + %div.976 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%div.975), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/vmap()/div" stack_frame_id=34} + %div.977 = f32[1,1,32,41,41]{4,3,2,1,0} divide(%exp.101, %div.976), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/vmap()/div" stack_frame_id=34} + %convert_element_type.1742 = bf16[1,1,32,41,41]{4,3,2,1,0} convert(%div.977), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/vmap()/convert_element_type" stack_frame_id=34} + %dot_general.1056 = bf16[1,32,128,1,41]{4,3,2,1,0} dot(%reshape.703, %convert_element_type.1742), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=34} + %transpose.99 = bf16[1,41,32,1,128]{1,3,4,2,0} transpose(%dot_general.1056), dimensions={0,4,1,3,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/vmap()/transpose;jit(compiled_generate_function)/transformer_layer_0/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %reshape.705 = bf16[1,41,32,128]{3,2,1,0} reshape(%transpose.99), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/reshape" stack_frame_id=34} + %state_1__5_.1 = bf16[32,128,4096]{2,1,0} parameter(8), metadata={op_name="state[1][5]"} + %dot_general.1057 = bf16[1,41,4096]{2,1,0} dot(%reshape.705, %state_1__5_.1), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=741} + %add.1127 = bf16[1,41,4096]{2,1,0} add(%dot_general.1057, %jit__take_.3), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/add" stack_frame_id=742} + %convert_element_type.1743 = f32[1,41,4096]{2,1,0} convert(%add.1127), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/feedforward_layernorm/convert_element_type" stack_frame_id=746} + %pow.387 = f32[1,41,4096]{2,1,0} power(%convert_element_type.1743, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/feedforward_layernorm/pow" stack_frame_id=749} + %reduce_sum.1451 = f32[1,41]{1,0} reduce(%pow.387, %constant.155), dimensions={2}, to_apply=%region_4.8, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/feedforward_layernorm/reduce_sum" stack_frame_id=43} + %broadcast_in_dim.1507 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1451), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/feedforward_layernorm/broadcast_in_dim" stack_frame_id=43} + %div.978 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1507, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/feedforward_layernorm/div" stack_frame_id=43} + %add.1128 = f32[1,41,1]{2,1,0} add(%div.978, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/feedforward_layernorm/add" stack_frame_id=750} + %rsqrt.194 = f32[1,41,1]{2,1,0} rsqrt(%add.1128), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/feedforward_layernorm/rsqrt" stack_frame_id=753} + %mul.3102 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.194), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3103 = f32[1,41]{1,0} reshape(%mul.3102), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3104 = f32[1,41,4096]{2,1,0} broadcast(%mul.3103), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3105 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.1743, %mul.3104), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/feedforward_layernorm/mul" stack_frame_id=754} + %state_1__10_.1 = bf16[4096]{0} parameter(13), metadata={op_name="state[1][10]"} + %convert_element_type.1744 = f32[4096]{0} convert(%state_1__10_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/feedforward_layernorm/convert_element_type" stack_frame_id=755} + %broadcast_in_dim.1508 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1744), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/feedforward_layernorm/broadcast_in_dim" stack_frame_id=755} + %mul.3106 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1508), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3107 = f32[1,4096]{1,0} reshape(%mul.3106), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3108 = f32[1,41,4096]{2,1,0} broadcast(%mul.3107), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3109 = f32[1,41,4096]{2,1,0} multiply(%mul.3105, %mul.3108), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/feedforward_layernorm/mul" stack_frame_id=755} + %convert_element_type.1745 = bf16[1,41,4096]{2,1,0} convert(%mul.3109), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/feedforward_layernorm/convert_element_type" stack_frame_id=759} + %state_1__7_.1 = bf16[4096,14336]{1,0} parameter(10), metadata={op_name="state[1][7]"} + %dot_general.1059 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.1745, %state_1__7_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/feedforward_intermediate_dense/dot_general" stack_frame_id=788} + %state_1__8_.1 = bf16[4096,14336]{1,0} parameter(11), metadata={op_name="state[1][8]"} + %dot_general.1058 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.1745, %state_1__8_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/feedforward_gate_dense/dot_general" stack_frame_id=768} + %convert_element_type.1746 = f32[1,41,14336]{2,1,0} convert(%dot_general.1058), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/convert_element_type" stack_frame_id=772} + %jit_silu_.95 = f32[1,41,14336]{2,1,0} call(%convert_element_type.1746), to_apply=%silu.9, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/jit(silu)" stack_frame_id=775} + %convert_element_type.1747 = bf16[1,41,14336]{2,1,0} convert(%jit_silu_.95), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/convert_element_type" stack_frame_id=779} + %mul.3110 = bf16[1,41,14336]{2,1,0} multiply(%dot_general.1059, %convert_element_type.1747), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/mul" stack_frame_id=791} + %state_1__9_.1 = bf16[14336,4096]{1,0} parameter(12), metadata={op_name="state[1][9]"} + %dot_general.1060 = bf16[1,41,4096]{2,1,0} dot(%mul.3110, %state_1__9_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/feedforward_output_dense/dot_general" stack_frame_id=800} + %add.1129 = bf16[1,41,4096]{2,1,0} add(%dot_general.1060, %add.1127), metadata={op_name="jit(compiled_generate_function)/transformer_layer_0/add" stack_frame_id=801} + %convert_element_type.1750 = f32[1,41,4096]{2,1,0} convert(%add.1129), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention_layernorm/convert_element_type" stack_frame_id=555} + %pow.388 = f32[1,41,4096]{2,1,0} power(%convert_element_type.1750, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention_layernorm/pow" stack_frame_id=558} + %reduce_sum.1452 = f32[1,41]{1,0} reduce(%pow.388, %constant.155), dimensions={2}, to_apply=%region_5.10, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention_layernorm/reduce_sum" stack_frame_id=24} + %broadcast_in_dim.1515 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1452), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention_layernorm/broadcast_in_dim" stack_frame_id=24} + %div.979 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1515, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention_layernorm/div" stack_frame_id=24} + %add.1131 = f32[1,41,1]{2,1,0} add(%div.979, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention_layernorm/add" stack_frame_id=559} + %rsqrt.195 = f32[1,41,1]{2,1,0} rsqrt(%add.1131), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention_layernorm/rsqrt" stack_frame_id=562} + %mul.3111 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.195), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3112 = f32[1,41]{1,0} reshape(%mul.3111), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3113 = f32[1,41,4096]{2,1,0} broadcast(%mul.3112), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3114 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.1750, %mul.3113), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention_layernorm/mul" stack_frame_id=563} + %state_1__15_.1 = bf16[4096]{0} parameter(18), metadata={op_name="state[1][15]"} + %convert_element_type.1751 = f32[4096]{0} convert(%state_1__15_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention_layernorm/convert_element_type" stack_frame_id=564} + %broadcast_in_dim.1516 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1751), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention_layernorm/broadcast_in_dim" stack_frame_id=564} + %mul.3115 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1516), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3116 = f32[1,4096]{1,0} reshape(%mul.3115), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3117 = f32[1,41,4096]{2,1,0} broadcast(%mul.3116), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3118 = f32[1,41,4096]{2,1,0} multiply(%mul.3114, %mul.3117), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention_layernorm/mul" stack_frame_id=564} + %convert_element_type.1752 = bf16[1,41,4096]{2,1,0} convert(%mul.3118), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention_layernorm/convert_element_type" stack_frame_id=568} + %state_1__12_.1 = bf16[4096,8,128]{2,1,0} parameter(15), metadata={op_name="state[1][12]"} + %dot_general.1063 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.1752, %state_1__12_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=647} + %iota.525 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/iota" stack_frame_id=677} + %broadcast_in_dim.1518 = f32[1,41]{1,0} reshape(%iota.525), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/broadcast_in_dim" stack_frame_id=680} + %iota.524 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/iota" stack_frame_id=667} + %mul.3128 = f32[64]{0} multiply(%iota.524, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=667} + %div.982 = f32[64]{0} divide(%mul.3128, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/div" stack_frame_id=668} + %neg.396 = f32[64]{0} negate(%div.982), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/neg" stack_frame_id=669} + %pow.390 = f32[64]{0} power(%broadcast.51, %neg.396), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/pow" stack_frame_id=672} + %div.983 = f32[64]{0} divide(%pow.390, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/div" stack_frame_id=673} + %dot_general.1065 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1518, %div.983), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/bi,j->bij/dot_general" stack_frame_id=683} + %stack.1554 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1065), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/stack" stack_frame_id=686} + %stack.1555 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1065), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/stack" stack_frame_id=686} + %stack.1556 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1554, %stack.1555), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/stack" stack_frame_id=686} + %reshape.708 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1556), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/reshape"} + %cos.194 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.708), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/cos" stack_frame_id=690} + %convert_element_type.1755 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.194), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/convert_element_type" stack_frame_id=694} + %mul.3129 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1755), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=714} + %mul.3130 = bf16[1,41,128]{2,1,0} reshape(%mul.3129), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=714} + %mul.3131 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3130), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=714} + %mul.3132 = bf16[1,41,8,128]{3,2,1,0} multiply(%dot_general.1063, %mul.3131), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=714} + %split.389 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1063), slice={[0:1], [0:41], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/split" stack_frame_id=706} + %neg.397 = bf16[1,41,8,64]{3,2,1,0} negate(%split.389), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/neg" stack_frame_id=707} + %stack.1557 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%neg.397), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/stack" stack_frame_id=710} + %split.388 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1063), slice={[0:1], [0:41], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/split" stack_frame_id=706} + %stack.1558 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%split.388), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/stack" stack_frame_id=710} + %stack.1559 = bf16[1,41,8,2,64]{4,3,2,1,0} concatenate(%stack.1557, %stack.1558), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/stack" stack_frame_id=710} + %reshape.709 = bf16[1,41,8,128]{3,2,1,0} reshape(%stack.1559), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/reshape" stack_frame_id=713} + %sin.194 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.708), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/sin" stack_frame_id=698} + %convert_element_type.1756 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.194), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/convert_element_type" stack_frame_id=702} + %mul.3133 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1756), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=715} + %mul.3134 = bf16[1,41,128]{2,1,0} reshape(%mul.3133), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=715} + %mul.3135 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3134), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=715} + %mul.3136 = bf16[1,41,8,128]{3,2,1,0} multiply(%reshape.709, %mul.3135), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=715} + %add.1133 = bf16[1,41,8,128]{3,2,1,0} add(%mul.3132, %mul.3136), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/add" stack_frame_id=716} + %stack.1560 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%add.1133), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/stack" stack_frame_id=719} + %state_1__13_.1 = bf16[4096,8,128]{2,1,0} parameter(16), metadata={op_name="state[1][13]"} + %dot_general.1064 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.1752, %state_1__13_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=656} + %stack.1561 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dot_general.1064), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/stack" stack_frame_id=719} + %stack.1562 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1560, %stack.1561), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/stack" stack_frame_id=719} + %stack.2008 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1562), metadata={op_name="jit(compiled_generate_function)/stack" stack_frame_id=804} + %broadcast_in_dim.1520 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dot_general.1064), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/broadcast_in_dim" stack_frame_id=725} + %reshape.711 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1520), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/reshape" stack_frame_id=725} + %iota.518 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/iota" stack_frame_id=525} + %broadcast_in_dim.1509 = f32[41,1]{1,0} reshape(%iota.518), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/broadcast_in_dim" stack_frame_id=528} + %ge.488 = f32[41,1]{1,0} broadcast(%broadcast_in_dim.1509), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/ge" stack_frame_id=532} + %ge.489 = f32[41]{0} reshape(%ge.488), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/ge" stack_frame_id=532} + %ge.490 = f32[41,41]{1,0} broadcast(%ge.489), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/ge" stack_frame_id=532} + %iota.519 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/iota" stack_frame_id=531} + %broadcast_in_dim.1510 = f32[1,41]{1,0} reshape(%iota.519), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/broadcast_in_dim" stack_frame_id=532} + %ge.491 = f32[1,41]{1,0} broadcast(%broadcast_in_dim.1510), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/ge" stack_frame_id=532} + %ge.492 = f32[41]{0} reshape(%ge.491), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/ge" stack_frame_id=532} + %ge.493 = f32[41,41]{1,0} broadcast(%ge.492), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/ge" stack_frame_id=532} + %ge.494 = pred[41,41]{1,0} compare(%ge.490, %ge.493), direction=GE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/ge" stack_frame_id=532} + %broadcast_in_dim.1511 = pred[1,41,41]{2,1,0} reshape(%ge.494), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/broadcast_in_dim" stack_frame_id=535} + %convert_element_type.1749 = s32[1,41,41]{2,1,0} convert(%broadcast_in_dim.1511), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/convert_element_type" stack_frame_id=551} + %iota.520 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/iota" stack_frame_id=538} + %broadcast_in_dim.1512 = s32[41,1]{1,0} reshape(%iota.520), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/broadcast_in_dim" stack_frame_id=536} + %lt.630 = s32[41,1]{1,0} broadcast(%broadcast_in_dim.1512), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/lt" stack_frame_id=543} + %lt.631 = s32[41]{0} reshape(%lt.630), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/lt" stack_frame_id=543} + %lt.632 = s32[41,41]{1,0} broadcast(%lt.631), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/lt" stack_frame_id=543} + %iota.521 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/iota" stack_frame_id=541} + %broadcast_in_dim.1513 = s32[1,41]{1,0} reshape(%iota.521), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/broadcast_in_dim" stack_frame_id=539} + %add.1130 = s32[1,41]{1,0} add(%broadcast_in_dim.1513, %broadcast.57), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/add" stack_frame_id=542} + %lt.633 = s32[1,41]{1,0} broadcast(%add.1130), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/lt" stack_frame_id=543} + %lt.634 = s32[41]{0} reshape(%lt.633), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/lt" stack_frame_id=543} + %lt.635 = s32[41,41]{1,0} broadcast(%lt.634), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/lt" stack_frame_id=543} + %lt.636 = pred[41,41]{1,0} compare(%lt.632, %lt.635), direction=LT, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/lt" stack_frame_id=543} + %convert_element_type.1748 = s32[41,41]{1,0} convert(%lt.636), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/convert_element_type" stack_frame_id=547} + %broadcast_in_dim.1514 = s32[1,41,41]{2,1,0} reshape(%convert_element_type.1748), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/broadcast_in_dim" stack_frame_id=551} + %min.96 = s32[1,41,41]{2,1,0} minimum(%convert_element_type.1749, %broadcast_in_dim.1514), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/min" stack_frame_id=551} + %broadcast_in_dim.1521 = s32[1,1,41,41]{3,2,1,0} reshape(%min.96), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/broadcast_in_dim" stack_frame_id=728} + %convert_element_type.1757 = pred[1,1,41,41]{3,2,1,0} compare(%broadcast_in_dim.1521, %broadcast.49), direction=NE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/convert_element_type" stack_frame_id=732} + %broadcast_in_dim.1522 = pred[1,1,1,41,41]{4,3,2,1,0} reshape(%convert_element_type.1757), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/broadcast_in_dim" stack_frame_id=34} + %and.301 = pred[1,1,1,41,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1522), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/vmap()/and" stack_frame_id=34} + %and.302 = pred[1,1,41,41]{3,2,1,0} reshape(%and.301), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/vmap()/and" stack_frame_id=34} + %and.303 = pred[1,1,32,41,41]{4,3,2,1,0} broadcast(%and.302), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/vmap()/and" stack_frame_id=34} + %state_1__11_.1 = bf16[4096,32,128]{2,1,0} parameter(14), metadata={op_name="state[1][11]"} + %dot_general.1061 = bf16[1,41,32,128]{3,2,1,0} dot(%convert_element_type.1752, %state_1__11_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=577} + %iota.523 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/iota" stack_frame_id=598} + %broadcast_in_dim.1517 = f32[1,41]{1,0} reshape(%iota.523), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/broadcast_in_dim" stack_frame_id=601} + %iota.522 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/iota" stack_frame_id=588} + %mul.3119 = f32[64]{0} multiply(%iota.522, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=588} + %div.980 = f32[64]{0} divide(%mul.3119, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/div" stack_frame_id=589} + %neg.394 = f32[64]{0} negate(%div.980), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/neg" stack_frame_id=590} + %pow.389 = f32[64]{0} power(%broadcast.51, %neg.394), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/pow" stack_frame_id=593} + %div.981 = f32[64]{0} divide(%pow.389, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/div" stack_frame_id=594} + %dot_general.1062 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1517, %div.981), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/bi,j->bij/dot_general" stack_frame_id=604} + %stack.1548 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1062), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/stack" stack_frame_id=607} + %stack.1549 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1062), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/stack" stack_frame_id=607} + %stack.1550 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1548, %stack.1549), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/stack" stack_frame_id=607} + %reshape.706 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1550), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/reshape"} + %cos.193 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.706), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/cos" stack_frame_id=611} + %convert_element_type.1753 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.193), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/convert_element_type" stack_frame_id=615} + %mul.3120 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1753), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=635} + %mul.3121 = bf16[1,41,128]{2,1,0} reshape(%mul.3120), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=635} + %mul.3122 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3121), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=635} + %mul.3123 = bf16[1,41,32,128]{3,2,1,0} multiply(%dot_general.1061, %mul.3122), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=635} + %split.387 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1061), slice={[0:1], [0:41], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/split" stack_frame_id=627} + %neg.395 = bf16[1,41,32,64]{3,2,1,0} negate(%split.387), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/neg" stack_frame_id=628} + %stack.1551 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%neg.395), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/stack" stack_frame_id=631} + %split.386 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1061), slice={[0:1], [0:41], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/split" stack_frame_id=627} + %stack.1552 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%split.386), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/stack" stack_frame_id=631} + %stack.1553 = bf16[1,41,32,2,64]{4,3,2,1,0} concatenate(%stack.1551, %stack.1552), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/stack" stack_frame_id=631} + %reshape.707 = bf16[1,41,32,128]{3,2,1,0} reshape(%stack.1553), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/reshape" stack_frame_id=634} + %sin.193 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.706), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/sin" stack_frame_id=619} + %convert_element_type.1754 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.193), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/convert_element_type" stack_frame_id=623} + %mul.3124 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1754), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=636} + %mul.3125 = bf16[1,41,128]{2,1,0} reshape(%mul.3124), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=636} + %mul.3126 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3125), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=636} + %mul.3127 = bf16[1,41,32,128]{3,2,1,0} multiply(%reshape.707, %mul.3126), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/mul" stack_frame_id=636} + %add.1132 = bf16[1,41,32,128]{3,2,1,0} add(%mul.3123, %mul.3127), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/rotary_embedding_1/add" stack_frame_id=637} + %reshape.712 = bf16[1,41,32,1,128]{4,3,2,1,0} reshape(%add.1132), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/reshape" stack_frame_id=34} + %broadcast_in_dim.1519 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%add.1133), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/broadcast_in_dim" stack_frame_id=722} + %reshape.710 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1519), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/reshape" stack_frame_id=722} + %dot_general.1066 = f32[1,32,41,1,41]{4,3,2,1,0} dot(%reshape.712, %reshape.710), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=34} + %mul.3137 = f32[1,32,41,1,41]{4,3,2,1,0} multiply(%dot_general.1066, %broadcast.48), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/vmap()/mul" stack_frame_id=34} + %vmap_jit__where__.96 = f32[1,1,32,41,41]{4,3,2,1,0} call(%and.303, %mul.3137, %constant.153), to_apply=%_where_0.5, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/vmap(jit(_where))" stack_frame_id=34} + %reduce_max.482 = f32[1,1,32,41]{3,2,1,0} reduce(%vmap_jit__where__.96, %constant.152), dimensions={4}, to_apply=%region_6.11, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/vmap()/reduce_max" stack_frame_id=34} + %max.106 = f32[1,1,32,41]{3,2,1,0} maximum(%reduce_max.482, %broadcast.47), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/vmap()/max" stack_frame_id=34} + %broadcast_in_dim.1523 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%max.106), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %sub.408 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1523), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/vmap()/sub" stack_frame_id=34} + %sub.409 = f32[1,1,32,41]{3,2,1,0} reshape(%sub.408), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/vmap()/sub" stack_frame_id=34} + %sub.410 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%sub.409), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/vmap()/sub" stack_frame_id=34} + %sub.411 = f32[1,1,32,41,41]{4,3,2,1,0} subtract(%vmap_jit__where__.96, %sub.410), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/vmap()/sub" stack_frame_id=34} + %exp.102 = f32[1,1,32,41,41]{4,3,2,1,0} exponential(%sub.411), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/vmap()/exp" stack_frame_id=34} + %reduce_sum.1453 = f32[1,1,32,41]{3,2,1,0} reduce(%exp.102, %constant.155), dimensions={4}, to_apply=%region_7.12, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/vmap()/reduce_sum" stack_frame_id=34} + %broadcast_in_dim.1524 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%reduce_sum.1453), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %div.984 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1524), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/vmap()/div" stack_frame_id=34} + %div.985 = f32[1,1,32,41]{3,2,1,0} reshape(%div.984), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/vmap()/div" stack_frame_id=34} + %div.986 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%div.985), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/vmap()/div" stack_frame_id=34} + %div.987 = f32[1,1,32,41,41]{4,3,2,1,0} divide(%exp.102, %div.986), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/vmap()/div" stack_frame_id=34} + %convert_element_type.1758 = bf16[1,1,32,41,41]{4,3,2,1,0} convert(%div.987), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/vmap()/convert_element_type" stack_frame_id=34} + %dot_general.1067 = bf16[1,32,128,1,41]{4,3,2,1,0} dot(%reshape.711, %convert_element_type.1758), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=34} + %transpose.100 = bf16[1,41,32,1,128]{1,3,4,2,0} transpose(%dot_general.1067), dimensions={0,4,1,3,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/vmap()/transpose;jit(compiled_generate_function)/transformer_layer_1/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %reshape.713 = bf16[1,41,32,128]{3,2,1,0} reshape(%transpose.100), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/reshape" stack_frame_id=34} + %state_1__14_.1 = bf16[32,128,4096]{2,1,0} parameter(17), metadata={op_name="state[1][14]"} + %dot_general.1068 = bf16[1,41,4096]{2,1,0} dot(%reshape.713, %state_1__14_.1), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=741} + %add.1134 = bf16[1,41,4096]{2,1,0} add(%dot_general.1068, %add.1129), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/add" stack_frame_id=742} + %convert_element_type.1759 = f32[1,41,4096]{2,1,0} convert(%add.1134), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/feedforward_layernorm/convert_element_type" stack_frame_id=746} + %pow.391 = f32[1,41,4096]{2,1,0} power(%convert_element_type.1759, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/feedforward_layernorm/pow" stack_frame_id=749} + %reduce_sum.1454 = f32[1,41]{1,0} reduce(%pow.391, %constant.155), dimensions={2}, to_apply=%region_8.13, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/feedforward_layernorm/reduce_sum" stack_frame_id=43} + %broadcast_in_dim.1525 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1454), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/feedforward_layernorm/broadcast_in_dim" stack_frame_id=43} + %div.988 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1525, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/feedforward_layernorm/div" stack_frame_id=43} + %add.1135 = f32[1,41,1]{2,1,0} add(%div.988, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/feedforward_layernorm/add" stack_frame_id=750} + %rsqrt.196 = f32[1,41,1]{2,1,0} rsqrt(%add.1135), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/feedforward_layernorm/rsqrt" stack_frame_id=753} + %mul.3138 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.196), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3139 = f32[1,41]{1,0} reshape(%mul.3138), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3140 = f32[1,41,4096]{2,1,0} broadcast(%mul.3139), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3141 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.1759, %mul.3140), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/feedforward_layernorm/mul" stack_frame_id=754} + %state_1__19_.1 = bf16[4096]{0} parameter(22), metadata={op_name="state[1][19]"} + %convert_element_type.1760 = f32[4096]{0} convert(%state_1__19_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/feedforward_layernorm/convert_element_type" stack_frame_id=755} + %broadcast_in_dim.1526 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1760), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/feedforward_layernorm/broadcast_in_dim" stack_frame_id=755} + %mul.3142 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1526), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3143 = f32[1,4096]{1,0} reshape(%mul.3142), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3144 = f32[1,41,4096]{2,1,0} broadcast(%mul.3143), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3145 = f32[1,41,4096]{2,1,0} multiply(%mul.3141, %mul.3144), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/feedforward_layernorm/mul" stack_frame_id=755} + %convert_element_type.1761 = bf16[1,41,4096]{2,1,0} convert(%mul.3145), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/feedforward_layernorm/convert_element_type" stack_frame_id=759} + %state_1__16_.1 = bf16[4096,14336]{1,0} parameter(19), metadata={op_name="state[1][16]"} + %dot_general.1070 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.1761, %state_1__16_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/feedforward_intermediate_dense/dot_general" stack_frame_id=788} + %state_1__17_.1 = bf16[4096,14336]{1,0} parameter(20), metadata={op_name="state[1][17]"} + %dot_general.1069 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.1761, %state_1__17_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/feedforward_gate_dense/dot_general" stack_frame_id=768} + %convert_element_type.1762 = f32[1,41,14336]{2,1,0} convert(%dot_general.1069), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/convert_element_type" stack_frame_id=772} + %jit_silu_.96 = f32[1,41,14336]{2,1,0} call(%convert_element_type.1762), to_apply=%silu.9, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/jit(silu)" stack_frame_id=775} + %convert_element_type.1763 = bf16[1,41,14336]{2,1,0} convert(%jit_silu_.96), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/convert_element_type" stack_frame_id=779} + %mul.3146 = bf16[1,41,14336]{2,1,0} multiply(%dot_general.1070, %convert_element_type.1763), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/mul" stack_frame_id=791} + %state_1__18_.1 = bf16[14336,4096]{1,0} parameter(21), metadata={op_name="state[1][18]"} + %dot_general.1071 = bf16[1,41,4096]{2,1,0} dot(%mul.3146, %state_1__18_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/feedforward_output_dense/dot_general" stack_frame_id=800} + %add.1136 = bf16[1,41,4096]{2,1,0} add(%dot_general.1071, %add.1134), metadata={op_name="jit(compiled_generate_function)/transformer_layer_1/add" stack_frame_id=801} + %convert_element_type.1766 = f32[1,41,4096]{2,1,0} convert(%add.1136), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention_layernorm/convert_element_type" stack_frame_id=555} + %pow.392 = f32[1,41,4096]{2,1,0} power(%convert_element_type.1766, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention_layernorm/pow" stack_frame_id=558} + %reduce_sum.1455 = f32[1,41]{1,0} reduce(%pow.392, %constant.155), dimensions={2}, to_apply=%region_9.14, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention_layernorm/reduce_sum" stack_frame_id=24} + %broadcast_in_dim.1533 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1455), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention_layernorm/broadcast_in_dim" stack_frame_id=24} + %div.989 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1533, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention_layernorm/div" stack_frame_id=24} + %add.1138 = f32[1,41,1]{2,1,0} add(%div.989, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention_layernorm/add" stack_frame_id=559} + %rsqrt.197 = f32[1,41,1]{2,1,0} rsqrt(%add.1138), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention_layernorm/rsqrt" stack_frame_id=562} + %mul.3147 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.197), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3148 = f32[1,41]{1,0} reshape(%mul.3147), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3149 = f32[1,41,4096]{2,1,0} broadcast(%mul.3148), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3150 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.1766, %mul.3149), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention_layernorm/mul" stack_frame_id=563} + %state_1__24_.1 = bf16[4096]{0} parameter(27), metadata={op_name="state[1][24]"} + %convert_element_type.1767 = f32[4096]{0} convert(%state_1__24_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention_layernorm/convert_element_type" stack_frame_id=564} + %broadcast_in_dim.1534 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1767), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention_layernorm/broadcast_in_dim" stack_frame_id=564} + %mul.3151 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1534), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3152 = f32[1,4096]{1,0} reshape(%mul.3151), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3153 = f32[1,41,4096]{2,1,0} broadcast(%mul.3152), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3154 = f32[1,41,4096]{2,1,0} multiply(%mul.3150, %mul.3153), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention_layernorm/mul" stack_frame_id=564} + %convert_element_type.1768 = bf16[1,41,4096]{2,1,0} convert(%mul.3154), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention_layernorm/convert_element_type" stack_frame_id=568} + %state_1__21_.1 = bf16[4096,8,128]{2,1,0} parameter(24), metadata={op_name="state[1][21]"} + %dot_general.1074 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.1768, %state_1__21_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=647} + %iota.533 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/iota" stack_frame_id=677} + %broadcast_in_dim.1536 = f32[1,41]{1,0} reshape(%iota.533), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/broadcast_in_dim" stack_frame_id=680} + %iota.532 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/iota" stack_frame_id=667} + %mul.3164 = f32[64]{0} multiply(%iota.532, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=667} + %div.992 = f32[64]{0} divide(%mul.3164, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/div" stack_frame_id=668} + %neg.400 = f32[64]{0} negate(%div.992), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/neg" stack_frame_id=669} + %pow.394 = f32[64]{0} power(%broadcast.51, %neg.400), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/pow" stack_frame_id=672} + %div.993 = f32[64]{0} divide(%pow.394, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/div" stack_frame_id=673} + %dot_general.1076 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1536, %div.993), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/bi,j->bij/dot_general" stack_frame_id=683} + %stack.1569 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1076), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/stack" stack_frame_id=686} + %stack.1570 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1076), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/stack" stack_frame_id=686} + %stack.1571 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1569, %stack.1570), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/stack" stack_frame_id=686} + %reshape.716 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1571), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/reshape"} + %cos.196 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.716), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/cos" stack_frame_id=690} + %convert_element_type.1771 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.196), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/convert_element_type" stack_frame_id=694} + %mul.3165 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1771), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=714} + %mul.3166 = bf16[1,41,128]{2,1,0} reshape(%mul.3165), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=714} + %mul.3167 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3166), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=714} + %mul.3168 = bf16[1,41,8,128]{3,2,1,0} multiply(%dot_general.1074, %mul.3167), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=714} + %split.393 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1074), slice={[0:1], [0:41], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/split" stack_frame_id=706} + %neg.401 = bf16[1,41,8,64]{3,2,1,0} negate(%split.393), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/neg" stack_frame_id=707} + %stack.1572 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%neg.401), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/stack" stack_frame_id=710} + %split.392 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1074), slice={[0:1], [0:41], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/split" stack_frame_id=706} + %stack.1573 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%split.392), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/stack" stack_frame_id=710} + %stack.1574 = bf16[1,41,8,2,64]{4,3,2,1,0} concatenate(%stack.1572, %stack.1573), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/stack" stack_frame_id=710} + %reshape.717 = bf16[1,41,8,128]{3,2,1,0} reshape(%stack.1574), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/reshape" stack_frame_id=713} + %sin.196 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.716), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/sin" stack_frame_id=698} + %convert_element_type.1772 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.196), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/convert_element_type" stack_frame_id=702} + %mul.3169 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1772), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=715} + %mul.3170 = bf16[1,41,128]{2,1,0} reshape(%mul.3169), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=715} + %mul.3171 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3170), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=715} + %mul.3172 = bf16[1,41,8,128]{3,2,1,0} multiply(%reshape.717, %mul.3171), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=715} + %add.1140 = bf16[1,41,8,128]{3,2,1,0} add(%mul.3168, %mul.3172), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/add" stack_frame_id=716} + %stack.1575 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%add.1140), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/stack" stack_frame_id=719} + %state_1__22_.1 = bf16[4096,8,128]{2,1,0} parameter(25), metadata={op_name="state[1][22]"} + %dot_general.1075 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.1768, %state_1__22_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=656} + %stack.1576 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dot_general.1075), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/stack" stack_frame_id=719} + %stack.1577 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1575, %stack.1576), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/stack" stack_frame_id=719} + %stack.2009 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1577), metadata={op_name="jit(compiled_generate_function)/stack" stack_frame_id=804} + %broadcast_in_dim.1538 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dot_general.1075), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/broadcast_in_dim" stack_frame_id=725} + %reshape.719 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1538), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/reshape" stack_frame_id=725} + %iota.526 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/iota" stack_frame_id=525} + %broadcast_in_dim.1527 = f32[41,1]{1,0} reshape(%iota.526), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/broadcast_in_dim" stack_frame_id=528} + %ge.495 = f32[41,1]{1,0} broadcast(%broadcast_in_dim.1527), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/ge" stack_frame_id=532} + %ge.496 = f32[41]{0} reshape(%ge.495), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/ge" stack_frame_id=532} + %ge.497 = f32[41,41]{1,0} broadcast(%ge.496), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/ge" stack_frame_id=532} + %iota.527 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/iota" stack_frame_id=531} + %broadcast_in_dim.1528 = f32[1,41]{1,0} reshape(%iota.527), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/broadcast_in_dim" stack_frame_id=532} + %ge.498 = f32[1,41]{1,0} broadcast(%broadcast_in_dim.1528), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/ge" stack_frame_id=532} + %ge.499 = f32[41]{0} reshape(%ge.498), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/ge" stack_frame_id=532} + %ge.500 = f32[41,41]{1,0} broadcast(%ge.499), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/ge" stack_frame_id=532} + %ge.501 = pred[41,41]{1,0} compare(%ge.497, %ge.500), direction=GE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/ge" stack_frame_id=532} + %broadcast_in_dim.1529 = pred[1,41,41]{2,1,0} reshape(%ge.501), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/broadcast_in_dim" stack_frame_id=535} + %convert_element_type.1765 = s32[1,41,41]{2,1,0} convert(%broadcast_in_dim.1529), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/convert_element_type" stack_frame_id=551} + %iota.528 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/iota" stack_frame_id=538} + %broadcast_in_dim.1530 = s32[41,1]{1,0} reshape(%iota.528), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/broadcast_in_dim" stack_frame_id=536} + %lt.637 = s32[41,1]{1,0} broadcast(%broadcast_in_dim.1530), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/lt" stack_frame_id=543} + %lt.638 = s32[41]{0} reshape(%lt.637), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/lt" stack_frame_id=543} + %lt.639 = s32[41,41]{1,0} broadcast(%lt.638), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/lt" stack_frame_id=543} + %iota.529 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/iota" stack_frame_id=541} + %broadcast_in_dim.1531 = s32[1,41]{1,0} reshape(%iota.529), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/broadcast_in_dim" stack_frame_id=539} + %add.1137 = s32[1,41]{1,0} add(%broadcast_in_dim.1531, %broadcast.57), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/add" stack_frame_id=542} + %lt.640 = s32[1,41]{1,0} broadcast(%add.1137), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/lt" stack_frame_id=543} + %lt.641 = s32[41]{0} reshape(%lt.640), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/lt" stack_frame_id=543} + %lt.642 = s32[41,41]{1,0} broadcast(%lt.641), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/lt" stack_frame_id=543} + %lt.643 = pred[41,41]{1,0} compare(%lt.639, %lt.642), direction=LT, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/lt" stack_frame_id=543} + %convert_element_type.1764 = s32[41,41]{1,0} convert(%lt.643), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/convert_element_type" stack_frame_id=547} + %broadcast_in_dim.1532 = s32[1,41,41]{2,1,0} reshape(%convert_element_type.1764), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/broadcast_in_dim" stack_frame_id=551} + %min.97 = s32[1,41,41]{2,1,0} minimum(%convert_element_type.1765, %broadcast_in_dim.1532), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/min" stack_frame_id=551} + %broadcast_in_dim.1539 = s32[1,1,41,41]{3,2,1,0} reshape(%min.97), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/broadcast_in_dim" stack_frame_id=728} + %convert_element_type.1773 = pred[1,1,41,41]{3,2,1,0} compare(%broadcast_in_dim.1539, %broadcast.49), direction=NE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/convert_element_type" stack_frame_id=732} + %broadcast_in_dim.1540 = pred[1,1,1,41,41]{4,3,2,1,0} reshape(%convert_element_type.1773), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/broadcast_in_dim" stack_frame_id=34} + %and.304 = pred[1,1,1,41,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1540), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/vmap()/and" stack_frame_id=34} + %and.305 = pred[1,1,41,41]{3,2,1,0} reshape(%and.304), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/vmap()/and" stack_frame_id=34} + %and.306 = pred[1,1,32,41,41]{4,3,2,1,0} broadcast(%and.305), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/vmap()/and" stack_frame_id=34} + %state_1__20_.1 = bf16[4096,32,128]{2,1,0} parameter(23), metadata={op_name="state[1][20]"} + %dot_general.1072 = bf16[1,41,32,128]{3,2,1,0} dot(%convert_element_type.1768, %state_1__20_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=577} + %iota.531 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/iota" stack_frame_id=598} + %broadcast_in_dim.1535 = f32[1,41]{1,0} reshape(%iota.531), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/broadcast_in_dim" stack_frame_id=601} + %iota.530 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/iota" stack_frame_id=588} + %mul.3155 = f32[64]{0} multiply(%iota.530, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=588} + %div.990 = f32[64]{0} divide(%mul.3155, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/div" stack_frame_id=589} + %neg.398 = f32[64]{0} negate(%div.990), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/neg" stack_frame_id=590} + %pow.393 = f32[64]{0} power(%broadcast.51, %neg.398), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/pow" stack_frame_id=593} + %div.991 = f32[64]{0} divide(%pow.393, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/div" stack_frame_id=594} + %dot_general.1073 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1535, %div.991), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/bi,j->bij/dot_general" stack_frame_id=604} + %stack.1563 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1073), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/stack" stack_frame_id=607} + %stack.1564 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1073), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/stack" stack_frame_id=607} + %stack.1565 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1563, %stack.1564), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/stack" stack_frame_id=607} + %reshape.714 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1565), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/reshape"} + %cos.195 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.714), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/cos" stack_frame_id=611} + %convert_element_type.1769 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.195), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/convert_element_type" stack_frame_id=615} + %mul.3156 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1769), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=635} + %mul.3157 = bf16[1,41,128]{2,1,0} reshape(%mul.3156), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=635} + %mul.3158 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3157), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=635} + %mul.3159 = bf16[1,41,32,128]{3,2,1,0} multiply(%dot_general.1072, %mul.3158), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=635} + %split.391 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1072), slice={[0:1], [0:41], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/split" stack_frame_id=627} + %neg.399 = bf16[1,41,32,64]{3,2,1,0} negate(%split.391), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/neg" stack_frame_id=628} + %stack.1566 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%neg.399), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/stack" stack_frame_id=631} + %split.390 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1072), slice={[0:1], [0:41], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/split" stack_frame_id=627} + %stack.1567 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%split.390), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/stack" stack_frame_id=631} + %stack.1568 = bf16[1,41,32,2,64]{4,3,2,1,0} concatenate(%stack.1566, %stack.1567), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/stack" stack_frame_id=631} + %reshape.715 = bf16[1,41,32,128]{3,2,1,0} reshape(%stack.1568), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/reshape" stack_frame_id=634} + %sin.195 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.714), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/sin" stack_frame_id=619} + %convert_element_type.1770 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.195), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/convert_element_type" stack_frame_id=623} + %mul.3160 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1770), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=636} + %mul.3161 = bf16[1,41,128]{2,1,0} reshape(%mul.3160), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=636} + %mul.3162 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3161), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=636} + %mul.3163 = bf16[1,41,32,128]{3,2,1,0} multiply(%reshape.715, %mul.3162), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/mul" stack_frame_id=636} + %add.1139 = bf16[1,41,32,128]{3,2,1,0} add(%mul.3159, %mul.3163), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/rotary_embedding_2/add" stack_frame_id=637} + %reshape.720 = bf16[1,41,32,1,128]{4,3,2,1,0} reshape(%add.1139), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/reshape" stack_frame_id=34} + %broadcast_in_dim.1537 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%add.1140), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/broadcast_in_dim" stack_frame_id=722} + %reshape.718 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1537), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/reshape" stack_frame_id=722} + %dot_general.1077 = f32[1,32,41,1,41]{4,3,2,1,0} dot(%reshape.720, %reshape.718), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=34} + %mul.3173 = f32[1,32,41,1,41]{4,3,2,1,0} multiply(%dot_general.1077, %broadcast.48), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/vmap()/mul" stack_frame_id=34} + %vmap_jit__where__.97 = f32[1,1,32,41,41]{4,3,2,1,0} call(%and.306, %mul.3173, %constant.153), to_apply=%_where_0.5, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/vmap(jit(_where))" stack_frame_id=34} + %reduce_max.483 = f32[1,1,32,41]{3,2,1,0} reduce(%vmap_jit__where__.97, %constant.152), dimensions={4}, to_apply=%region_10.15, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/vmap()/reduce_max" stack_frame_id=34} + %max.107 = f32[1,1,32,41]{3,2,1,0} maximum(%reduce_max.483, %broadcast.47), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/vmap()/max" stack_frame_id=34} + %broadcast_in_dim.1541 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%max.107), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %sub.412 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1541), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/vmap()/sub" stack_frame_id=34} + %sub.413 = f32[1,1,32,41]{3,2,1,0} reshape(%sub.412), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/vmap()/sub" stack_frame_id=34} + %sub.414 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%sub.413), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/vmap()/sub" stack_frame_id=34} + %sub.415 = f32[1,1,32,41,41]{4,3,2,1,0} subtract(%vmap_jit__where__.97, %sub.414), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/vmap()/sub" stack_frame_id=34} + %exp.103 = f32[1,1,32,41,41]{4,3,2,1,0} exponential(%sub.415), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/vmap()/exp" stack_frame_id=34} + %reduce_sum.1456 = f32[1,1,32,41]{3,2,1,0} reduce(%exp.103, %constant.155), dimensions={4}, to_apply=%region_11.16, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/vmap()/reduce_sum" stack_frame_id=34} + %broadcast_in_dim.1542 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%reduce_sum.1456), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %div.994 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1542), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/vmap()/div" stack_frame_id=34} + %div.995 = f32[1,1,32,41]{3,2,1,0} reshape(%div.994), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/vmap()/div" stack_frame_id=34} + %div.996 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%div.995), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/vmap()/div" stack_frame_id=34} + %div.997 = f32[1,1,32,41,41]{4,3,2,1,0} divide(%exp.103, %div.996), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/vmap()/div" stack_frame_id=34} + %convert_element_type.1774 = bf16[1,1,32,41,41]{4,3,2,1,0} convert(%div.997), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/vmap()/convert_element_type" stack_frame_id=34} + %dot_general.1078 = bf16[1,32,128,1,41]{4,3,2,1,0} dot(%reshape.719, %convert_element_type.1774), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=34} + %transpose.101 = bf16[1,41,32,1,128]{1,3,4,2,0} transpose(%dot_general.1078), dimensions={0,4,1,3,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/vmap()/transpose;jit(compiled_generate_function)/transformer_layer_2/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %reshape.721 = bf16[1,41,32,128]{3,2,1,0} reshape(%transpose.101), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/reshape" stack_frame_id=34} + %state_1__23_.1 = bf16[32,128,4096]{2,1,0} parameter(26), metadata={op_name="state[1][23]"} + %dot_general.1079 = bf16[1,41,4096]{2,1,0} dot(%reshape.721, %state_1__23_.1), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=741} + %add.1141 = bf16[1,41,4096]{2,1,0} add(%dot_general.1079, %add.1136), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/add" stack_frame_id=742} + %convert_element_type.1775 = f32[1,41,4096]{2,1,0} convert(%add.1141), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/feedforward_layernorm/convert_element_type" stack_frame_id=746} + %pow.395 = f32[1,41,4096]{2,1,0} power(%convert_element_type.1775, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/feedforward_layernorm/pow" stack_frame_id=749} + %reduce_sum.1457 = f32[1,41]{1,0} reduce(%pow.395, %constant.155), dimensions={2}, to_apply=%region_12.17, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/feedforward_layernorm/reduce_sum" stack_frame_id=43} + %broadcast_in_dim.1543 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1457), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/feedforward_layernorm/broadcast_in_dim" stack_frame_id=43} + %div.998 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1543, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/feedforward_layernorm/div" stack_frame_id=43} + %add.1142 = f32[1,41,1]{2,1,0} add(%div.998, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/feedforward_layernorm/add" stack_frame_id=750} + %rsqrt.198 = f32[1,41,1]{2,1,0} rsqrt(%add.1142), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/feedforward_layernorm/rsqrt" stack_frame_id=753} + %mul.3174 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.198), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3175 = f32[1,41]{1,0} reshape(%mul.3174), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3176 = f32[1,41,4096]{2,1,0} broadcast(%mul.3175), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3177 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.1775, %mul.3176), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/feedforward_layernorm/mul" stack_frame_id=754} + %state_1__28_.1 = bf16[4096]{0} parameter(31), metadata={op_name="state[1][28]"} + %convert_element_type.1776 = f32[4096]{0} convert(%state_1__28_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/feedforward_layernorm/convert_element_type" stack_frame_id=755} + %broadcast_in_dim.1544 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1776), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/feedforward_layernorm/broadcast_in_dim" stack_frame_id=755} + %mul.3178 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1544), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3179 = f32[1,4096]{1,0} reshape(%mul.3178), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3180 = f32[1,41,4096]{2,1,0} broadcast(%mul.3179), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3181 = f32[1,41,4096]{2,1,0} multiply(%mul.3177, %mul.3180), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/feedforward_layernorm/mul" stack_frame_id=755} + %convert_element_type.1777 = bf16[1,41,4096]{2,1,0} convert(%mul.3181), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/feedforward_layernorm/convert_element_type" stack_frame_id=759} + %state_1__25_.1 = bf16[4096,14336]{1,0} parameter(28), metadata={op_name="state[1][25]"} + %dot_general.1081 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.1777, %state_1__25_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/feedforward_intermediate_dense/dot_general" stack_frame_id=788} + %state_1__26_.1 = bf16[4096,14336]{1,0} parameter(29), metadata={op_name="state[1][26]"} + %dot_general.1080 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.1777, %state_1__26_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/feedforward_gate_dense/dot_general" stack_frame_id=768} + %convert_element_type.1778 = f32[1,41,14336]{2,1,0} convert(%dot_general.1080), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/convert_element_type" stack_frame_id=772} + %jit_silu_.97 = f32[1,41,14336]{2,1,0} call(%convert_element_type.1778), to_apply=%silu.9, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/jit(silu)" stack_frame_id=775} + %convert_element_type.1779 = bf16[1,41,14336]{2,1,0} convert(%jit_silu_.97), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/convert_element_type" stack_frame_id=779} + %mul.3182 = bf16[1,41,14336]{2,1,0} multiply(%dot_general.1081, %convert_element_type.1779), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/mul" stack_frame_id=791} + %state_1__27_.1 = bf16[14336,4096]{1,0} parameter(30), metadata={op_name="state[1][27]"} + %dot_general.1082 = bf16[1,41,4096]{2,1,0} dot(%mul.3182, %state_1__27_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/feedforward_output_dense/dot_general" stack_frame_id=800} + %add.1143 = bf16[1,41,4096]{2,1,0} add(%dot_general.1082, %add.1141), metadata={op_name="jit(compiled_generate_function)/transformer_layer_2/add" stack_frame_id=801} + %convert_element_type.1782 = f32[1,41,4096]{2,1,0} convert(%add.1143), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention_layernorm/convert_element_type" stack_frame_id=555} + %pow.396 = f32[1,41,4096]{2,1,0} power(%convert_element_type.1782, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention_layernorm/pow" stack_frame_id=558} + %reduce_sum.1458 = f32[1,41]{1,0} reduce(%pow.396, %constant.155), dimensions={2}, to_apply=%region_13.18, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention_layernorm/reduce_sum" stack_frame_id=24} + %broadcast_in_dim.1551 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1458), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention_layernorm/broadcast_in_dim" stack_frame_id=24} + %div.999 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1551, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention_layernorm/div" stack_frame_id=24} + %add.1145 = f32[1,41,1]{2,1,0} add(%div.999, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention_layernorm/add" stack_frame_id=559} + %rsqrt.199 = f32[1,41,1]{2,1,0} rsqrt(%add.1145), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention_layernorm/rsqrt" stack_frame_id=562} + %mul.3183 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.199), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3184 = f32[1,41]{1,0} reshape(%mul.3183), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3185 = f32[1,41,4096]{2,1,0} broadcast(%mul.3184), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3186 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.1782, %mul.3185), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention_layernorm/mul" stack_frame_id=563} + %state_1__33_.1 = bf16[4096]{0} parameter(36), metadata={op_name="state[1][33]"} + %convert_element_type.1783 = f32[4096]{0} convert(%state_1__33_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention_layernorm/convert_element_type" stack_frame_id=564} + %broadcast_in_dim.1552 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1783), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention_layernorm/broadcast_in_dim" stack_frame_id=564} + %mul.3187 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1552), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3188 = f32[1,4096]{1,0} reshape(%mul.3187), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3189 = f32[1,41,4096]{2,1,0} broadcast(%mul.3188), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3190 = f32[1,41,4096]{2,1,0} multiply(%mul.3186, %mul.3189), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention_layernorm/mul" stack_frame_id=564} + %convert_element_type.1784 = bf16[1,41,4096]{2,1,0} convert(%mul.3190), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention_layernorm/convert_element_type" stack_frame_id=568} + %state_1__30_.1 = bf16[4096,8,128]{2,1,0} parameter(33), metadata={op_name="state[1][30]"} + %dot_general.1085 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.1784, %state_1__30_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=647} + %iota.541 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/iota" stack_frame_id=677} + %broadcast_in_dim.1554 = f32[1,41]{1,0} reshape(%iota.541), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/broadcast_in_dim" stack_frame_id=680} + %iota.540 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/iota" stack_frame_id=667} + %mul.3200 = f32[64]{0} multiply(%iota.540, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=667} + %div.1002 = f32[64]{0} divide(%mul.3200, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/div" stack_frame_id=668} + %neg.404 = f32[64]{0} negate(%div.1002), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/neg" stack_frame_id=669} + %pow.398 = f32[64]{0} power(%broadcast.51, %neg.404), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/pow" stack_frame_id=672} + %div.1003 = f32[64]{0} divide(%pow.398, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/div" stack_frame_id=673} + %dot_general.1087 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1554, %div.1003), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/bi,j->bij/dot_general" stack_frame_id=683} + %stack.1584 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1087), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/stack" stack_frame_id=686} + %stack.1585 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1087), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/stack" stack_frame_id=686} + %stack.1586 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1584, %stack.1585), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/stack" stack_frame_id=686} + %reshape.724 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1586), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/reshape"} + %cos.198 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.724), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/cos" stack_frame_id=690} + %convert_element_type.1787 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.198), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/convert_element_type" stack_frame_id=694} + %mul.3201 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1787), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=714} + %mul.3202 = bf16[1,41,128]{2,1,0} reshape(%mul.3201), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=714} + %mul.3203 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3202), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=714} + %mul.3204 = bf16[1,41,8,128]{3,2,1,0} multiply(%dot_general.1085, %mul.3203), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=714} + %split.397 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1085), slice={[0:1], [0:41], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/split" stack_frame_id=706} + %neg.405 = bf16[1,41,8,64]{3,2,1,0} negate(%split.397), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/neg" stack_frame_id=707} + %stack.1587 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%neg.405), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/stack" stack_frame_id=710} + %split.396 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1085), slice={[0:1], [0:41], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/split" stack_frame_id=706} + %stack.1588 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%split.396), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/stack" stack_frame_id=710} + %stack.1589 = bf16[1,41,8,2,64]{4,3,2,1,0} concatenate(%stack.1587, %stack.1588), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/stack" stack_frame_id=710} + %reshape.725 = bf16[1,41,8,128]{3,2,1,0} reshape(%stack.1589), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/reshape" stack_frame_id=713} + %sin.198 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.724), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/sin" stack_frame_id=698} + %convert_element_type.1788 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.198), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/convert_element_type" stack_frame_id=702} + %mul.3205 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1788), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=715} + %mul.3206 = bf16[1,41,128]{2,1,0} reshape(%mul.3205), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=715} + %mul.3207 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3206), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=715} + %mul.3208 = bf16[1,41,8,128]{3,2,1,0} multiply(%reshape.725, %mul.3207), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=715} + %add.1147 = bf16[1,41,8,128]{3,2,1,0} add(%mul.3204, %mul.3208), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/add" stack_frame_id=716} + %stack.1590 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%add.1147), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/stack" stack_frame_id=719} + %state_1__31_.1 = bf16[4096,8,128]{2,1,0} parameter(34), metadata={op_name="state[1][31]"} + %dot_general.1086 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.1784, %state_1__31_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=656} + %stack.1591 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dot_general.1086), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/stack" stack_frame_id=719} + %stack.1592 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1590, %stack.1591), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/stack" stack_frame_id=719} + %stack.2010 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1592), metadata={op_name="jit(compiled_generate_function)/stack" stack_frame_id=804} + %broadcast_in_dim.1556 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dot_general.1086), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/broadcast_in_dim" stack_frame_id=725} + %reshape.727 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1556), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/reshape" stack_frame_id=725} + %iota.534 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/iota" stack_frame_id=525} + %broadcast_in_dim.1545 = f32[41,1]{1,0} reshape(%iota.534), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/broadcast_in_dim" stack_frame_id=528} + %ge.502 = f32[41,1]{1,0} broadcast(%broadcast_in_dim.1545), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/ge" stack_frame_id=532} + %ge.503 = f32[41]{0} reshape(%ge.502), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/ge" stack_frame_id=532} + %ge.504 = f32[41,41]{1,0} broadcast(%ge.503), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/ge" stack_frame_id=532} + %iota.535 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/iota" stack_frame_id=531} + %broadcast_in_dim.1546 = f32[1,41]{1,0} reshape(%iota.535), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/broadcast_in_dim" stack_frame_id=532} + %ge.505 = f32[1,41]{1,0} broadcast(%broadcast_in_dim.1546), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/ge" stack_frame_id=532} + %ge.506 = f32[41]{0} reshape(%ge.505), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/ge" stack_frame_id=532} + %ge.507 = f32[41,41]{1,0} broadcast(%ge.506), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/ge" stack_frame_id=532} + %ge.508 = pred[41,41]{1,0} compare(%ge.504, %ge.507), direction=GE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/ge" stack_frame_id=532} + %broadcast_in_dim.1547 = pred[1,41,41]{2,1,0} reshape(%ge.508), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/broadcast_in_dim" stack_frame_id=535} + %convert_element_type.1781 = s32[1,41,41]{2,1,0} convert(%broadcast_in_dim.1547), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/convert_element_type" stack_frame_id=551} + %iota.536 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/iota" stack_frame_id=538} + %broadcast_in_dim.1548 = s32[41,1]{1,0} reshape(%iota.536), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/broadcast_in_dim" stack_frame_id=536} + %lt.644 = s32[41,1]{1,0} broadcast(%broadcast_in_dim.1548), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/lt" stack_frame_id=543} + %lt.645 = s32[41]{0} reshape(%lt.644), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/lt" stack_frame_id=543} + %lt.646 = s32[41,41]{1,0} broadcast(%lt.645), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/lt" stack_frame_id=543} + %iota.537 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/iota" stack_frame_id=541} + %broadcast_in_dim.1549 = s32[1,41]{1,0} reshape(%iota.537), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/broadcast_in_dim" stack_frame_id=539} + %add.1144 = s32[1,41]{1,0} add(%broadcast_in_dim.1549, %broadcast.57), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/add" stack_frame_id=542} + %lt.647 = s32[1,41]{1,0} broadcast(%add.1144), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/lt" stack_frame_id=543} + %lt.648 = s32[41]{0} reshape(%lt.647), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/lt" stack_frame_id=543} + %lt.649 = s32[41,41]{1,0} broadcast(%lt.648), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/lt" stack_frame_id=543} + %lt.650 = pred[41,41]{1,0} compare(%lt.646, %lt.649), direction=LT, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/lt" stack_frame_id=543} + %convert_element_type.1780 = s32[41,41]{1,0} convert(%lt.650), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/convert_element_type" stack_frame_id=547} + %broadcast_in_dim.1550 = s32[1,41,41]{2,1,0} reshape(%convert_element_type.1780), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/broadcast_in_dim" stack_frame_id=551} + %min.98 = s32[1,41,41]{2,1,0} minimum(%convert_element_type.1781, %broadcast_in_dim.1550), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/min" stack_frame_id=551} + %broadcast_in_dim.1557 = s32[1,1,41,41]{3,2,1,0} reshape(%min.98), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/broadcast_in_dim" stack_frame_id=728} + %convert_element_type.1789 = pred[1,1,41,41]{3,2,1,0} compare(%broadcast_in_dim.1557, %broadcast.49), direction=NE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/convert_element_type" stack_frame_id=732} + %broadcast_in_dim.1558 = pred[1,1,1,41,41]{4,3,2,1,0} reshape(%convert_element_type.1789), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/broadcast_in_dim" stack_frame_id=34} + %and.307 = pred[1,1,1,41,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1558), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/vmap()/and" stack_frame_id=34} + %and.308 = pred[1,1,41,41]{3,2,1,0} reshape(%and.307), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/vmap()/and" stack_frame_id=34} + %and.309 = pred[1,1,32,41,41]{4,3,2,1,0} broadcast(%and.308), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/vmap()/and" stack_frame_id=34} + %state_1__29_.1 = bf16[4096,32,128]{2,1,0} parameter(32), metadata={op_name="state[1][29]"} + %dot_general.1083 = bf16[1,41,32,128]{3,2,1,0} dot(%convert_element_type.1784, %state_1__29_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=577} + %iota.539 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/iota" stack_frame_id=598} + %broadcast_in_dim.1553 = f32[1,41]{1,0} reshape(%iota.539), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/broadcast_in_dim" stack_frame_id=601} + %iota.538 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/iota" stack_frame_id=588} + %mul.3191 = f32[64]{0} multiply(%iota.538, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=588} + %div.1000 = f32[64]{0} divide(%mul.3191, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/div" stack_frame_id=589} + %neg.402 = f32[64]{0} negate(%div.1000), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/neg" stack_frame_id=590} + %pow.397 = f32[64]{0} power(%broadcast.51, %neg.402), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/pow" stack_frame_id=593} + %div.1001 = f32[64]{0} divide(%pow.397, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/div" stack_frame_id=594} + %dot_general.1084 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1553, %div.1001), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/bi,j->bij/dot_general" stack_frame_id=604} + %stack.1578 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1084), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/stack" stack_frame_id=607} + %stack.1579 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1084), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/stack" stack_frame_id=607} + %stack.1580 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1578, %stack.1579), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/stack" stack_frame_id=607} + %reshape.722 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1580), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/reshape"} + %cos.197 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.722), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/cos" stack_frame_id=611} + %convert_element_type.1785 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.197), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/convert_element_type" stack_frame_id=615} + %mul.3192 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1785), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=635} + %mul.3193 = bf16[1,41,128]{2,1,0} reshape(%mul.3192), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=635} + %mul.3194 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3193), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=635} + %mul.3195 = bf16[1,41,32,128]{3,2,1,0} multiply(%dot_general.1083, %mul.3194), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=635} + %split.395 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1083), slice={[0:1], [0:41], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/split" stack_frame_id=627} + %neg.403 = bf16[1,41,32,64]{3,2,1,0} negate(%split.395), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/neg" stack_frame_id=628} + %stack.1581 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%neg.403), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/stack" stack_frame_id=631} + %split.394 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1083), slice={[0:1], [0:41], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/split" stack_frame_id=627} + %stack.1582 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%split.394), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/stack" stack_frame_id=631} + %stack.1583 = bf16[1,41,32,2,64]{4,3,2,1,0} concatenate(%stack.1581, %stack.1582), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/stack" stack_frame_id=631} + %reshape.723 = bf16[1,41,32,128]{3,2,1,0} reshape(%stack.1583), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/reshape" stack_frame_id=634} + %sin.197 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.722), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/sin" stack_frame_id=619} + %convert_element_type.1786 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.197), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/convert_element_type" stack_frame_id=623} + %mul.3196 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1786), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=636} + %mul.3197 = bf16[1,41,128]{2,1,0} reshape(%mul.3196), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=636} + %mul.3198 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3197), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=636} + %mul.3199 = bf16[1,41,32,128]{3,2,1,0} multiply(%reshape.723, %mul.3198), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/mul" stack_frame_id=636} + %add.1146 = bf16[1,41,32,128]{3,2,1,0} add(%mul.3195, %mul.3199), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/rotary_embedding_3/add" stack_frame_id=637} + %reshape.728 = bf16[1,41,32,1,128]{4,3,2,1,0} reshape(%add.1146), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/reshape" stack_frame_id=34} + %broadcast_in_dim.1555 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%add.1147), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/broadcast_in_dim" stack_frame_id=722} + %reshape.726 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1555), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/reshape" stack_frame_id=722} + %dot_general.1088 = f32[1,32,41,1,41]{4,3,2,1,0} dot(%reshape.728, %reshape.726), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=34} + %mul.3209 = f32[1,32,41,1,41]{4,3,2,1,0} multiply(%dot_general.1088, %broadcast.48), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/vmap()/mul" stack_frame_id=34} + %vmap_jit__where__.98 = f32[1,1,32,41,41]{4,3,2,1,0} call(%and.309, %mul.3209, %constant.153), to_apply=%_where_0.5, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/vmap(jit(_where))" stack_frame_id=34} + %reduce_max.484 = f32[1,1,32,41]{3,2,1,0} reduce(%vmap_jit__where__.98, %constant.152), dimensions={4}, to_apply=%region_14.19, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/vmap()/reduce_max" stack_frame_id=34} + %max.108 = f32[1,1,32,41]{3,2,1,0} maximum(%reduce_max.484, %broadcast.47), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/vmap()/max" stack_frame_id=34} + %broadcast_in_dim.1559 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%max.108), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %sub.416 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1559), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/vmap()/sub" stack_frame_id=34} + %sub.417 = f32[1,1,32,41]{3,2,1,0} reshape(%sub.416), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/vmap()/sub" stack_frame_id=34} + %sub.418 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%sub.417), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/vmap()/sub" stack_frame_id=34} + %sub.419 = f32[1,1,32,41,41]{4,3,2,1,0} subtract(%vmap_jit__where__.98, %sub.418), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/vmap()/sub" stack_frame_id=34} + %exp.104 = f32[1,1,32,41,41]{4,3,2,1,0} exponential(%sub.419), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/vmap()/exp" stack_frame_id=34} + %reduce_sum.1459 = f32[1,1,32,41]{3,2,1,0} reduce(%exp.104, %constant.155), dimensions={4}, to_apply=%region_15.20, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/vmap()/reduce_sum" stack_frame_id=34} + %broadcast_in_dim.1560 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%reduce_sum.1459), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %div.1004 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1560), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/vmap()/div" stack_frame_id=34} + %div.1005 = f32[1,1,32,41]{3,2,1,0} reshape(%div.1004), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/vmap()/div" stack_frame_id=34} + %div.1006 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%div.1005), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/vmap()/div" stack_frame_id=34} + %div.1007 = f32[1,1,32,41,41]{4,3,2,1,0} divide(%exp.104, %div.1006), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/vmap()/div" stack_frame_id=34} + %convert_element_type.1790 = bf16[1,1,32,41,41]{4,3,2,1,0} convert(%div.1007), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/vmap()/convert_element_type" stack_frame_id=34} + %dot_general.1089 = bf16[1,32,128,1,41]{4,3,2,1,0} dot(%reshape.727, %convert_element_type.1790), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=34} + %transpose.102 = bf16[1,41,32,1,128]{1,3,4,2,0} transpose(%dot_general.1089), dimensions={0,4,1,3,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/vmap()/transpose;jit(compiled_generate_function)/transformer_layer_3/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %reshape.729 = bf16[1,41,32,128]{3,2,1,0} reshape(%transpose.102), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/reshape" stack_frame_id=34} + %state_1__32_.1 = bf16[32,128,4096]{2,1,0} parameter(35), metadata={op_name="state[1][32]"} + %dot_general.1090 = bf16[1,41,4096]{2,1,0} dot(%reshape.729, %state_1__32_.1), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=741} + %add.1148 = bf16[1,41,4096]{2,1,0} add(%dot_general.1090, %add.1143), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/add" stack_frame_id=742} + %convert_element_type.1791 = f32[1,41,4096]{2,1,0} convert(%add.1148), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/feedforward_layernorm/convert_element_type" stack_frame_id=746} + %pow.399 = f32[1,41,4096]{2,1,0} power(%convert_element_type.1791, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/feedforward_layernorm/pow" stack_frame_id=749} + %reduce_sum.1460 = f32[1,41]{1,0} reduce(%pow.399, %constant.155), dimensions={2}, to_apply=%region_16.21, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/feedforward_layernorm/reduce_sum" stack_frame_id=43} + %broadcast_in_dim.1561 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1460), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/feedforward_layernorm/broadcast_in_dim" stack_frame_id=43} + %div.1008 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1561, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/feedforward_layernorm/div" stack_frame_id=43} + %add.1149 = f32[1,41,1]{2,1,0} add(%div.1008, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/feedforward_layernorm/add" stack_frame_id=750} + %rsqrt.200 = f32[1,41,1]{2,1,0} rsqrt(%add.1149), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/feedforward_layernorm/rsqrt" stack_frame_id=753} + %mul.3210 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.200), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3211 = f32[1,41]{1,0} reshape(%mul.3210), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3212 = f32[1,41,4096]{2,1,0} broadcast(%mul.3211), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3213 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.1791, %mul.3212), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/feedforward_layernorm/mul" stack_frame_id=754} + %state_1__37_.1 = bf16[4096]{0} parameter(40), metadata={op_name="state[1][37]"} + %convert_element_type.1792 = f32[4096]{0} convert(%state_1__37_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/feedforward_layernorm/convert_element_type" stack_frame_id=755} + %broadcast_in_dim.1562 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1792), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/feedforward_layernorm/broadcast_in_dim" stack_frame_id=755} + %mul.3214 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1562), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3215 = f32[1,4096]{1,0} reshape(%mul.3214), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3216 = f32[1,41,4096]{2,1,0} broadcast(%mul.3215), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3217 = f32[1,41,4096]{2,1,0} multiply(%mul.3213, %mul.3216), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/feedforward_layernorm/mul" stack_frame_id=755} + %convert_element_type.1793 = bf16[1,41,4096]{2,1,0} convert(%mul.3217), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/feedforward_layernorm/convert_element_type" stack_frame_id=759} + %state_1__34_.1 = bf16[4096,14336]{1,0} parameter(37), metadata={op_name="state[1][34]"} + %dot_general.1092 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.1793, %state_1__34_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/feedforward_intermediate_dense/dot_general" stack_frame_id=788} + %state_1__35_.1 = bf16[4096,14336]{1,0} parameter(38), metadata={op_name="state[1][35]"} + %dot_general.1091 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.1793, %state_1__35_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/feedforward_gate_dense/dot_general" stack_frame_id=768} + %convert_element_type.1794 = f32[1,41,14336]{2,1,0} convert(%dot_general.1091), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/convert_element_type" stack_frame_id=772} + %jit_silu_.98 = f32[1,41,14336]{2,1,0} call(%convert_element_type.1794), to_apply=%silu.9, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/jit(silu)" stack_frame_id=775} + %convert_element_type.1795 = bf16[1,41,14336]{2,1,0} convert(%jit_silu_.98), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/convert_element_type" stack_frame_id=779} + %mul.3218 = bf16[1,41,14336]{2,1,0} multiply(%dot_general.1092, %convert_element_type.1795), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/mul" stack_frame_id=791} + %state_1__36_.1 = bf16[14336,4096]{1,0} parameter(39), metadata={op_name="state[1][36]"} + %dot_general.1093 = bf16[1,41,4096]{2,1,0} dot(%mul.3218, %state_1__36_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/feedforward_output_dense/dot_general" stack_frame_id=800} + %add.1150 = bf16[1,41,4096]{2,1,0} add(%dot_general.1093, %add.1148), metadata={op_name="jit(compiled_generate_function)/transformer_layer_3/add" stack_frame_id=801} + %convert_element_type.1798 = f32[1,41,4096]{2,1,0} convert(%add.1150), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention_layernorm/convert_element_type" stack_frame_id=555} + %pow.400 = f32[1,41,4096]{2,1,0} power(%convert_element_type.1798, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention_layernorm/pow" stack_frame_id=558} + %reduce_sum.1461 = f32[1,41]{1,0} reduce(%pow.400, %constant.155), dimensions={2}, to_apply=%region_17.22, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention_layernorm/reduce_sum" stack_frame_id=24} + %broadcast_in_dim.1569 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1461), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention_layernorm/broadcast_in_dim" stack_frame_id=24} + %div.1009 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1569, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention_layernorm/div" stack_frame_id=24} + %add.1152 = f32[1,41,1]{2,1,0} add(%div.1009, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention_layernorm/add" stack_frame_id=559} + %rsqrt.201 = f32[1,41,1]{2,1,0} rsqrt(%add.1152), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention_layernorm/rsqrt" stack_frame_id=562} + %mul.3219 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.201), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3220 = f32[1,41]{1,0} reshape(%mul.3219), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3221 = f32[1,41,4096]{2,1,0} broadcast(%mul.3220), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3222 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.1798, %mul.3221), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention_layernorm/mul" stack_frame_id=563} + %state_1__42_.1 = bf16[4096]{0} parameter(45), metadata={op_name="state[1][42]"} + %convert_element_type.1799 = f32[4096]{0} convert(%state_1__42_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention_layernorm/convert_element_type" stack_frame_id=564} + %broadcast_in_dim.1570 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1799), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention_layernorm/broadcast_in_dim" stack_frame_id=564} + %mul.3223 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1570), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3224 = f32[1,4096]{1,0} reshape(%mul.3223), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3225 = f32[1,41,4096]{2,1,0} broadcast(%mul.3224), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3226 = f32[1,41,4096]{2,1,0} multiply(%mul.3222, %mul.3225), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention_layernorm/mul" stack_frame_id=564} + %convert_element_type.1800 = bf16[1,41,4096]{2,1,0} convert(%mul.3226), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention_layernorm/convert_element_type" stack_frame_id=568} + %state_1__39_.1 = bf16[4096,8,128]{2,1,0} parameter(42), metadata={op_name="state[1][39]"} + %dot_general.1096 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.1800, %state_1__39_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=647} + %iota.549 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/iota" stack_frame_id=677} + %broadcast_in_dim.1572 = f32[1,41]{1,0} reshape(%iota.549), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/broadcast_in_dim" stack_frame_id=680} + %iota.548 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/iota" stack_frame_id=667} + %mul.3236 = f32[64]{0} multiply(%iota.548, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=667} + %div.1012 = f32[64]{0} divide(%mul.3236, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/div" stack_frame_id=668} + %neg.408 = f32[64]{0} negate(%div.1012), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/neg" stack_frame_id=669} + %pow.402 = f32[64]{0} power(%broadcast.51, %neg.408), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/pow" stack_frame_id=672} + %div.1013 = f32[64]{0} divide(%pow.402, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/div" stack_frame_id=673} + %dot_general.1098 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1572, %div.1013), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/bi,j->bij/dot_general" stack_frame_id=683} + %stack.1599 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1098), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/stack" stack_frame_id=686} + %stack.1600 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1098), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/stack" stack_frame_id=686} + %stack.1601 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1599, %stack.1600), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/stack" stack_frame_id=686} + %reshape.732 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1601), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/reshape"} + %cos.200 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.732), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/cos" stack_frame_id=690} + %convert_element_type.1803 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.200), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/convert_element_type" stack_frame_id=694} + %mul.3237 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1803), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=714} + %mul.3238 = bf16[1,41,128]{2,1,0} reshape(%mul.3237), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=714} + %mul.3239 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3238), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=714} + %mul.3240 = bf16[1,41,8,128]{3,2,1,0} multiply(%dot_general.1096, %mul.3239), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=714} + %split.401 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1096), slice={[0:1], [0:41], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/split" stack_frame_id=706} + %neg.409 = bf16[1,41,8,64]{3,2,1,0} negate(%split.401), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/neg" stack_frame_id=707} + %stack.1602 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%neg.409), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/stack" stack_frame_id=710} + %split.400 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1096), slice={[0:1], [0:41], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/split" stack_frame_id=706} + %stack.1603 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%split.400), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/stack" stack_frame_id=710} + %stack.1604 = bf16[1,41,8,2,64]{4,3,2,1,0} concatenate(%stack.1602, %stack.1603), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/stack" stack_frame_id=710} + %reshape.733 = bf16[1,41,8,128]{3,2,1,0} reshape(%stack.1604), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/reshape" stack_frame_id=713} + %sin.200 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.732), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/sin" stack_frame_id=698} + %convert_element_type.1804 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.200), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/convert_element_type" stack_frame_id=702} + %mul.3241 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1804), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=715} + %mul.3242 = bf16[1,41,128]{2,1,0} reshape(%mul.3241), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=715} + %mul.3243 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3242), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=715} + %mul.3244 = bf16[1,41,8,128]{3,2,1,0} multiply(%reshape.733, %mul.3243), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=715} + %add.1154 = bf16[1,41,8,128]{3,2,1,0} add(%mul.3240, %mul.3244), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/add" stack_frame_id=716} + %stack.1605 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%add.1154), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/stack" stack_frame_id=719} + %state_1__40_.1 = bf16[4096,8,128]{2,1,0} parameter(43), metadata={op_name="state[1][40]"} + %dot_general.1097 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.1800, %state_1__40_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=656} + %stack.1606 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dot_general.1097), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/stack" stack_frame_id=719} + %stack.1607 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1605, %stack.1606), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/stack" stack_frame_id=719} + %stack.2011 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1607), metadata={op_name="jit(compiled_generate_function)/stack" stack_frame_id=804} + %broadcast_in_dim.1574 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dot_general.1097), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/broadcast_in_dim" stack_frame_id=725} + %reshape.735 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1574), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/reshape" stack_frame_id=725} + %iota.542 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/iota" stack_frame_id=525} + %broadcast_in_dim.1563 = f32[41,1]{1,0} reshape(%iota.542), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/broadcast_in_dim" stack_frame_id=528} + %ge.509 = f32[41,1]{1,0} broadcast(%broadcast_in_dim.1563), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/ge" stack_frame_id=532} + %ge.510 = f32[41]{0} reshape(%ge.509), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/ge" stack_frame_id=532} + %ge.511 = f32[41,41]{1,0} broadcast(%ge.510), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/ge" stack_frame_id=532} + %iota.543 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/iota" stack_frame_id=531} + %broadcast_in_dim.1564 = f32[1,41]{1,0} reshape(%iota.543), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/broadcast_in_dim" stack_frame_id=532} + %ge.512 = f32[1,41]{1,0} broadcast(%broadcast_in_dim.1564), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/ge" stack_frame_id=532} + %ge.513 = f32[41]{0} reshape(%ge.512), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/ge" stack_frame_id=532} + %ge.514 = f32[41,41]{1,0} broadcast(%ge.513), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/ge" stack_frame_id=532} + %ge.515 = pred[41,41]{1,0} compare(%ge.511, %ge.514), direction=GE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/ge" stack_frame_id=532} + %broadcast_in_dim.1565 = pred[1,41,41]{2,1,0} reshape(%ge.515), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/broadcast_in_dim" stack_frame_id=535} + %convert_element_type.1797 = s32[1,41,41]{2,1,0} convert(%broadcast_in_dim.1565), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/convert_element_type" stack_frame_id=551} + %iota.544 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/iota" stack_frame_id=538} + %broadcast_in_dim.1566 = s32[41,1]{1,0} reshape(%iota.544), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/broadcast_in_dim" stack_frame_id=536} + %lt.651 = s32[41,1]{1,0} broadcast(%broadcast_in_dim.1566), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/lt" stack_frame_id=543} + %lt.652 = s32[41]{0} reshape(%lt.651), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/lt" stack_frame_id=543} + %lt.653 = s32[41,41]{1,0} broadcast(%lt.652), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/lt" stack_frame_id=543} + %iota.545 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/iota" stack_frame_id=541} + %broadcast_in_dim.1567 = s32[1,41]{1,0} reshape(%iota.545), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/broadcast_in_dim" stack_frame_id=539} + %add.1151 = s32[1,41]{1,0} add(%broadcast_in_dim.1567, %broadcast.57), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/add" stack_frame_id=542} + %lt.654 = s32[1,41]{1,0} broadcast(%add.1151), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/lt" stack_frame_id=543} + %lt.655 = s32[41]{0} reshape(%lt.654), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/lt" stack_frame_id=543} + %lt.656 = s32[41,41]{1,0} broadcast(%lt.655), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/lt" stack_frame_id=543} + %lt.657 = pred[41,41]{1,0} compare(%lt.653, %lt.656), direction=LT, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/lt" stack_frame_id=543} + %convert_element_type.1796 = s32[41,41]{1,0} convert(%lt.657), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/convert_element_type" stack_frame_id=547} + %broadcast_in_dim.1568 = s32[1,41,41]{2,1,0} reshape(%convert_element_type.1796), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/broadcast_in_dim" stack_frame_id=551} + %min.99 = s32[1,41,41]{2,1,0} minimum(%convert_element_type.1797, %broadcast_in_dim.1568), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/min" stack_frame_id=551} + %broadcast_in_dim.1575 = s32[1,1,41,41]{3,2,1,0} reshape(%min.99), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/broadcast_in_dim" stack_frame_id=728} + %convert_element_type.1805 = pred[1,1,41,41]{3,2,1,0} compare(%broadcast_in_dim.1575, %broadcast.49), direction=NE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/convert_element_type" stack_frame_id=732} + %broadcast_in_dim.1576 = pred[1,1,1,41,41]{4,3,2,1,0} reshape(%convert_element_type.1805), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/broadcast_in_dim" stack_frame_id=34} + %and.310 = pred[1,1,1,41,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1576), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/vmap()/and" stack_frame_id=34} + %and.311 = pred[1,1,41,41]{3,2,1,0} reshape(%and.310), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/vmap()/and" stack_frame_id=34} + %and.312 = pred[1,1,32,41,41]{4,3,2,1,0} broadcast(%and.311), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/vmap()/and" stack_frame_id=34} + %state_1__38_.1 = bf16[4096,32,128]{2,1,0} parameter(41), metadata={op_name="state[1][38]"} + %dot_general.1094 = bf16[1,41,32,128]{3,2,1,0} dot(%convert_element_type.1800, %state_1__38_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=577} + %iota.547 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/iota" stack_frame_id=598} + %broadcast_in_dim.1571 = f32[1,41]{1,0} reshape(%iota.547), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/broadcast_in_dim" stack_frame_id=601} + %iota.546 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/iota" stack_frame_id=588} + %mul.3227 = f32[64]{0} multiply(%iota.546, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=588} + %div.1010 = f32[64]{0} divide(%mul.3227, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/div" stack_frame_id=589} + %neg.406 = f32[64]{0} negate(%div.1010), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/neg" stack_frame_id=590} + %pow.401 = f32[64]{0} power(%broadcast.51, %neg.406), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/pow" stack_frame_id=593} + %div.1011 = f32[64]{0} divide(%pow.401, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/div" stack_frame_id=594} + %dot_general.1095 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1571, %div.1011), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/bi,j->bij/dot_general" stack_frame_id=604} + %stack.1593 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1095), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/stack" stack_frame_id=607} + %stack.1594 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1095), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/stack" stack_frame_id=607} + %stack.1595 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1593, %stack.1594), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/stack" stack_frame_id=607} + %reshape.730 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1595), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/reshape"} + %cos.199 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.730), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/cos" stack_frame_id=611} + %convert_element_type.1801 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.199), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/convert_element_type" stack_frame_id=615} + %mul.3228 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1801), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=635} + %mul.3229 = bf16[1,41,128]{2,1,0} reshape(%mul.3228), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=635} + %mul.3230 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3229), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=635} + %mul.3231 = bf16[1,41,32,128]{3,2,1,0} multiply(%dot_general.1094, %mul.3230), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=635} + %split.399 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1094), slice={[0:1], [0:41], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/split" stack_frame_id=627} + %neg.407 = bf16[1,41,32,64]{3,2,1,0} negate(%split.399), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/neg" stack_frame_id=628} + %stack.1596 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%neg.407), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/stack" stack_frame_id=631} + %split.398 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1094), slice={[0:1], [0:41], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/split" stack_frame_id=627} + %stack.1597 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%split.398), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/stack" stack_frame_id=631} + %stack.1598 = bf16[1,41,32,2,64]{4,3,2,1,0} concatenate(%stack.1596, %stack.1597), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/stack" stack_frame_id=631} + %reshape.731 = bf16[1,41,32,128]{3,2,1,0} reshape(%stack.1598), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/reshape" stack_frame_id=634} + %sin.199 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.730), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/sin" stack_frame_id=619} + %convert_element_type.1802 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.199), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/convert_element_type" stack_frame_id=623} + %mul.3232 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1802), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=636} + %mul.3233 = bf16[1,41,128]{2,1,0} reshape(%mul.3232), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=636} + %mul.3234 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3233), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=636} + %mul.3235 = bf16[1,41,32,128]{3,2,1,0} multiply(%reshape.731, %mul.3234), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/mul" stack_frame_id=636} + %add.1153 = bf16[1,41,32,128]{3,2,1,0} add(%mul.3231, %mul.3235), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/rotary_embedding_4/add" stack_frame_id=637} + %reshape.736 = bf16[1,41,32,1,128]{4,3,2,1,0} reshape(%add.1153), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/reshape" stack_frame_id=34} + %broadcast_in_dim.1573 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%add.1154), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/broadcast_in_dim" stack_frame_id=722} + %reshape.734 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1573), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/reshape" stack_frame_id=722} + %dot_general.1099 = f32[1,32,41,1,41]{4,3,2,1,0} dot(%reshape.736, %reshape.734), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=34} + %mul.3245 = f32[1,32,41,1,41]{4,3,2,1,0} multiply(%dot_general.1099, %broadcast.48), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/vmap()/mul" stack_frame_id=34} + %vmap_jit__where__.99 = f32[1,1,32,41,41]{4,3,2,1,0} call(%and.312, %mul.3245, %constant.153), to_apply=%_where_0.5, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/vmap(jit(_where))" stack_frame_id=34} + %reduce_max.485 = f32[1,1,32,41]{3,2,1,0} reduce(%vmap_jit__where__.99, %constant.152), dimensions={4}, to_apply=%region_18.23, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/vmap()/reduce_max" stack_frame_id=34} + %max.109 = f32[1,1,32,41]{3,2,1,0} maximum(%reduce_max.485, %broadcast.47), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/vmap()/max" stack_frame_id=34} + %broadcast_in_dim.1577 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%max.109), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %sub.420 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1577), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/vmap()/sub" stack_frame_id=34} + %sub.421 = f32[1,1,32,41]{3,2,1,0} reshape(%sub.420), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/vmap()/sub" stack_frame_id=34} + %sub.422 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%sub.421), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/vmap()/sub" stack_frame_id=34} + %sub.423 = f32[1,1,32,41,41]{4,3,2,1,0} subtract(%vmap_jit__where__.99, %sub.422), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/vmap()/sub" stack_frame_id=34} + %exp.105 = f32[1,1,32,41,41]{4,3,2,1,0} exponential(%sub.423), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/vmap()/exp" stack_frame_id=34} + %reduce_sum.1462 = f32[1,1,32,41]{3,2,1,0} reduce(%exp.105, %constant.155), dimensions={4}, to_apply=%region_19.24, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/vmap()/reduce_sum" stack_frame_id=34} + %broadcast_in_dim.1578 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%reduce_sum.1462), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %div.1014 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1578), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/vmap()/div" stack_frame_id=34} + %div.1015 = f32[1,1,32,41]{3,2,1,0} reshape(%div.1014), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/vmap()/div" stack_frame_id=34} + %div.1016 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%div.1015), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/vmap()/div" stack_frame_id=34} + %div.1017 = f32[1,1,32,41,41]{4,3,2,1,0} divide(%exp.105, %div.1016), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/vmap()/div" stack_frame_id=34} + %convert_element_type.1806 = bf16[1,1,32,41,41]{4,3,2,1,0} convert(%div.1017), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/vmap()/convert_element_type" stack_frame_id=34} + %dot_general.1100 = bf16[1,32,128,1,41]{4,3,2,1,0} dot(%reshape.735, %convert_element_type.1806), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=34} + %transpose.103 = bf16[1,41,32,1,128]{1,3,4,2,0} transpose(%dot_general.1100), dimensions={0,4,1,3,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/vmap()/transpose;jit(compiled_generate_function)/transformer_layer_4/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %reshape.737 = bf16[1,41,32,128]{3,2,1,0} reshape(%transpose.103), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/reshape" stack_frame_id=34} + %state_1__41_.1 = bf16[32,128,4096]{2,1,0} parameter(44), metadata={op_name="state[1][41]"} + %dot_general.1101 = bf16[1,41,4096]{2,1,0} dot(%reshape.737, %state_1__41_.1), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=741} + %add.1155 = bf16[1,41,4096]{2,1,0} add(%dot_general.1101, %add.1150), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/add" stack_frame_id=742} + %convert_element_type.1807 = f32[1,41,4096]{2,1,0} convert(%add.1155), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/feedforward_layernorm/convert_element_type" stack_frame_id=746} + %pow.403 = f32[1,41,4096]{2,1,0} power(%convert_element_type.1807, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/feedforward_layernorm/pow" stack_frame_id=749} + %reduce_sum.1463 = f32[1,41]{1,0} reduce(%pow.403, %constant.155), dimensions={2}, to_apply=%region_20.25, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/feedforward_layernorm/reduce_sum" stack_frame_id=43} + %broadcast_in_dim.1579 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1463), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/feedforward_layernorm/broadcast_in_dim" stack_frame_id=43} + %div.1018 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1579, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/feedforward_layernorm/div" stack_frame_id=43} + %add.1156 = f32[1,41,1]{2,1,0} add(%div.1018, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/feedforward_layernorm/add" stack_frame_id=750} + %rsqrt.202 = f32[1,41,1]{2,1,0} rsqrt(%add.1156), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/feedforward_layernorm/rsqrt" stack_frame_id=753} + %mul.3246 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.202), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3247 = f32[1,41]{1,0} reshape(%mul.3246), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3248 = f32[1,41,4096]{2,1,0} broadcast(%mul.3247), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3249 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.1807, %mul.3248), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/feedforward_layernorm/mul" stack_frame_id=754} + %state_1__46_.1 = bf16[4096]{0} parameter(49), metadata={op_name="state[1][46]"} + %convert_element_type.1808 = f32[4096]{0} convert(%state_1__46_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/feedforward_layernorm/convert_element_type" stack_frame_id=755} + %broadcast_in_dim.1580 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1808), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/feedforward_layernorm/broadcast_in_dim" stack_frame_id=755} + %mul.3250 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1580), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3251 = f32[1,4096]{1,0} reshape(%mul.3250), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3252 = f32[1,41,4096]{2,1,0} broadcast(%mul.3251), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3253 = f32[1,41,4096]{2,1,0} multiply(%mul.3249, %mul.3252), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/feedforward_layernorm/mul" stack_frame_id=755} + %convert_element_type.1809 = bf16[1,41,4096]{2,1,0} convert(%mul.3253), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/feedforward_layernorm/convert_element_type" stack_frame_id=759} + %state_1__43_.1 = bf16[4096,14336]{1,0} parameter(46), metadata={op_name="state[1][43]"} + %dot_general.1103 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.1809, %state_1__43_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/feedforward_intermediate_dense/dot_general" stack_frame_id=788} + %state_1__44_.1 = bf16[4096,14336]{1,0} parameter(47), metadata={op_name="state[1][44]"} + %dot_general.1102 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.1809, %state_1__44_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/feedforward_gate_dense/dot_general" stack_frame_id=768} + %convert_element_type.1810 = f32[1,41,14336]{2,1,0} convert(%dot_general.1102), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/convert_element_type" stack_frame_id=772} + %jit_silu_.99 = f32[1,41,14336]{2,1,0} call(%convert_element_type.1810), to_apply=%silu.9, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/jit(silu)" stack_frame_id=775} + %convert_element_type.1811 = bf16[1,41,14336]{2,1,0} convert(%jit_silu_.99), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/convert_element_type" stack_frame_id=779} + %mul.3254 = bf16[1,41,14336]{2,1,0} multiply(%dot_general.1103, %convert_element_type.1811), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/mul" stack_frame_id=791} + %state_1__45_.1 = bf16[14336,4096]{1,0} parameter(48), metadata={op_name="state[1][45]"} + %dot_general.1104 = bf16[1,41,4096]{2,1,0} dot(%mul.3254, %state_1__45_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/feedforward_output_dense/dot_general" stack_frame_id=800} + %add.1157 = bf16[1,41,4096]{2,1,0} add(%dot_general.1104, %add.1155), metadata={op_name="jit(compiled_generate_function)/transformer_layer_4/add" stack_frame_id=801} + %convert_element_type.1814 = f32[1,41,4096]{2,1,0} convert(%add.1157), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention_layernorm/convert_element_type" stack_frame_id=555} + %pow.404 = f32[1,41,4096]{2,1,0} power(%convert_element_type.1814, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention_layernorm/pow" stack_frame_id=558} + %reduce_sum.1464 = f32[1,41]{1,0} reduce(%pow.404, %constant.155), dimensions={2}, to_apply=%region_21.26, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention_layernorm/reduce_sum" stack_frame_id=24} + %broadcast_in_dim.1587 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1464), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention_layernorm/broadcast_in_dim" stack_frame_id=24} + %div.1019 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1587, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention_layernorm/div" stack_frame_id=24} + %add.1159 = f32[1,41,1]{2,1,0} add(%div.1019, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention_layernorm/add" stack_frame_id=559} + %rsqrt.203 = f32[1,41,1]{2,1,0} rsqrt(%add.1159), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention_layernorm/rsqrt" stack_frame_id=562} + %mul.3255 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.203), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3256 = f32[1,41]{1,0} reshape(%mul.3255), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3257 = f32[1,41,4096]{2,1,0} broadcast(%mul.3256), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3258 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.1814, %mul.3257), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention_layernorm/mul" stack_frame_id=563} + %state_1__51_.1 = bf16[4096]{0} parameter(54), metadata={op_name="state[1][51]"} + %convert_element_type.1815 = f32[4096]{0} convert(%state_1__51_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention_layernorm/convert_element_type" stack_frame_id=564} + %broadcast_in_dim.1588 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1815), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention_layernorm/broadcast_in_dim" stack_frame_id=564} + %mul.3259 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1588), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3260 = f32[1,4096]{1,0} reshape(%mul.3259), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3261 = f32[1,41,4096]{2,1,0} broadcast(%mul.3260), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3262 = f32[1,41,4096]{2,1,0} multiply(%mul.3258, %mul.3261), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention_layernorm/mul" stack_frame_id=564} + %convert_element_type.1816 = bf16[1,41,4096]{2,1,0} convert(%mul.3262), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention_layernorm/convert_element_type" stack_frame_id=568} + %state_1__48_.1 = bf16[4096,8,128]{2,1,0} parameter(51), metadata={op_name="state[1][48]"} + %dot_general.1107 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.1816, %state_1__48_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=647} + %iota.557 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/iota" stack_frame_id=677} + %broadcast_in_dim.1590 = f32[1,41]{1,0} reshape(%iota.557), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/broadcast_in_dim" stack_frame_id=680} + %iota.556 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/iota" stack_frame_id=667} + %mul.3272 = f32[64]{0} multiply(%iota.556, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=667} + %div.1022 = f32[64]{0} divide(%mul.3272, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/div" stack_frame_id=668} + %neg.412 = f32[64]{0} negate(%div.1022), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/neg" stack_frame_id=669} + %pow.406 = f32[64]{0} power(%broadcast.51, %neg.412), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/pow" stack_frame_id=672} + %div.1023 = f32[64]{0} divide(%pow.406, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/div" stack_frame_id=673} + %dot_general.1109 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1590, %div.1023), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/bi,j->bij/dot_general" stack_frame_id=683} + %stack.1614 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1109), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/stack" stack_frame_id=686} + %stack.1615 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1109), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/stack" stack_frame_id=686} + %stack.1616 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1614, %stack.1615), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/stack" stack_frame_id=686} + %reshape.740 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1616), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/reshape"} + %cos.202 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.740), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/cos" stack_frame_id=690} + %convert_element_type.1819 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.202), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/convert_element_type" stack_frame_id=694} + %mul.3273 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1819), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=714} + %mul.3274 = bf16[1,41,128]{2,1,0} reshape(%mul.3273), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=714} + %mul.3275 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3274), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=714} + %mul.3276 = bf16[1,41,8,128]{3,2,1,0} multiply(%dot_general.1107, %mul.3275), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=714} + %split.405 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1107), slice={[0:1], [0:41], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/split" stack_frame_id=706} + %neg.413 = bf16[1,41,8,64]{3,2,1,0} negate(%split.405), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/neg" stack_frame_id=707} + %stack.1617 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%neg.413), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/stack" stack_frame_id=710} + %split.404 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1107), slice={[0:1], [0:41], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/split" stack_frame_id=706} + %stack.1618 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%split.404), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/stack" stack_frame_id=710} + %stack.1619 = bf16[1,41,8,2,64]{4,3,2,1,0} concatenate(%stack.1617, %stack.1618), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/stack" stack_frame_id=710} + %reshape.741 = bf16[1,41,8,128]{3,2,1,0} reshape(%stack.1619), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/reshape" stack_frame_id=713} + %sin.202 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.740), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/sin" stack_frame_id=698} + %convert_element_type.1820 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.202), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/convert_element_type" stack_frame_id=702} + %mul.3277 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1820), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=715} + %mul.3278 = bf16[1,41,128]{2,1,0} reshape(%mul.3277), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=715} + %mul.3279 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3278), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=715} + %mul.3280 = bf16[1,41,8,128]{3,2,1,0} multiply(%reshape.741, %mul.3279), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=715} + %add.1161 = bf16[1,41,8,128]{3,2,1,0} add(%mul.3276, %mul.3280), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/add" stack_frame_id=716} + %stack.1620 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%add.1161), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/stack" stack_frame_id=719} + %state_1__49_.1 = bf16[4096,8,128]{2,1,0} parameter(52), metadata={op_name="state[1][49]"} + %dot_general.1108 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.1816, %state_1__49_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=656} + %stack.1621 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dot_general.1108), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/stack" stack_frame_id=719} + %stack.1622 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1620, %stack.1621), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/stack" stack_frame_id=719} + %stack.2012 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1622), metadata={op_name="jit(compiled_generate_function)/stack" stack_frame_id=804} + %broadcast_in_dim.1592 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dot_general.1108), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/broadcast_in_dim" stack_frame_id=725} + %reshape.743 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1592), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/reshape" stack_frame_id=725} + %iota.550 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/iota" stack_frame_id=525} + %broadcast_in_dim.1581 = f32[41,1]{1,0} reshape(%iota.550), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/broadcast_in_dim" stack_frame_id=528} + %ge.516 = f32[41,1]{1,0} broadcast(%broadcast_in_dim.1581), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/ge" stack_frame_id=532} + %ge.517 = f32[41]{0} reshape(%ge.516), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/ge" stack_frame_id=532} + %ge.518 = f32[41,41]{1,0} broadcast(%ge.517), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/ge" stack_frame_id=532} + %iota.551 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/iota" stack_frame_id=531} + %broadcast_in_dim.1582 = f32[1,41]{1,0} reshape(%iota.551), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/broadcast_in_dim" stack_frame_id=532} + %ge.519 = f32[1,41]{1,0} broadcast(%broadcast_in_dim.1582), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/ge" stack_frame_id=532} + %ge.520 = f32[41]{0} reshape(%ge.519), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/ge" stack_frame_id=532} + %ge.521 = f32[41,41]{1,0} broadcast(%ge.520), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/ge" stack_frame_id=532} + %ge.522 = pred[41,41]{1,0} compare(%ge.518, %ge.521), direction=GE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/ge" stack_frame_id=532} + %broadcast_in_dim.1583 = pred[1,41,41]{2,1,0} reshape(%ge.522), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/broadcast_in_dim" stack_frame_id=535} + %convert_element_type.1813 = s32[1,41,41]{2,1,0} convert(%broadcast_in_dim.1583), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/convert_element_type" stack_frame_id=551} + %iota.552 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/iota" stack_frame_id=538} + %broadcast_in_dim.1584 = s32[41,1]{1,0} reshape(%iota.552), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/broadcast_in_dim" stack_frame_id=536} + %lt.658 = s32[41,1]{1,0} broadcast(%broadcast_in_dim.1584), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/lt" stack_frame_id=543} + %lt.659 = s32[41]{0} reshape(%lt.658), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/lt" stack_frame_id=543} + %lt.660 = s32[41,41]{1,0} broadcast(%lt.659), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/lt" stack_frame_id=543} + %iota.553 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/iota" stack_frame_id=541} + %broadcast_in_dim.1585 = s32[1,41]{1,0} reshape(%iota.553), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/broadcast_in_dim" stack_frame_id=539} + %add.1158 = s32[1,41]{1,0} add(%broadcast_in_dim.1585, %broadcast.57), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/add" stack_frame_id=542} + %lt.661 = s32[1,41]{1,0} broadcast(%add.1158), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/lt" stack_frame_id=543} + %lt.662 = s32[41]{0} reshape(%lt.661), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/lt" stack_frame_id=543} + %lt.663 = s32[41,41]{1,0} broadcast(%lt.662), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/lt" stack_frame_id=543} + %lt.664 = pred[41,41]{1,0} compare(%lt.660, %lt.663), direction=LT, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/lt" stack_frame_id=543} + %convert_element_type.1812 = s32[41,41]{1,0} convert(%lt.664), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/convert_element_type" stack_frame_id=547} + %broadcast_in_dim.1586 = s32[1,41,41]{2,1,0} reshape(%convert_element_type.1812), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/broadcast_in_dim" stack_frame_id=551} + %min.100 = s32[1,41,41]{2,1,0} minimum(%convert_element_type.1813, %broadcast_in_dim.1586), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/min" stack_frame_id=551} + %broadcast_in_dim.1593 = s32[1,1,41,41]{3,2,1,0} reshape(%min.100), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/broadcast_in_dim" stack_frame_id=728} + %convert_element_type.1821 = pred[1,1,41,41]{3,2,1,0} compare(%broadcast_in_dim.1593, %broadcast.49), direction=NE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/convert_element_type" stack_frame_id=732} + %broadcast_in_dim.1594 = pred[1,1,1,41,41]{4,3,2,1,0} reshape(%convert_element_type.1821), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/broadcast_in_dim" stack_frame_id=34} + %and.313 = pred[1,1,1,41,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1594), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/vmap()/and" stack_frame_id=34} + %and.314 = pred[1,1,41,41]{3,2,1,0} reshape(%and.313), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/vmap()/and" stack_frame_id=34} + %and.315 = pred[1,1,32,41,41]{4,3,2,1,0} broadcast(%and.314), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/vmap()/and" stack_frame_id=34} + %state_1__47_.1 = bf16[4096,32,128]{2,1,0} parameter(50), metadata={op_name="state[1][47]"} + %dot_general.1105 = bf16[1,41,32,128]{3,2,1,0} dot(%convert_element_type.1816, %state_1__47_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=577} + %iota.555 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/iota" stack_frame_id=598} + %broadcast_in_dim.1589 = f32[1,41]{1,0} reshape(%iota.555), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/broadcast_in_dim" stack_frame_id=601} + %iota.554 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/iota" stack_frame_id=588} + %mul.3263 = f32[64]{0} multiply(%iota.554, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=588} + %div.1020 = f32[64]{0} divide(%mul.3263, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/div" stack_frame_id=589} + %neg.410 = f32[64]{0} negate(%div.1020), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/neg" stack_frame_id=590} + %pow.405 = f32[64]{0} power(%broadcast.51, %neg.410), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/pow" stack_frame_id=593} + %div.1021 = f32[64]{0} divide(%pow.405, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/div" stack_frame_id=594} + %dot_general.1106 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1589, %div.1021), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/bi,j->bij/dot_general" stack_frame_id=604} + %stack.1608 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1106), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/stack" stack_frame_id=607} + %stack.1609 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1106), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/stack" stack_frame_id=607} + %stack.1610 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1608, %stack.1609), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/stack" stack_frame_id=607} + %reshape.738 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1610), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/reshape"} + %cos.201 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.738), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/cos" stack_frame_id=611} + %convert_element_type.1817 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.201), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/convert_element_type" stack_frame_id=615} + %mul.3264 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1817), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=635} + %mul.3265 = bf16[1,41,128]{2,1,0} reshape(%mul.3264), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=635} + %mul.3266 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3265), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=635} + %mul.3267 = bf16[1,41,32,128]{3,2,1,0} multiply(%dot_general.1105, %mul.3266), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=635} + %split.403 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1105), slice={[0:1], [0:41], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/split" stack_frame_id=627} + %neg.411 = bf16[1,41,32,64]{3,2,1,0} negate(%split.403), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/neg" stack_frame_id=628} + %stack.1611 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%neg.411), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/stack" stack_frame_id=631} + %split.402 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1105), slice={[0:1], [0:41], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/split" stack_frame_id=627} + %stack.1612 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%split.402), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/stack" stack_frame_id=631} + %stack.1613 = bf16[1,41,32,2,64]{4,3,2,1,0} concatenate(%stack.1611, %stack.1612), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/stack" stack_frame_id=631} + %reshape.739 = bf16[1,41,32,128]{3,2,1,0} reshape(%stack.1613), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/reshape" stack_frame_id=634} + %sin.201 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.738), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/sin" stack_frame_id=619} + %convert_element_type.1818 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.201), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/convert_element_type" stack_frame_id=623} + %mul.3268 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1818), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=636} + %mul.3269 = bf16[1,41,128]{2,1,0} reshape(%mul.3268), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=636} + %mul.3270 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3269), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=636} + %mul.3271 = bf16[1,41,32,128]{3,2,1,0} multiply(%reshape.739, %mul.3270), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/mul" stack_frame_id=636} + %add.1160 = bf16[1,41,32,128]{3,2,1,0} add(%mul.3267, %mul.3271), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/rotary_embedding_5/add" stack_frame_id=637} + %reshape.744 = bf16[1,41,32,1,128]{4,3,2,1,0} reshape(%add.1160), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/reshape" stack_frame_id=34} + %broadcast_in_dim.1591 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%add.1161), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/broadcast_in_dim" stack_frame_id=722} + %reshape.742 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1591), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/reshape" stack_frame_id=722} + %dot_general.1110 = f32[1,32,41,1,41]{4,3,2,1,0} dot(%reshape.744, %reshape.742), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=34} + %mul.3281 = f32[1,32,41,1,41]{4,3,2,1,0} multiply(%dot_general.1110, %broadcast.48), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/vmap()/mul" stack_frame_id=34} + %vmap_jit__where__.100 = f32[1,1,32,41,41]{4,3,2,1,0} call(%and.315, %mul.3281, %constant.153), to_apply=%_where_0.5, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/vmap(jit(_where))" stack_frame_id=34} + %reduce_max.486 = f32[1,1,32,41]{3,2,1,0} reduce(%vmap_jit__where__.100, %constant.152), dimensions={4}, to_apply=%region_22.27, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/vmap()/reduce_max" stack_frame_id=34} + %max.110 = f32[1,1,32,41]{3,2,1,0} maximum(%reduce_max.486, %broadcast.47), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/vmap()/max" stack_frame_id=34} + %broadcast_in_dim.1595 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%max.110), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %sub.424 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1595), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/vmap()/sub" stack_frame_id=34} + %sub.425 = f32[1,1,32,41]{3,2,1,0} reshape(%sub.424), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/vmap()/sub" stack_frame_id=34} + %sub.426 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%sub.425), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/vmap()/sub" stack_frame_id=34} + %sub.427 = f32[1,1,32,41,41]{4,3,2,1,0} subtract(%vmap_jit__where__.100, %sub.426), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/vmap()/sub" stack_frame_id=34} + %exp.106 = f32[1,1,32,41,41]{4,3,2,1,0} exponential(%sub.427), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/vmap()/exp" stack_frame_id=34} + %reduce_sum.1465 = f32[1,1,32,41]{3,2,1,0} reduce(%exp.106, %constant.155), dimensions={4}, to_apply=%region_23.28, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/vmap()/reduce_sum" stack_frame_id=34} + %broadcast_in_dim.1596 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%reduce_sum.1465), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %div.1024 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1596), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/vmap()/div" stack_frame_id=34} + %div.1025 = f32[1,1,32,41]{3,2,1,0} reshape(%div.1024), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/vmap()/div" stack_frame_id=34} + %div.1026 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%div.1025), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/vmap()/div" stack_frame_id=34} + %div.1027 = f32[1,1,32,41,41]{4,3,2,1,0} divide(%exp.106, %div.1026), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/vmap()/div" stack_frame_id=34} + %convert_element_type.1822 = bf16[1,1,32,41,41]{4,3,2,1,0} convert(%div.1027), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/vmap()/convert_element_type" stack_frame_id=34} + %dot_general.1111 = bf16[1,32,128,1,41]{4,3,2,1,0} dot(%reshape.743, %convert_element_type.1822), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=34} + %transpose.104 = bf16[1,41,32,1,128]{1,3,4,2,0} transpose(%dot_general.1111), dimensions={0,4,1,3,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/vmap()/transpose;jit(compiled_generate_function)/transformer_layer_5/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %reshape.745 = bf16[1,41,32,128]{3,2,1,0} reshape(%transpose.104), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/reshape" stack_frame_id=34} + %state_1__50_.1 = bf16[32,128,4096]{2,1,0} parameter(53), metadata={op_name="state[1][50]"} + %dot_general.1112 = bf16[1,41,4096]{2,1,0} dot(%reshape.745, %state_1__50_.1), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=741} + %add.1162 = bf16[1,41,4096]{2,1,0} add(%dot_general.1112, %add.1157), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/add" stack_frame_id=742} + %convert_element_type.1823 = f32[1,41,4096]{2,1,0} convert(%add.1162), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/feedforward_layernorm/convert_element_type" stack_frame_id=746} + %pow.407 = f32[1,41,4096]{2,1,0} power(%convert_element_type.1823, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/feedforward_layernorm/pow" stack_frame_id=749} + %reduce_sum.1466 = f32[1,41]{1,0} reduce(%pow.407, %constant.155), dimensions={2}, to_apply=%region_24.29, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/feedforward_layernorm/reduce_sum" stack_frame_id=43} + %broadcast_in_dim.1597 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1466), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/feedforward_layernorm/broadcast_in_dim" stack_frame_id=43} + %div.1028 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1597, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/feedforward_layernorm/div" stack_frame_id=43} + %add.1163 = f32[1,41,1]{2,1,0} add(%div.1028, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/feedforward_layernorm/add" stack_frame_id=750} + %rsqrt.204 = f32[1,41,1]{2,1,0} rsqrt(%add.1163), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/feedforward_layernorm/rsqrt" stack_frame_id=753} + %mul.3282 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.204), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3283 = f32[1,41]{1,0} reshape(%mul.3282), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3284 = f32[1,41,4096]{2,1,0} broadcast(%mul.3283), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3285 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.1823, %mul.3284), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/feedforward_layernorm/mul" stack_frame_id=754} + %state_1__55_.1 = bf16[4096]{0} parameter(58), metadata={op_name="state[1][55]"} + %convert_element_type.1824 = f32[4096]{0} convert(%state_1__55_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/feedforward_layernorm/convert_element_type" stack_frame_id=755} + %broadcast_in_dim.1598 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1824), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/feedforward_layernorm/broadcast_in_dim" stack_frame_id=755} + %mul.3286 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1598), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3287 = f32[1,4096]{1,0} reshape(%mul.3286), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3288 = f32[1,41,4096]{2,1,0} broadcast(%mul.3287), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3289 = f32[1,41,4096]{2,1,0} multiply(%mul.3285, %mul.3288), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/feedforward_layernorm/mul" stack_frame_id=755} + %convert_element_type.1825 = bf16[1,41,4096]{2,1,0} convert(%mul.3289), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/feedforward_layernorm/convert_element_type" stack_frame_id=759} + %state_1__52_.1 = bf16[4096,14336]{1,0} parameter(55), metadata={op_name="state[1][52]"} + %dot_general.1114 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.1825, %state_1__52_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/feedforward_intermediate_dense/dot_general" stack_frame_id=788} + %state_1__53_.1 = bf16[4096,14336]{1,0} parameter(56), metadata={op_name="state[1][53]"} + %dot_general.1113 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.1825, %state_1__53_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/feedforward_gate_dense/dot_general" stack_frame_id=768} + %convert_element_type.1826 = f32[1,41,14336]{2,1,0} convert(%dot_general.1113), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/convert_element_type" stack_frame_id=772} + %jit_silu_.100 = f32[1,41,14336]{2,1,0} call(%convert_element_type.1826), to_apply=%silu.9, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/jit(silu)" stack_frame_id=775} + %convert_element_type.1827 = bf16[1,41,14336]{2,1,0} convert(%jit_silu_.100), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/convert_element_type" stack_frame_id=779} + %mul.3290 = bf16[1,41,14336]{2,1,0} multiply(%dot_general.1114, %convert_element_type.1827), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/mul" stack_frame_id=791} + %state_1__54_.1 = bf16[14336,4096]{1,0} parameter(57), metadata={op_name="state[1][54]"} + %dot_general.1115 = bf16[1,41,4096]{2,1,0} dot(%mul.3290, %state_1__54_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/feedforward_output_dense/dot_general" stack_frame_id=800} + %add.1164 = bf16[1,41,4096]{2,1,0} add(%dot_general.1115, %add.1162), metadata={op_name="jit(compiled_generate_function)/transformer_layer_5/add" stack_frame_id=801} + %convert_element_type.1830 = f32[1,41,4096]{2,1,0} convert(%add.1164), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention_layernorm/convert_element_type" stack_frame_id=555} + %pow.408 = f32[1,41,4096]{2,1,0} power(%convert_element_type.1830, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention_layernorm/pow" stack_frame_id=558} + %reduce_sum.1467 = f32[1,41]{1,0} reduce(%pow.408, %constant.155), dimensions={2}, to_apply=%region_25.30, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention_layernorm/reduce_sum" stack_frame_id=24} + %broadcast_in_dim.1605 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1467), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention_layernorm/broadcast_in_dim" stack_frame_id=24} + %div.1029 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1605, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention_layernorm/div" stack_frame_id=24} + %add.1166 = f32[1,41,1]{2,1,0} add(%div.1029, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention_layernorm/add" stack_frame_id=559} + %rsqrt.205 = f32[1,41,1]{2,1,0} rsqrt(%add.1166), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention_layernorm/rsqrt" stack_frame_id=562} + %mul.3291 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.205), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3292 = f32[1,41]{1,0} reshape(%mul.3291), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3293 = f32[1,41,4096]{2,1,0} broadcast(%mul.3292), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3294 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.1830, %mul.3293), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention_layernorm/mul" stack_frame_id=563} + %state_1__60_.1 = bf16[4096]{0} parameter(63), metadata={op_name="state[1][60]"} + %convert_element_type.1831 = f32[4096]{0} convert(%state_1__60_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention_layernorm/convert_element_type" stack_frame_id=564} + %broadcast_in_dim.1606 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1831), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention_layernorm/broadcast_in_dim" stack_frame_id=564} + %mul.3295 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1606), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3296 = f32[1,4096]{1,0} reshape(%mul.3295), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3297 = f32[1,41,4096]{2,1,0} broadcast(%mul.3296), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3298 = f32[1,41,4096]{2,1,0} multiply(%mul.3294, %mul.3297), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention_layernorm/mul" stack_frame_id=564} + %convert_element_type.1832 = bf16[1,41,4096]{2,1,0} convert(%mul.3298), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention_layernorm/convert_element_type" stack_frame_id=568} + %state_1__57_.1 = bf16[4096,8,128]{2,1,0} parameter(60), metadata={op_name="state[1][57]"} + %dot_general.1118 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.1832, %state_1__57_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=647} + %iota.565 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/iota" stack_frame_id=677} + %broadcast_in_dim.1608 = f32[1,41]{1,0} reshape(%iota.565), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/broadcast_in_dim" stack_frame_id=680} + %iota.564 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/iota" stack_frame_id=667} + %mul.3308 = f32[64]{0} multiply(%iota.564, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=667} + %div.1032 = f32[64]{0} divide(%mul.3308, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/div" stack_frame_id=668} + %neg.416 = f32[64]{0} negate(%div.1032), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/neg" stack_frame_id=669} + %pow.410 = f32[64]{0} power(%broadcast.51, %neg.416), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/pow" stack_frame_id=672} + %div.1033 = f32[64]{0} divide(%pow.410, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/div" stack_frame_id=673} + %dot_general.1120 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1608, %div.1033), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/bi,j->bij/dot_general" stack_frame_id=683} + %stack.1629 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1120), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/stack" stack_frame_id=686} + %stack.1630 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1120), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/stack" stack_frame_id=686} + %stack.1631 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1629, %stack.1630), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/stack" stack_frame_id=686} + %reshape.748 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1631), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/reshape"} + %cos.204 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.748), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/cos" stack_frame_id=690} + %convert_element_type.1835 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.204), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/convert_element_type" stack_frame_id=694} + %mul.3309 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1835), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=714} + %mul.3310 = bf16[1,41,128]{2,1,0} reshape(%mul.3309), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=714} + %mul.3311 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3310), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=714} + %mul.3312 = bf16[1,41,8,128]{3,2,1,0} multiply(%dot_general.1118, %mul.3311), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=714} + %split.409 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1118), slice={[0:1], [0:41], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/split" stack_frame_id=706} + %neg.417 = bf16[1,41,8,64]{3,2,1,0} negate(%split.409), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/neg" stack_frame_id=707} + %stack.1632 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%neg.417), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/stack" stack_frame_id=710} + %split.408 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1118), slice={[0:1], [0:41], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/split" stack_frame_id=706} + %stack.1633 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%split.408), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/stack" stack_frame_id=710} + %stack.1634 = bf16[1,41,8,2,64]{4,3,2,1,0} concatenate(%stack.1632, %stack.1633), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/stack" stack_frame_id=710} + %reshape.749 = bf16[1,41,8,128]{3,2,1,0} reshape(%stack.1634), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/reshape" stack_frame_id=713} + %sin.204 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.748), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/sin" stack_frame_id=698} + %convert_element_type.1836 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.204), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/convert_element_type" stack_frame_id=702} + %mul.3313 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1836), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=715} + %mul.3314 = bf16[1,41,128]{2,1,0} reshape(%mul.3313), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=715} + %mul.3315 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3314), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=715} + %mul.3316 = bf16[1,41,8,128]{3,2,1,0} multiply(%reshape.749, %mul.3315), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=715} + %add.1168 = bf16[1,41,8,128]{3,2,1,0} add(%mul.3312, %mul.3316), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/add" stack_frame_id=716} + %stack.1635 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%add.1168), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/stack" stack_frame_id=719} + %state_1__58_.1 = bf16[4096,8,128]{2,1,0} parameter(61), metadata={op_name="state[1][58]"} + %dot_general.1119 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.1832, %state_1__58_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=656} + %stack.1636 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dot_general.1119), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/stack" stack_frame_id=719} + %stack.1637 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1635, %stack.1636), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/stack" stack_frame_id=719} + %stack.2013 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1637), metadata={op_name="jit(compiled_generate_function)/stack" stack_frame_id=804} + %broadcast_in_dim.1610 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dot_general.1119), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/broadcast_in_dim" stack_frame_id=725} + %reshape.751 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1610), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/reshape" stack_frame_id=725} + %iota.558 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/iota" stack_frame_id=525} + %broadcast_in_dim.1599 = f32[41,1]{1,0} reshape(%iota.558), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/broadcast_in_dim" stack_frame_id=528} + %ge.523 = f32[41,1]{1,0} broadcast(%broadcast_in_dim.1599), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/ge" stack_frame_id=532} + %ge.524 = f32[41]{0} reshape(%ge.523), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/ge" stack_frame_id=532} + %ge.525 = f32[41,41]{1,0} broadcast(%ge.524), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/ge" stack_frame_id=532} + %iota.559 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/iota" stack_frame_id=531} + %broadcast_in_dim.1600 = f32[1,41]{1,0} reshape(%iota.559), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/broadcast_in_dim" stack_frame_id=532} + %ge.526 = f32[1,41]{1,0} broadcast(%broadcast_in_dim.1600), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/ge" stack_frame_id=532} + %ge.527 = f32[41]{0} reshape(%ge.526), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/ge" stack_frame_id=532} + %ge.528 = f32[41,41]{1,0} broadcast(%ge.527), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/ge" stack_frame_id=532} + %ge.529 = pred[41,41]{1,0} compare(%ge.525, %ge.528), direction=GE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/ge" stack_frame_id=532} + %broadcast_in_dim.1601 = pred[1,41,41]{2,1,0} reshape(%ge.529), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/broadcast_in_dim" stack_frame_id=535} + %convert_element_type.1829 = s32[1,41,41]{2,1,0} convert(%broadcast_in_dim.1601), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/convert_element_type" stack_frame_id=551} + %iota.560 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/iota" stack_frame_id=538} + %broadcast_in_dim.1602 = s32[41,1]{1,0} reshape(%iota.560), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/broadcast_in_dim" stack_frame_id=536} + %lt.665 = s32[41,1]{1,0} broadcast(%broadcast_in_dim.1602), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/lt" stack_frame_id=543} + %lt.666 = s32[41]{0} reshape(%lt.665), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/lt" stack_frame_id=543} + %lt.667 = s32[41,41]{1,0} broadcast(%lt.666), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/lt" stack_frame_id=543} + %iota.561 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/iota" stack_frame_id=541} + %broadcast_in_dim.1603 = s32[1,41]{1,0} reshape(%iota.561), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/broadcast_in_dim" stack_frame_id=539} + %add.1165 = s32[1,41]{1,0} add(%broadcast_in_dim.1603, %broadcast.57), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/add" stack_frame_id=542} + %lt.668 = s32[1,41]{1,0} broadcast(%add.1165), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/lt" stack_frame_id=543} + %lt.669 = s32[41]{0} reshape(%lt.668), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/lt" stack_frame_id=543} + %lt.670 = s32[41,41]{1,0} broadcast(%lt.669), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/lt" stack_frame_id=543} + %lt.671 = pred[41,41]{1,0} compare(%lt.667, %lt.670), direction=LT, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/lt" stack_frame_id=543} + %convert_element_type.1828 = s32[41,41]{1,0} convert(%lt.671), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/convert_element_type" stack_frame_id=547} + %broadcast_in_dim.1604 = s32[1,41,41]{2,1,0} reshape(%convert_element_type.1828), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/broadcast_in_dim" stack_frame_id=551} + %min.101 = s32[1,41,41]{2,1,0} minimum(%convert_element_type.1829, %broadcast_in_dim.1604), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/min" stack_frame_id=551} + %broadcast_in_dim.1611 = s32[1,1,41,41]{3,2,1,0} reshape(%min.101), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/broadcast_in_dim" stack_frame_id=728} + %convert_element_type.1837 = pred[1,1,41,41]{3,2,1,0} compare(%broadcast_in_dim.1611, %broadcast.49), direction=NE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/convert_element_type" stack_frame_id=732} + %broadcast_in_dim.1612 = pred[1,1,1,41,41]{4,3,2,1,0} reshape(%convert_element_type.1837), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/broadcast_in_dim" stack_frame_id=34} + %and.316 = pred[1,1,1,41,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1612), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/vmap()/and" stack_frame_id=34} + %and.317 = pred[1,1,41,41]{3,2,1,0} reshape(%and.316), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/vmap()/and" stack_frame_id=34} + %and.318 = pred[1,1,32,41,41]{4,3,2,1,0} broadcast(%and.317), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/vmap()/and" stack_frame_id=34} + %state_1__56_.1 = bf16[4096,32,128]{2,1,0} parameter(59), metadata={op_name="state[1][56]"} + %dot_general.1116 = bf16[1,41,32,128]{3,2,1,0} dot(%convert_element_type.1832, %state_1__56_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=577} + %iota.563 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/iota" stack_frame_id=598} + %broadcast_in_dim.1607 = f32[1,41]{1,0} reshape(%iota.563), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/broadcast_in_dim" stack_frame_id=601} + %iota.562 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/iota" stack_frame_id=588} + %mul.3299 = f32[64]{0} multiply(%iota.562, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=588} + %div.1030 = f32[64]{0} divide(%mul.3299, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/div" stack_frame_id=589} + %neg.414 = f32[64]{0} negate(%div.1030), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/neg" stack_frame_id=590} + %pow.409 = f32[64]{0} power(%broadcast.51, %neg.414), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/pow" stack_frame_id=593} + %div.1031 = f32[64]{0} divide(%pow.409, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/div" stack_frame_id=594} + %dot_general.1117 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1607, %div.1031), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/bi,j->bij/dot_general" stack_frame_id=604} + %stack.1623 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1117), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/stack" stack_frame_id=607} + %stack.1624 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1117), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/stack" stack_frame_id=607} + %stack.1625 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1623, %stack.1624), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/stack" stack_frame_id=607} + %reshape.746 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1625), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/reshape"} + %cos.203 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.746), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/cos" stack_frame_id=611} + %convert_element_type.1833 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.203), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/convert_element_type" stack_frame_id=615} + %mul.3300 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1833), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=635} + %mul.3301 = bf16[1,41,128]{2,1,0} reshape(%mul.3300), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=635} + %mul.3302 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3301), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=635} + %mul.3303 = bf16[1,41,32,128]{3,2,1,0} multiply(%dot_general.1116, %mul.3302), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=635} + %split.407 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1116), slice={[0:1], [0:41], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/split" stack_frame_id=627} + %neg.415 = bf16[1,41,32,64]{3,2,1,0} negate(%split.407), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/neg" stack_frame_id=628} + %stack.1626 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%neg.415), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/stack" stack_frame_id=631} + %split.406 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1116), slice={[0:1], [0:41], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/split" stack_frame_id=627} + %stack.1627 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%split.406), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/stack" stack_frame_id=631} + %stack.1628 = bf16[1,41,32,2,64]{4,3,2,1,0} concatenate(%stack.1626, %stack.1627), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/stack" stack_frame_id=631} + %reshape.747 = bf16[1,41,32,128]{3,2,1,0} reshape(%stack.1628), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/reshape" stack_frame_id=634} + %sin.203 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.746), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/sin" stack_frame_id=619} + %convert_element_type.1834 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.203), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/convert_element_type" stack_frame_id=623} + %mul.3304 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1834), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=636} + %mul.3305 = bf16[1,41,128]{2,1,0} reshape(%mul.3304), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=636} + %mul.3306 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3305), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=636} + %mul.3307 = bf16[1,41,32,128]{3,2,1,0} multiply(%reshape.747, %mul.3306), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/mul" stack_frame_id=636} + %add.1167 = bf16[1,41,32,128]{3,2,1,0} add(%mul.3303, %mul.3307), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/rotary_embedding_6/add" stack_frame_id=637} + %reshape.752 = bf16[1,41,32,1,128]{4,3,2,1,0} reshape(%add.1167), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/reshape" stack_frame_id=34} + %broadcast_in_dim.1609 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%add.1168), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/broadcast_in_dim" stack_frame_id=722} + %reshape.750 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1609), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/reshape" stack_frame_id=722} + %dot_general.1121 = f32[1,32,41,1,41]{4,3,2,1,0} dot(%reshape.752, %reshape.750), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=34} + %mul.3317 = f32[1,32,41,1,41]{4,3,2,1,0} multiply(%dot_general.1121, %broadcast.48), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/vmap()/mul" stack_frame_id=34} + %vmap_jit__where__.101 = f32[1,1,32,41,41]{4,3,2,1,0} call(%and.318, %mul.3317, %constant.153), to_apply=%_where_0.5, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/vmap(jit(_where))" stack_frame_id=34} + %reduce_max.487 = f32[1,1,32,41]{3,2,1,0} reduce(%vmap_jit__where__.101, %constant.152), dimensions={4}, to_apply=%region_26.31, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/vmap()/reduce_max" stack_frame_id=34} + %max.111 = f32[1,1,32,41]{3,2,1,0} maximum(%reduce_max.487, %broadcast.47), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/vmap()/max" stack_frame_id=34} + %broadcast_in_dim.1613 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%max.111), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %sub.428 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1613), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/vmap()/sub" stack_frame_id=34} + %sub.429 = f32[1,1,32,41]{3,2,1,0} reshape(%sub.428), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/vmap()/sub" stack_frame_id=34} + %sub.430 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%sub.429), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/vmap()/sub" stack_frame_id=34} + %sub.431 = f32[1,1,32,41,41]{4,3,2,1,0} subtract(%vmap_jit__where__.101, %sub.430), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/vmap()/sub" stack_frame_id=34} + %exp.107 = f32[1,1,32,41,41]{4,3,2,1,0} exponential(%sub.431), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/vmap()/exp" stack_frame_id=34} + %reduce_sum.1468 = f32[1,1,32,41]{3,2,1,0} reduce(%exp.107, %constant.155), dimensions={4}, to_apply=%region_27.32, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/vmap()/reduce_sum" stack_frame_id=34} + %broadcast_in_dim.1614 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%reduce_sum.1468), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %div.1034 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1614), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/vmap()/div" stack_frame_id=34} + %div.1035 = f32[1,1,32,41]{3,2,1,0} reshape(%div.1034), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/vmap()/div" stack_frame_id=34} + %div.1036 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%div.1035), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/vmap()/div" stack_frame_id=34} + %div.1037 = f32[1,1,32,41,41]{4,3,2,1,0} divide(%exp.107, %div.1036), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/vmap()/div" stack_frame_id=34} + %convert_element_type.1838 = bf16[1,1,32,41,41]{4,3,2,1,0} convert(%div.1037), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/vmap()/convert_element_type" stack_frame_id=34} + %dot_general.1122 = bf16[1,32,128,1,41]{4,3,2,1,0} dot(%reshape.751, %convert_element_type.1838), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=34} + %transpose.105 = bf16[1,41,32,1,128]{1,3,4,2,0} transpose(%dot_general.1122), dimensions={0,4,1,3,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/vmap()/transpose;jit(compiled_generate_function)/transformer_layer_6/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %reshape.753 = bf16[1,41,32,128]{3,2,1,0} reshape(%transpose.105), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/reshape" stack_frame_id=34} + %state_1__59_.1 = bf16[32,128,4096]{2,1,0} parameter(62), metadata={op_name="state[1][59]"} + %dot_general.1123 = bf16[1,41,4096]{2,1,0} dot(%reshape.753, %state_1__59_.1), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=741} + %add.1169 = bf16[1,41,4096]{2,1,0} add(%dot_general.1123, %add.1164), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/add" stack_frame_id=742} + %convert_element_type.1839 = f32[1,41,4096]{2,1,0} convert(%add.1169), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/feedforward_layernorm/convert_element_type" stack_frame_id=746} + %pow.411 = f32[1,41,4096]{2,1,0} power(%convert_element_type.1839, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/feedforward_layernorm/pow" stack_frame_id=749} + %reduce_sum.1469 = f32[1,41]{1,0} reduce(%pow.411, %constant.155), dimensions={2}, to_apply=%region_28.33, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/feedforward_layernorm/reduce_sum" stack_frame_id=43} + %broadcast_in_dim.1615 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1469), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/feedforward_layernorm/broadcast_in_dim" stack_frame_id=43} + %div.1038 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1615, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/feedforward_layernorm/div" stack_frame_id=43} + %add.1170 = f32[1,41,1]{2,1,0} add(%div.1038, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/feedforward_layernorm/add" stack_frame_id=750} + %rsqrt.206 = f32[1,41,1]{2,1,0} rsqrt(%add.1170), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/feedforward_layernorm/rsqrt" stack_frame_id=753} + %mul.3318 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.206), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3319 = f32[1,41]{1,0} reshape(%mul.3318), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3320 = f32[1,41,4096]{2,1,0} broadcast(%mul.3319), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3321 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.1839, %mul.3320), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/feedforward_layernorm/mul" stack_frame_id=754} + %state_1__64_.1 = bf16[4096]{0} parameter(67), metadata={op_name="state[1][64]"} + %convert_element_type.1840 = f32[4096]{0} convert(%state_1__64_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/feedforward_layernorm/convert_element_type" stack_frame_id=755} + %broadcast_in_dim.1616 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1840), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/feedforward_layernorm/broadcast_in_dim" stack_frame_id=755} + %mul.3322 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1616), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3323 = f32[1,4096]{1,0} reshape(%mul.3322), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3324 = f32[1,41,4096]{2,1,0} broadcast(%mul.3323), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3325 = f32[1,41,4096]{2,1,0} multiply(%mul.3321, %mul.3324), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/feedforward_layernorm/mul" stack_frame_id=755} + %convert_element_type.1841 = bf16[1,41,4096]{2,1,0} convert(%mul.3325), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/feedforward_layernorm/convert_element_type" stack_frame_id=759} + %state_1__61_.1 = bf16[4096,14336]{1,0} parameter(64), metadata={op_name="state[1][61]"} + %dot_general.1125 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.1841, %state_1__61_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/feedforward_intermediate_dense/dot_general" stack_frame_id=788} + %state_1__62_.1 = bf16[4096,14336]{1,0} parameter(65), metadata={op_name="state[1][62]"} + %dot_general.1124 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.1841, %state_1__62_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/feedforward_gate_dense/dot_general" stack_frame_id=768} + %convert_element_type.1842 = f32[1,41,14336]{2,1,0} convert(%dot_general.1124), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/convert_element_type" stack_frame_id=772} + %jit_silu_.101 = f32[1,41,14336]{2,1,0} call(%convert_element_type.1842), to_apply=%silu.9, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/jit(silu)" stack_frame_id=775} + %convert_element_type.1843 = bf16[1,41,14336]{2,1,0} convert(%jit_silu_.101), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/convert_element_type" stack_frame_id=779} + %mul.3326 = bf16[1,41,14336]{2,1,0} multiply(%dot_general.1125, %convert_element_type.1843), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/mul" stack_frame_id=791} + %state_1__63_.1 = bf16[14336,4096]{1,0} parameter(66), metadata={op_name="state[1][63]"} + %dot_general.1126 = bf16[1,41,4096]{2,1,0} dot(%mul.3326, %state_1__63_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/feedforward_output_dense/dot_general" stack_frame_id=800} + %add.1171 = bf16[1,41,4096]{2,1,0} add(%dot_general.1126, %add.1169), metadata={op_name="jit(compiled_generate_function)/transformer_layer_6/add" stack_frame_id=801} + %convert_element_type.1846 = f32[1,41,4096]{2,1,0} convert(%add.1171), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention_layernorm/convert_element_type" stack_frame_id=555} + %pow.412 = f32[1,41,4096]{2,1,0} power(%convert_element_type.1846, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention_layernorm/pow" stack_frame_id=558} + %reduce_sum.1470 = f32[1,41]{1,0} reduce(%pow.412, %constant.155), dimensions={2}, to_apply=%region_29.34, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention_layernorm/reduce_sum" stack_frame_id=24} + %broadcast_in_dim.1623 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1470), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention_layernorm/broadcast_in_dim" stack_frame_id=24} + %div.1039 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1623, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention_layernorm/div" stack_frame_id=24} + %add.1173 = f32[1,41,1]{2,1,0} add(%div.1039, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention_layernorm/add" stack_frame_id=559} + %rsqrt.207 = f32[1,41,1]{2,1,0} rsqrt(%add.1173), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention_layernorm/rsqrt" stack_frame_id=562} + %mul.3327 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.207), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3328 = f32[1,41]{1,0} reshape(%mul.3327), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3329 = f32[1,41,4096]{2,1,0} broadcast(%mul.3328), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3330 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.1846, %mul.3329), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention_layernorm/mul" stack_frame_id=563} + %state_1__69_.1 = bf16[4096]{0} parameter(72), metadata={op_name="state[1][69]"} + %convert_element_type.1847 = f32[4096]{0} convert(%state_1__69_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention_layernorm/convert_element_type" stack_frame_id=564} + %broadcast_in_dim.1624 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1847), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention_layernorm/broadcast_in_dim" stack_frame_id=564} + %mul.3331 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1624), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3332 = f32[1,4096]{1,0} reshape(%mul.3331), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3333 = f32[1,41,4096]{2,1,0} broadcast(%mul.3332), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3334 = f32[1,41,4096]{2,1,0} multiply(%mul.3330, %mul.3333), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention_layernorm/mul" stack_frame_id=564} + %convert_element_type.1848 = bf16[1,41,4096]{2,1,0} convert(%mul.3334), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention_layernorm/convert_element_type" stack_frame_id=568} + %state_1__66_.1 = bf16[4096,8,128]{2,1,0} parameter(69), metadata={op_name="state[1][66]"} + %dot_general.1129 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.1848, %state_1__66_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=647} + %iota.573 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/iota" stack_frame_id=677} + %broadcast_in_dim.1626 = f32[1,41]{1,0} reshape(%iota.573), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/broadcast_in_dim" stack_frame_id=680} + %iota.572 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/iota" stack_frame_id=667} + %mul.3344 = f32[64]{0} multiply(%iota.572, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=667} + %div.1042 = f32[64]{0} divide(%mul.3344, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/div" stack_frame_id=668} + %neg.420 = f32[64]{0} negate(%div.1042), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/neg" stack_frame_id=669} + %pow.414 = f32[64]{0} power(%broadcast.51, %neg.420), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/pow" stack_frame_id=672} + %div.1043 = f32[64]{0} divide(%pow.414, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/div" stack_frame_id=673} + %dot_general.1131 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1626, %div.1043), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/bi,j->bij/dot_general" stack_frame_id=683} + %stack.1644 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1131), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/stack" stack_frame_id=686} + %stack.1645 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1131), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/stack" stack_frame_id=686} + %stack.1646 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1644, %stack.1645), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/stack" stack_frame_id=686} + %reshape.756 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1646), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/reshape"} + %cos.206 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.756), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/cos" stack_frame_id=690} + %convert_element_type.1851 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.206), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/convert_element_type" stack_frame_id=694} + %mul.3345 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1851), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=714} + %mul.3346 = bf16[1,41,128]{2,1,0} reshape(%mul.3345), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=714} + %mul.3347 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3346), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=714} + %mul.3348 = bf16[1,41,8,128]{3,2,1,0} multiply(%dot_general.1129, %mul.3347), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=714} + %split.413 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1129), slice={[0:1], [0:41], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/split" stack_frame_id=706} + %neg.421 = bf16[1,41,8,64]{3,2,1,0} negate(%split.413), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/neg" stack_frame_id=707} + %stack.1647 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%neg.421), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/stack" stack_frame_id=710} + %split.412 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1129), slice={[0:1], [0:41], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/split" stack_frame_id=706} + %stack.1648 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%split.412), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/stack" stack_frame_id=710} + %stack.1649 = bf16[1,41,8,2,64]{4,3,2,1,0} concatenate(%stack.1647, %stack.1648), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/stack" stack_frame_id=710} + %reshape.757 = bf16[1,41,8,128]{3,2,1,0} reshape(%stack.1649), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/reshape" stack_frame_id=713} + %sin.206 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.756), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/sin" stack_frame_id=698} + %convert_element_type.1852 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.206), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/convert_element_type" stack_frame_id=702} + %mul.3349 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1852), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=715} + %mul.3350 = bf16[1,41,128]{2,1,0} reshape(%mul.3349), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=715} + %mul.3351 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3350), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=715} + %mul.3352 = bf16[1,41,8,128]{3,2,1,0} multiply(%reshape.757, %mul.3351), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=715} + %add.1175 = bf16[1,41,8,128]{3,2,1,0} add(%mul.3348, %mul.3352), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/add" stack_frame_id=716} + %stack.1650 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%add.1175), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/stack" stack_frame_id=719} + %state_1__67_.1 = bf16[4096,8,128]{2,1,0} parameter(70), metadata={op_name="state[1][67]"} + %dot_general.1130 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.1848, %state_1__67_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=656} + %stack.1651 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dot_general.1130), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/stack" stack_frame_id=719} + %stack.1652 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1650, %stack.1651), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/stack" stack_frame_id=719} + %stack.2014 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1652), metadata={op_name="jit(compiled_generate_function)/stack" stack_frame_id=804} + %broadcast_in_dim.1628 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dot_general.1130), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/broadcast_in_dim" stack_frame_id=725} + %reshape.759 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1628), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/reshape" stack_frame_id=725} + %iota.566 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/iota" stack_frame_id=525} + %broadcast_in_dim.1617 = f32[41,1]{1,0} reshape(%iota.566), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/broadcast_in_dim" stack_frame_id=528} + %ge.530 = f32[41,1]{1,0} broadcast(%broadcast_in_dim.1617), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/ge" stack_frame_id=532} + %ge.531 = f32[41]{0} reshape(%ge.530), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/ge" stack_frame_id=532} + %ge.532 = f32[41,41]{1,0} broadcast(%ge.531), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/ge" stack_frame_id=532} + %iota.567 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/iota" stack_frame_id=531} + %broadcast_in_dim.1618 = f32[1,41]{1,0} reshape(%iota.567), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/broadcast_in_dim" stack_frame_id=532} + %ge.533 = f32[1,41]{1,0} broadcast(%broadcast_in_dim.1618), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/ge" stack_frame_id=532} + %ge.534 = f32[41]{0} reshape(%ge.533), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/ge" stack_frame_id=532} + %ge.535 = f32[41,41]{1,0} broadcast(%ge.534), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/ge" stack_frame_id=532} + %ge.536 = pred[41,41]{1,0} compare(%ge.532, %ge.535), direction=GE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/ge" stack_frame_id=532} + %broadcast_in_dim.1619 = pred[1,41,41]{2,1,0} reshape(%ge.536), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/broadcast_in_dim" stack_frame_id=535} + %convert_element_type.1845 = s32[1,41,41]{2,1,0} convert(%broadcast_in_dim.1619), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/convert_element_type" stack_frame_id=551} + %iota.568 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/iota" stack_frame_id=538} + %broadcast_in_dim.1620 = s32[41,1]{1,0} reshape(%iota.568), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/broadcast_in_dim" stack_frame_id=536} + %lt.672 = s32[41,1]{1,0} broadcast(%broadcast_in_dim.1620), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/lt" stack_frame_id=543} + %lt.673 = s32[41]{0} reshape(%lt.672), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/lt" stack_frame_id=543} + %lt.674 = s32[41,41]{1,0} broadcast(%lt.673), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/lt" stack_frame_id=543} + %iota.569 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/iota" stack_frame_id=541} + %broadcast_in_dim.1621 = s32[1,41]{1,0} reshape(%iota.569), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/broadcast_in_dim" stack_frame_id=539} + %add.1172 = s32[1,41]{1,0} add(%broadcast_in_dim.1621, %broadcast.57), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/add" stack_frame_id=542} + %lt.675 = s32[1,41]{1,0} broadcast(%add.1172), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/lt" stack_frame_id=543} + %lt.676 = s32[41]{0} reshape(%lt.675), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/lt" stack_frame_id=543} + %lt.677 = s32[41,41]{1,0} broadcast(%lt.676), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/lt" stack_frame_id=543} + %lt.678 = pred[41,41]{1,0} compare(%lt.674, %lt.677), direction=LT, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/lt" stack_frame_id=543} + %convert_element_type.1844 = s32[41,41]{1,0} convert(%lt.678), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/convert_element_type" stack_frame_id=547} + %broadcast_in_dim.1622 = s32[1,41,41]{2,1,0} reshape(%convert_element_type.1844), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/broadcast_in_dim" stack_frame_id=551} + %min.102 = s32[1,41,41]{2,1,0} minimum(%convert_element_type.1845, %broadcast_in_dim.1622), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/min" stack_frame_id=551} + %broadcast_in_dim.1629 = s32[1,1,41,41]{3,2,1,0} reshape(%min.102), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/broadcast_in_dim" stack_frame_id=728} + %convert_element_type.1853 = pred[1,1,41,41]{3,2,1,0} compare(%broadcast_in_dim.1629, %broadcast.49), direction=NE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/convert_element_type" stack_frame_id=732} + %broadcast_in_dim.1630 = pred[1,1,1,41,41]{4,3,2,1,0} reshape(%convert_element_type.1853), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/broadcast_in_dim" stack_frame_id=34} + %and.319 = pred[1,1,1,41,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1630), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/vmap()/and" stack_frame_id=34} + %and.320 = pred[1,1,41,41]{3,2,1,0} reshape(%and.319), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/vmap()/and" stack_frame_id=34} + %and.321 = pred[1,1,32,41,41]{4,3,2,1,0} broadcast(%and.320), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/vmap()/and" stack_frame_id=34} + %state_1__65_.1 = bf16[4096,32,128]{2,1,0} parameter(68), metadata={op_name="state[1][65]"} + %dot_general.1127 = bf16[1,41,32,128]{3,2,1,0} dot(%convert_element_type.1848, %state_1__65_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=577} + %iota.571 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/iota" stack_frame_id=598} + %broadcast_in_dim.1625 = f32[1,41]{1,0} reshape(%iota.571), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/broadcast_in_dim" stack_frame_id=601} + %iota.570 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/iota" stack_frame_id=588} + %mul.3335 = f32[64]{0} multiply(%iota.570, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=588} + %div.1040 = f32[64]{0} divide(%mul.3335, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/div" stack_frame_id=589} + %neg.418 = f32[64]{0} negate(%div.1040), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/neg" stack_frame_id=590} + %pow.413 = f32[64]{0} power(%broadcast.51, %neg.418), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/pow" stack_frame_id=593} + %div.1041 = f32[64]{0} divide(%pow.413, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/div" stack_frame_id=594} + %dot_general.1128 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1625, %div.1041), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/bi,j->bij/dot_general" stack_frame_id=604} + %stack.1638 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1128), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/stack" stack_frame_id=607} + %stack.1639 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1128), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/stack" stack_frame_id=607} + %stack.1640 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1638, %stack.1639), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/stack" stack_frame_id=607} + %reshape.754 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1640), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/reshape"} + %cos.205 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.754), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/cos" stack_frame_id=611} + %convert_element_type.1849 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.205), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/convert_element_type" stack_frame_id=615} + %mul.3336 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1849), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=635} + %mul.3337 = bf16[1,41,128]{2,1,0} reshape(%mul.3336), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=635} + %mul.3338 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3337), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=635} + %mul.3339 = bf16[1,41,32,128]{3,2,1,0} multiply(%dot_general.1127, %mul.3338), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=635} + %split.411 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1127), slice={[0:1], [0:41], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/split" stack_frame_id=627} + %neg.419 = bf16[1,41,32,64]{3,2,1,0} negate(%split.411), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/neg" stack_frame_id=628} + %stack.1641 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%neg.419), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/stack" stack_frame_id=631} + %split.410 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1127), slice={[0:1], [0:41], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/split" stack_frame_id=627} + %stack.1642 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%split.410), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/stack" stack_frame_id=631} + %stack.1643 = bf16[1,41,32,2,64]{4,3,2,1,0} concatenate(%stack.1641, %stack.1642), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/stack" stack_frame_id=631} + %reshape.755 = bf16[1,41,32,128]{3,2,1,0} reshape(%stack.1643), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/reshape" stack_frame_id=634} + %sin.205 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.754), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/sin" stack_frame_id=619} + %convert_element_type.1850 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.205), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/convert_element_type" stack_frame_id=623} + %mul.3340 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1850), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=636} + %mul.3341 = bf16[1,41,128]{2,1,0} reshape(%mul.3340), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=636} + %mul.3342 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3341), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=636} + %mul.3343 = bf16[1,41,32,128]{3,2,1,0} multiply(%reshape.755, %mul.3342), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/mul" stack_frame_id=636} + %add.1174 = bf16[1,41,32,128]{3,2,1,0} add(%mul.3339, %mul.3343), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/rotary_embedding_7/add" stack_frame_id=637} + %reshape.760 = bf16[1,41,32,1,128]{4,3,2,1,0} reshape(%add.1174), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/reshape" stack_frame_id=34} + %broadcast_in_dim.1627 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%add.1175), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/broadcast_in_dim" stack_frame_id=722} + %reshape.758 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1627), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/reshape" stack_frame_id=722} + %dot_general.1132 = f32[1,32,41,1,41]{4,3,2,1,0} dot(%reshape.760, %reshape.758), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=34} + %mul.3353 = f32[1,32,41,1,41]{4,3,2,1,0} multiply(%dot_general.1132, %broadcast.48), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/vmap()/mul" stack_frame_id=34} + %vmap_jit__where__.102 = f32[1,1,32,41,41]{4,3,2,1,0} call(%and.321, %mul.3353, %constant.153), to_apply=%_where_0.5, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/vmap(jit(_where))" stack_frame_id=34} + %reduce_max.488 = f32[1,1,32,41]{3,2,1,0} reduce(%vmap_jit__where__.102, %constant.152), dimensions={4}, to_apply=%region_30.35, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/vmap()/reduce_max" stack_frame_id=34} + %max.112 = f32[1,1,32,41]{3,2,1,0} maximum(%reduce_max.488, %broadcast.47), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/vmap()/max" stack_frame_id=34} + %broadcast_in_dim.1631 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%max.112), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %sub.432 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1631), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/vmap()/sub" stack_frame_id=34} + %sub.433 = f32[1,1,32,41]{3,2,1,0} reshape(%sub.432), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/vmap()/sub" stack_frame_id=34} + %sub.434 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%sub.433), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/vmap()/sub" stack_frame_id=34} + %sub.435 = f32[1,1,32,41,41]{4,3,2,1,0} subtract(%vmap_jit__where__.102, %sub.434), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/vmap()/sub" stack_frame_id=34} + %exp.108 = f32[1,1,32,41,41]{4,3,2,1,0} exponential(%sub.435), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/vmap()/exp" stack_frame_id=34} + %reduce_sum.1471 = f32[1,1,32,41]{3,2,1,0} reduce(%exp.108, %constant.155), dimensions={4}, to_apply=%region_31.36, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/vmap()/reduce_sum" stack_frame_id=34} + %broadcast_in_dim.1632 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%reduce_sum.1471), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %div.1044 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1632), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/vmap()/div" stack_frame_id=34} + %div.1045 = f32[1,1,32,41]{3,2,1,0} reshape(%div.1044), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/vmap()/div" stack_frame_id=34} + %div.1046 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%div.1045), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/vmap()/div" stack_frame_id=34} + %div.1047 = f32[1,1,32,41,41]{4,3,2,1,0} divide(%exp.108, %div.1046), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/vmap()/div" stack_frame_id=34} + %convert_element_type.1854 = bf16[1,1,32,41,41]{4,3,2,1,0} convert(%div.1047), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/vmap()/convert_element_type" stack_frame_id=34} + %dot_general.1133 = bf16[1,32,128,1,41]{4,3,2,1,0} dot(%reshape.759, %convert_element_type.1854), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=34} + %transpose.106 = bf16[1,41,32,1,128]{1,3,4,2,0} transpose(%dot_general.1133), dimensions={0,4,1,3,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/vmap()/transpose;jit(compiled_generate_function)/transformer_layer_7/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %reshape.761 = bf16[1,41,32,128]{3,2,1,0} reshape(%transpose.106), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/reshape" stack_frame_id=34} + %state_1__68_.1 = bf16[32,128,4096]{2,1,0} parameter(71), metadata={op_name="state[1][68]"} + %dot_general.1134 = bf16[1,41,4096]{2,1,0} dot(%reshape.761, %state_1__68_.1), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=741} + %add.1176 = bf16[1,41,4096]{2,1,0} add(%dot_general.1134, %add.1171), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/add" stack_frame_id=742} + %convert_element_type.1855 = f32[1,41,4096]{2,1,0} convert(%add.1176), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/feedforward_layernorm/convert_element_type" stack_frame_id=746} + %pow.415 = f32[1,41,4096]{2,1,0} power(%convert_element_type.1855, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/feedforward_layernorm/pow" stack_frame_id=749} + %reduce_sum.1472 = f32[1,41]{1,0} reduce(%pow.415, %constant.155), dimensions={2}, to_apply=%region_32.37, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/feedforward_layernorm/reduce_sum" stack_frame_id=43} + %broadcast_in_dim.1633 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1472), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/feedforward_layernorm/broadcast_in_dim" stack_frame_id=43} + %div.1048 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1633, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/feedforward_layernorm/div" stack_frame_id=43} + %add.1177 = f32[1,41,1]{2,1,0} add(%div.1048, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/feedforward_layernorm/add" stack_frame_id=750} + %rsqrt.208 = f32[1,41,1]{2,1,0} rsqrt(%add.1177), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/feedforward_layernorm/rsqrt" stack_frame_id=753} + %mul.3354 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.208), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3355 = f32[1,41]{1,0} reshape(%mul.3354), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3356 = f32[1,41,4096]{2,1,0} broadcast(%mul.3355), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3357 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.1855, %mul.3356), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/feedforward_layernorm/mul" stack_frame_id=754} + %state_1__73_.1 = bf16[4096]{0} parameter(76), metadata={op_name="state[1][73]"} + %convert_element_type.1856 = f32[4096]{0} convert(%state_1__73_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/feedforward_layernorm/convert_element_type" stack_frame_id=755} + %broadcast_in_dim.1634 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1856), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/feedforward_layernorm/broadcast_in_dim" stack_frame_id=755} + %mul.3358 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1634), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3359 = f32[1,4096]{1,0} reshape(%mul.3358), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3360 = f32[1,41,4096]{2,1,0} broadcast(%mul.3359), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3361 = f32[1,41,4096]{2,1,0} multiply(%mul.3357, %mul.3360), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/feedforward_layernorm/mul" stack_frame_id=755} + %convert_element_type.1857 = bf16[1,41,4096]{2,1,0} convert(%mul.3361), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/feedforward_layernorm/convert_element_type" stack_frame_id=759} + %state_1__70_.1 = bf16[4096,14336]{1,0} parameter(73), metadata={op_name="state[1][70]"} + %dot_general.1136 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.1857, %state_1__70_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/feedforward_intermediate_dense/dot_general" stack_frame_id=788} + %state_1__71_.1 = bf16[4096,14336]{1,0} parameter(74), metadata={op_name="state[1][71]"} + %dot_general.1135 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.1857, %state_1__71_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/feedforward_gate_dense/dot_general" stack_frame_id=768} + %convert_element_type.1858 = f32[1,41,14336]{2,1,0} convert(%dot_general.1135), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/convert_element_type" stack_frame_id=772} + %jit_silu_.102 = f32[1,41,14336]{2,1,0} call(%convert_element_type.1858), to_apply=%silu.9, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/jit(silu)" stack_frame_id=775} + %convert_element_type.1859 = bf16[1,41,14336]{2,1,0} convert(%jit_silu_.102), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/convert_element_type" stack_frame_id=779} + %mul.3362 = bf16[1,41,14336]{2,1,0} multiply(%dot_general.1136, %convert_element_type.1859), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/mul" stack_frame_id=791} + %state_1__72_.1 = bf16[14336,4096]{1,0} parameter(75), metadata={op_name="state[1][72]"} + %dot_general.1137 = bf16[1,41,4096]{2,1,0} dot(%mul.3362, %state_1__72_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/feedforward_output_dense/dot_general" stack_frame_id=800} + %add.1178 = bf16[1,41,4096]{2,1,0} add(%dot_general.1137, %add.1176), metadata={op_name="jit(compiled_generate_function)/transformer_layer_7/add" stack_frame_id=801} + %convert_element_type.1862 = f32[1,41,4096]{2,1,0} convert(%add.1178), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention_layernorm/convert_element_type" stack_frame_id=555} + %pow.416 = f32[1,41,4096]{2,1,0} power(%convert_element_type.1862, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention_layernorm/pow" stack_frame_id=558} + %reduce_sum.1473 = f32[1,41]{1,0} reduce(%pow.416, %constant.155), dimensions={2}, to_apply=%region_33.38, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention_layernorm/reduce_sum" stack_frame_id=24} + %broadcast_in_dim.1641 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1473), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention_layernorm/broadcast_in_dim" stack_frame_id=24} + %div.1049 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1641, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention_layernorm/div" stack_frame_id=24} + %add.1180 = f32[1,41,1]{2,1,0} add(%div.1049, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention_layernorm/add" stack_frame_id=559} + %rsqrt.209 = f32[1,41,1]{2,1,0} rsqrt(%add.1180), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention_layernorm/rsqrt" stack_frame_id=562} + %mul.3363 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.209), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3364 = f32[1,41]{1,0} reshape(%mul.3363), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3365 = f32[1,41,4096]{2,1,0} broadcast(%mul.3364), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3366 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.1862, %mul.3365), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention_layernorm/mul" stack_frame_id=563} + %state_1__78_.1 = bf16[4096]{0} parameter(81), metadata={op_name="state[1][78]"} + %convert_element_type.1863 = f32[4096]{0} convert(%state_1__78_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention_layernorm/convert_element_type" stack_frame_id=564} + %broadcast_in_dim.1642 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1863), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention_layernorm/broadcast_in_dim" stack_frame_id=564} + %mul.3367 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1642), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3368 = f32[1,4096]{1,0} reshape(%mul.3367), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3369 = f32[1,41,4096]{2,1,0} broadcast(%mul.3368), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3370 = f32[1,41,4096]{2,1,0} multiply(%mul.3366, %mul.3369), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention_layernorm/mul" stack_frame_id=564} + %convert_element_type.1864 = bf16[1,41,4096]{2,1,0} convert(%mul.3370), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention_layernorm/convert_element_type" stack_frame_id=568} + %state_1__75_.1 = bf16[4096,8,128]{2,1,0} parameter(78), metadata={op_name="state[1][75]"} + %dot_general.1140 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.1864, %state_1__75_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=647} + %iota.581 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/iota" stack_frame_id=677} + %broadcast_in_dim.1644 = f32[1,41]{1,0} reshape(%iota.581), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/broadcast_in_dim" stack_frame_id=680} + %iota.580 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/iota" stack_frame_id=667} + %mul.3380 = f32[64]{0} multiply(%iota.580, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=667} + %div.1052 = f32[64]{0} divide(%mul.3380, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/div" stack_frame_id=668} + %neg.424 = f32[64]{0} negate(%div.1052), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/neg" stack_frame_id=669} + %pow.418 = f32[64]{0} power(%broadcast.51, %neg.424), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/pow" stack_frame_id=672} + %div.1053 = f32[64]{0} divide(%pow.418, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/div" stack_frame_id=673} + %dot_general.1142 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1644, %div.1053), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/bi,j->bij/dot_general" stack_frame_id=683} + %stack.1659 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1142), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/stack" stack_frame_id=686} + %stack.1660 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1142), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/stack" stack_frame_id=686} + %stack.1661 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1659, %stack.1660), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/stack" stack_frame_id=686} + %reshape.764 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1661), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/reshape"} + %cos.208 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.764), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/cos" stack_frame_id=690} + %convert_element_type.1867 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.208), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/convert_element_type" stack_frame_id=694} + %mul.3381 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1867), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=714} + %mul.3382 = bf16[1,41,128]{2,1,0} reshape(%mul.3381), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=714} + %mul.3383 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3382), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=714} + %mul.3384 = bf16[1,41,8,128]{3,2,1,0} multiply(%dot_general.1140, %mul.3383), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=714} + %split.417 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1140), slice={[0:1], [0:41], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/split" stack_frame_id=706} + %neg.425 = bf16[1,41,8,64]{3,2,1,0} negate(%split.417), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/neg" stack_frame_id=707} + %stack.1662 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%neg.425), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/stack" stack_frame_id=710} + %split.416 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1140), slice={[0:1], [0:41], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/split" stack_frame_id=706} + %stack.1663 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%split.416), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/stack" stack_frame_id=710} + %stack.1664 = bf16[1,41,8,2,64]{4,3,2,1,0} concatenate(%stack.1662, %stack.1663), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/stack" stack_frame_id=710} + %reshape.765 = bf16[1,41,8,128]{3,2,1,0} reshape(%stack.1664), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/reshape" stack_frame_id=713} + %sin.208 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.764), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/sin" stack_frame_id=698} + %convert_element_type.1868 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.208), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/convert_element_type" stack_frame_id=702} + %mul.3385 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1868), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=715} + %mul.3386 = bf16[1,41,128]{2,1,0} reshape(%mul.3385), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=715} + %mul.3387 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3386), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=715} + %mul.3388 = bf16[1,41,8,128]{3,2,1,0} multiply(%reshape.765, %mul.3387), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=715} + %add.1182 = bf16[1,41,8,128]{3,2,1,0} add(%mul.3384, %mul.3388), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/add" stack_frame_id=716} + %stack.1665 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%add.1182), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/stack" stack_frame_id=719} + %state_1__76_.1 = bf16[4096,8,128]{2,1,0} parameter(79), metadata={op_name="state[1][76]"} + %dot_general.1141 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.1864, %state_1__76_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=656} + %stack.1666 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dot_general.1141), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/stack" stack_frame_id=719} + %stack.1667 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1665, %stack.1666), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/stack" stack_frame_id=719} + %stack.2015 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1667), metadata={op_name="jit(compiled_generate_function)/stack" stack_frame_id=804} + %broadcast_in_dim.1646 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dot_general.1141), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/broadcast_in_dim" stack_frame_id=725} + %reshape.767 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1646), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/reshape" stack_frame_id=725} + %iota.574 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/iota" stack_frame_id=525} + %broadcast_in_dim.1635 = f32[41,1]{1,0} reshape(%iota.574), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/broadcast_in_dim" stack_frame_id=528} + %ge.537 = f32[41,1]{1,0} broadcast(%broadcast_in_dim.1635), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/ge" stack_frame_id=532} + %ge.538 = f32[41]{0} reshape(%ge.537), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/ge" stack_frame_id=532} + %ge.539 = f32[41,41]{1,0} broadcast(%ge.538), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/ge" stack_frame_id=532} + %iota.575 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/iota" stack_frame_id=531} + %broadcast_in_dim.1636 = f32[1,41]{1,0} reshape(%iota.575), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/broadcast_in_dim" stack_frame_id=532} + %ge.540 = f32[1,41]{1,0} broadcast(%broadcast_in_dim.1636), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/ge" stack_frame_id=532} + %ge.541 = f32[41]{0} reshape(%ge.540), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/ge" stack_frame_id=532} + %ge.542 = f32[41,41]{1,0} broadcast(%ge.541), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/ge" stack_frame_id=532} + %ge.543 = pred[41,41]{1,0} compare(%ge.539, %ge.542), direction=GE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/ge" stack_frame_id=532} + %broadcast_in_dim.1637 = pred[1,41,41]{2,1,0} reshape(%ge.543), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/broadcast_in_dim" stack_frame_id=535} + %convert_element_type.1861 = s32[1,41,41]{2,1,0} convert(%broadcast_in_dim.1637), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/convert_element_type" stack_frame_id=551} + %iota.576 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/iota" stack_frame_id=538} + %broadcast_in_dim.1638 = s32[41,1]{1,0} reshape(%iota.576), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/broadcast_in_dim" stack_frame_id=536} + %lt.679 = s32[41,1]{1,0} broadcast(%broadcast_in_dim.1638), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/lt" stack_frame_id=543} + %lt.680 = s32[41]{0} reshape(%lt.679), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/lt" stack_frame_id=543} + %lt.681 = s32[41,41]{1,0} broadcast(%lt.680), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/lt" stack_frame_id=543} + %iota.577 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/iota" stack_frame_id=541} + %broadcast_in_dim.1639 = s32[1,41]{1,0} reshape(%iota.577), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/broadcast_in_dim" stack_frame_id=539} + %add.1179 = s32[1,41]{1,0} add(%broadcast_in_dim.1639, %broadcast.57), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/add" stack_frame_id=542} + %lt.682 = s32[1,41]{1,0} broadcast(%add.1179), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/lt" stack_frame_id=543} + %lt.683 = s32[41]{0} reshape(%lt.682), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/lt" stack_frame_id=543} + %lt.684 = s32[41,41]{1,0} broadcast(%lt.683), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/lt" stack_frame_id=543} + %lt.685 = pred[41,41]{1,0} compare(%lt.681, %lt.684), direction=LT, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/lt" stack_frame_id=543} + %convert_element_type.1860 = s32[41,41]{1,0} convert(%lt.685), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/convert_element_type" stack_frame_id=547} + %broadcast_in_dim.1640 = s32[1,41,41]{2,1,0} reshape(%convert_element_type.1860), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/broadcast_in_dim" stack_frame_id=551} + %min.103 = s32[1,41,41]{2,1,0} minimum(%convert_element_type.1861, %broadcast_in_dim.1640), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/min" stack_frame_id=551} + %broadcast_in_dim.1647 = s32[1,1,41,41]{3,2,1,0} reshape(%min.103), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/broadcast_in_dim" stack_frame_id=728} + %convert_element_type.1869 = pred[1,1,41,41]{3,2,1,0} compare(%broadcast_in_dim.1647, %broadcast.49), direction=NE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/convert_element_type" stack_frame_id=732} + %broadcast_in_dim.1648 = pred[1,1,1,41,41]{4,3,2,1,0} reshape(%convert_element_type.1869), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/broadcast_in_dim" stack_frame_id=34} + %and.322 = pred[1,1,1,41,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1648), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/vmap()/and" stack_frame_id=34} + %and.323 = pred[1,1,41,41]{3,2,1,0} reshape(%and.322), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/vmap()/and" stack_frame_id=34} + %and.324 = pred[1,1,32,41,41]{4,3,2,1,0} broadcast(%and.323), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/vmap()/and" stack_frame_id=34} + %state_1__74_.1 = bf16[4096,32,128]{2,1,0} parameter(77), metadata={op_name="state[1][74]"} + %dot_general.1138 = bf16[1,41,32,128]{3,2,1,0} dot(%convert_element_type.1864, %state_1__74_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=577} + %iota.579 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/iota" stack_frame_id=598} + %broadcast_in_dim.1643 = f32[1,41]{1,0} reshape(%iota.579), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/broadcast_in_dim" stack_frame_id=601} + %iota.578 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/iota" stack_frame_id=588} + %mul.3371 = f32[64]{0} multiply(%iota.578, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=588} + %div.1050 = f32[64]{0} divide(%mul.3371, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/div" stack_frame_id=589} + %neg.422 = f32[64]{0} negate(%div.1050), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/neg" stack_frame_id=590} + %pow.417 = f32[64]{0} power(%broadcast.51, %neg.422), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/pow" stack_frame_id=593} + %div.1051 = f32[64]{0} divide(%pow.417, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/div" stack_frame_id=594} + %dot_general.1139 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1643, %div.1051), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/bi,j->bij/dot_general" stack_frame_id=604} + %stack.1653 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1139), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/stack" stack_frame_id=607} + %stack.1654 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1139), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/stack" stack_frame_id=607} + %stack.1655 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1653, %stack.1654), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/stack" stack_frame_id=607} + %reshape.762 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1655), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/reshape"} + %cos.207 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.762), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/cos" stack_frame_id=611} + %convert_element_type.1865 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.207), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/convert_element_type" stack_frame_id=615} + %mul.3372 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1865), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=635} + %mul.3373 = bf16[1,41,128]{2,1,0} reshape(%mul.3372), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=635} + %mul.3374 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3373), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=635} + %mul.3375 = bf16[1,41,32,128]{3,2,1,0} multiply(%dot_general.1138, %mul.3374), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=635} + %split.415 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1138), slice={[0:1], [0:41], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/split" stack_frame_id=627} + %neg.423 = bf16[1,41,32,64]{3,2,1,0} negate(%split.415), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/neg" stack_frame_id=628} + %stack.1656 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%neg.423), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/stack" stack_frame_id=631} + %split.414 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1138), slice={[0:1], [0:41], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/split" stack_frame_id=627} + %stack.1657 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%split.414), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/stack" stack_frame_id=631} + %stack.1658 = bf16[1,41,32,2,64]{4,3,2,1,0} concatenate(%stack.1656, %stack.1657), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/stack" stack_frame_id=631} + %reshape.763 = bf16[1,41,32,128]{3,2,1,0} reshape(%stack.1658), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/reshape" stack_frame_id=634} + %sin.207 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.762), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/sin" stack_frame_id=619} + %convert_element_type.1866 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.207), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/convert_element_type" stack_frame_id=623} + %mul.3376 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1866), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=636} + %mul.3377 = bf16[1,41,128]{2,1,0} reshape(%mul.3376), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=636} + %mul.3378 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3377), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=636} + %mul.3379 = bf16[1,41,32,128]{3,2,1,0} multiply(%reshape.763, %mul.3378), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/mul" stack_frame_id=636} + %add.1181 = bf16[1,41,32,128]{3,2,1,0} add(%mul.3375, %mul.3379), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/rotary_embedding_8/add" stack_frame_id=637} + %reshape.768 = bf16[1,41,32,1,128]{4,3,2,1,0} reshape(%add.1181), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/reshape" stack_frame_id=34} + %broadcast_in_dim.1645 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%add.1182), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/broadcast_in_dim" stack_frame_id=722} + %reshape.766 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1645), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/reshape" stack_frame_id=722} + %dot_general.1143 = f32[1,32,41,1,41]{4,3,2,1,0} dot(%reshape.768, %reshape.766), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=34} + %mul.3389 = f32[1,32,41,1,41]{4,3,2,1,0} multiply(%dot_general.1143, %broadcast.48), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/vmap()/mul" stack_frame_id=34} + %vmap_jit__where__.103 = f32[1,1,32,41,41]{4,3,2,1,0} call(%and.324, %mul.3389, %constant.153), to_apply=%_where_0.5, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/vmap(jit(_where))" stack_frame_id=34} + %reduce_max.489 = f32[1,1,32,41]{3,2,1,0} reduce(%vmap_jit__where__.103, %constant.152), dimensions={4}, to_apply=%region_34.39, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/vmap()/reduce_max" stack_frame_id=34} + %max.113 = f32[1,1,32,41]{3,2,1,0} maximum(%reduce_max.489, %broadcast.47), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/vmap()/max" stack_frame_id=34} + %broadcast_in_dim.1649 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%max.113), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %sub.436 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1649), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/vmap()/sub" stack_frame_id=34} + %sub.437 = f32[1,1,32,41]{3,2,1,0} reshape(%sub.436), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/vmap()/sub" stack_frame_id=34} + %sub.438 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%sub.437), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/vmap()/sub" stack_frame_id=34} + %sub.439 = f32[1,1,32,41,41]{4,3,2,1,0} subtract(%vmap_jit__where__.103, %sub.438), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/vmap()/sub" stack_frame_id=34} + %exp.109 = f32[1,1,32,41,41]{4,3,2,1,0} exponential(%sub.439), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/vmap()/exp" stack_frame_id=34} + %reduce_sum.1474 = f32[1,1,32,41]{3,2,1,0} reduce(%exp.109, %constant.155), dimensions={4}, to_apply=%region_35.40, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/vmap()/reduce_sum" stack_frame_id=34} + %broadcast_in_dim.1650 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%reduce_sum.1474), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %div.1054 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1650), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/vmap()/div" stack_frame_id=34} + %div.1055 = f32[1,1,32,41]{3,2,1,0} reshape(%div.1054), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/vmap()/div" stack_frame_id=34} + %div.1056 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%div.1055), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/vmap()/div" stack_frame_id=34} + %div.1057 = f32[1,1,32,41,41]{4,3,2,1,0} divide(%exp.109, %div.1056), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/vmap()/div" stack_frame_id=34} + %convert_element_type.1870 = bf16[1,1,32,41,41]{4,3,2,1,0} convert(%div.1057), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/vmap()/convert_element_type" stack_frame_id=34} + %dot_general.1144 = bf16[1,32,128,1,41]{4,3,2,1,0} dot(%reshape.767, %convert_element_type.1870), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=34} + %transpose.107 = bf16[1,41,32,1,128]{1,3,4,2,0} transpose(%dot_general.1144), dimensions={0,4,1,3,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/vmap()/transpose;jit(compiled_generate_function)/transformer_layer_8/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %reshape.769 = bf16[1,41,32,128]{3,2,1,0} reshape(%transpose.107), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/reshape" stack_frame_id=34} + %state_1__77_.1 = bf16[32,128,4096]{2,1,0} parameter(80), metadata={op_name="state[1][77]"} + %dot_general.1145 = bf16[1,41,4096]{2,1,0} dot(%reshape.769, %state_1__77_.1), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=741} + %add.1183 = bf16[1,41,4096]{2,1,0} add(%dot_general.1145, %add.1178), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/add" stack_frame_id=742} + %convert_element_type.1871 = f32[1,41,4096]{2,1,0} convert(%add.1183), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/feedforward_layernorm/convert_element_type" stack_frame_id=746} + %pow.419 = f32[1,41,4096]{2,1,0} power(%convert_element_type.1871, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/feedforward_layernorm/pow" stack_frame_id=749} + %reduce_sum.1475 = f32[1,41]{1,0} reduce(%pow.419, %constant.155), dimensions={2}, to_apply=%region_36.41, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/feedforward_layernorm/reduce_sum" stack_frame_id=43} + %broadcast_in_dim.1651 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1475), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/feedforward_layernorm/broadcast_in_dim" stack_frame_id=43} + %div.1058 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1651, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/feedforward_layernorm/div" stack_frame_id=43} + %add.1184 = f32[1,41,1]{2,1,0} add(%div.1058, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/feedforward_layernorm/add" stack_frame_id=750} + %rsqrt.210 = f32[1,41,1]{2,1,0} rsqrt(%add.1184), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/feedforward_layernorm/rsqrt" stack_frame_id=753} + %mul.3390 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.210), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3391 = f32[1,41]{1,0} reshape(%mul.3390), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3392 = f32[1,41,4096]{2,1,0} broadcast(%mul.3391), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3393 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.1871, %mul.3392), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/feedforward_layernorm/mul" stack_frame_id=754} + %state_1__82_.1 = bf16[4096]{0} parameter(85), metadata={op_name="state[1][82]"} + %convert_element_type.1872 = f32[4096]{0} convert(%state_1__82_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/feedforward_layernorm/convert_element_type" stack_frame_id=755} + %broadcast_in_dim.1652 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1872), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/feedforward_layernorm/broadcast_in_dim" stack_frame_id=755} + %mul.3394 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1652), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3395 = f32[1,4096]{1,0} reshape(%mul.3394), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3396 = f32[1,41,4096]{2,1,0} broadcast(%mul.3395), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3397 = f32[1,41,4096]{2,1,0} multiply(%mul.3393, %mul.3396), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/feedforward_layernorm/mul" stack_frame_id=755} + %convert_element_type.1873 = bf16[1,41,4096]{2,1,0} convert(%mul.3397), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/feedforward_layernorm/convert_element_type" stack_frame_id=759} + %state_1__79_.1 = bf16[4096,14336]{1,0} parameter(82), metadata={op_name="state[1][79]"} + %dot_general.1147 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.1873, %state_1__79_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/feedforward_intermediate_dense/dot_general" stack_frame_id=788} + %state_1__80_.1 = bf16[4096,14336]{1,0} parameter(83), metadata={op_name="state[1][80]"} + %dot_general.1146 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.1873, %state_1__80_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/feedforward_gate_dense/dot_general" stack_frame_id=768} + %convert_element_type.1874 = f32[1,41,14336]{2,1,0} convert(%dot_general.1146), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/convert_element_type" stack_frame_id=772} + %jit_silu_.103 = f32[1,41,14336]{2,1,0} call(%convert_element_type.1874), to_apply=%silu.9, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/jit(silu)" stack_frame_id=775} + %convert_element_type.1875 = bf16[1,41,14336]{2,1,0} convert(%jit_silu_.103), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/convert_element_type" stack_frame_id=779} + %mul.3398 = bf16[1,41,14336]{2,1,0} multiply(%dot_general.1147, %convert_element_type.1875), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/mul" stack_frame_id=791} + %state_1__81_.1 = bf16[14336,4096]{1,0} parameter(84), metadata={op_name="state[1][81]"} + %dot_general.1148 = bf16[1,41,4096]{2,1,0} dot(%mul.3398, %state_1__81_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/feedforward_output_dense/dot_general" stack_frame_id=800} + %add.1185 = bf16[1,41,4096]{2,1,0} add(%dot_general.1148, %add.1183), metadata={op_name="jit(compiled_generate_function)/transformer_layer_8/add" stack_frame_id=801} + %convert_element_type.1878 = f32[1,41,4096]{2,1,0} convert(%add.1185), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention_layernorm/convert_element_type" stack_frame_id=555} + %pow.420 = f32[1,41,4096]{2,1,0} power(%convert_element_type.1878, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention_layernorm/pow" stack_frame_id=558} + %reduce_sum.1476 = f32[1,41]{1,0} reduce(%pow.420, %constant.155), dimensions={2}, to_apply=%region_37.42, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention_layernorm/reduce_sum" stack_frame_id=24} + %broadcast_in_dim.1659 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1476), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention_layernorm/broadcast_in_dim" stack_frame_id=24} + %div.1059 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1659, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention_layernorm/div" stack_frame_id=24} + %add.1187 = f32[1,41,1]{2,1,0} add(%div.1059, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention_layernorm/add" stack_frame_id=559} + %rsqrt.211 = f32[1,41,1]{2,1,0} rsqrt(%add.1187), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention_layernorm/rsqrt" stack_frame_id=562} + %mul.3399 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.211), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3400 = f32[1,41]{1,0} reshape(%mul.3399), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3401 = f32[1,41,4096]{2,1,0} broadcast(%mul.3400), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3402 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.1878, %mul.3401), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention_layernorm/mul" stack_frame_id=563} + %state_1__87_.1 = bf16[4096]{0} parameter(90), metadata={op_name="state[1][87]"} + %convert_element_type.1879 = f32[4096]{0} convert(%state_1__87_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention_layernorm/convert_element_type" stack_frame_id=564} + %broadcast_in_dim.1660 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1879), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention_layernorm/broadcast_in_dim" stack_frame_id=564} + %mul.3403 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1660), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3404 = f32[1,4096]{1,0} reshape(%mul.3403), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3405 = f32[1,41,4096]{2,1,0} broadcast(%mul.3404), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3406 = f32[1,41,4096]{2,1,0} multiply(%mul.3402, %mul.3405), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention_layernorm/mul" stack_frame_id=564} + %convert_element_type.1880 = bf16[1,41,4096]{2,1,0} convert(%mul.3406), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention_layernorm/convert_element_type" stack_frame_id=568} + %state_1__84_.1 = bf16[4096,8,128]{2,1,0} parameter(87), metadata={op_name="state[1][84]"} + %dot_general.1151 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.1880, %state_1__84_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=647} + %iota.589 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/iota" stack_frame_id=677} + %broadcast_in_dim.1662 = f32[1,41]{1,0} reshape(%iota.589), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/broadcast_in_dim" stack_frame_id=680} + %iota.588 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/iota" stack_frame_id=667} + %mul.3416 = f32[64]{0} multiply(%iota.588, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=667} + %div.1062 = f32[64]{0} divide(%mul.3416, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/div" stack_frame_id=668} + %neg.428 = f32[64]{0} negate(%div.1062), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/neg" stack_frame_id=669} + %pow.422 = f32[64]{0} power(%broadcast.51, %neg.428), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/pow" stack_frame_id=672} + %div.1063 = f32[64]{0} divide(%pow.422, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/div" stack_frame_id=673} + %dot_general.1153 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1662, %div.1063), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/bi,j->bij/dot_general" stack_frame_id=683} + %stack.1674 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1153), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/stack" stack_frame_id=686} + %stack.1675 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1153), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/stack" stack_frame_id=686} + %stack.1676 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1674, %stack.1675), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/stack" stack_frame_id=686} + %reshape.772 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1676), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/reshape"} + %cos.210 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.772), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/cos" stack_frame_id=690} + %convert_element_type.1883 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.210), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/convert_element_type" stack_frame_id=694} + %mul.3417 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1883), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=714} + %mul.3418 = bf16[1,41,128]{2,1,0} reshape(%mul.3417), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=714} + %mul.3419 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3418), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=714} + %mul.3420 = bf16[1,41,8,128]{3,2,1,0} multiply(%dot_general.1151, %mul.3419), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=714} + %split.421 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1151), slice={[0:1], [0:41], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/split" stack_frame_id=706} + %neg.429 = bf16[1,41,8,64]{3,2,1,0} negate(%split.421), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/neg" stack_frame_id=707} + %stack.1677 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%neg.429), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/stack" stack_frame_id=710} + %split.420 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1151), slice={[0:1], [0:41], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/split" stack_frame_id=706} + %stack.1678 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%split.420), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/stack" stack_frame_id=710} + %stack.1679 = bf16[1,41,8,2,64]{4,3,2,1,0} concatenate(%stack.1677, %stack.1678), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/stack" stack_frame_id=710} + %reshape.773 = bf16[1,41,8,128]{3,2,1,0} reshape(%stack.1679), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/reshape" stack_frame_id=713} + %sin.210 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.772), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/sin" stack_frame_id=698} + %convert_element_type.1884 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.210), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/convert_element_type" stack_frame_id=702} + %mul.3421 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1884), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=715} + %mul.3422 = bf16[1,41,128]{2,1,0} reshape(%mul.3421), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=715} + %mul.3423 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3422), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=715} + %mul.3424 = bf16[1,41,8,128]{3,2,1,0} multiply(%reshape.773, %mul.3423), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=715} + %add.1189 = bf16[1,41,8,128]{3,2,1,0} add(%mul.3420, %mul.3424), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/add" stack_frame_id=716} + %stack.1680 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%add.1189), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/stack" stack_frame_id=719} + %state_1__85_.1 = bf16[4096,8,128]{2,1,0} parameter(88), metadata={op_name="state[1][85]"} + %dot_general.1152 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.1880, %state_1__85_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=656} + %stack.1681 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dot_general.1152), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/stack" stack_frame_id=719} + %stack.1682 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1680, %stack.1681), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/stack" stack_frame_id=719} + %stack.2016 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1682), metadata={op_name="jit(compiled_generate_function)/stack" stack_frame_id=804} + %broadcast_in_dim.1664 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dot_general.1152), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/broadcast_in_dim" stack_frame_id=725} + %reshape.775 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1664), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/reshape" stack_frame_id=725} + %iota.582 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/iota" stack_frame_id=525} + %broadcast_in_dim.1653 = f32[41,1]{1,0} reshape(%iota.582), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/broadcast_in_dim" stack_frame_id=528} + %ge.544 = f32[41,1]{1,0} broadcast(%broadcast_in_dim.1653), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/ge" stack_frame_id=532} + %ge.545 = f32[41]{0} reshape(%ge.544), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/ge" stack_frame_id=532} + %ge.546 = f32[41,41]{1,0} broadcast(%ge.545), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/ge" stack_frame_id=532} + %iota.583 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/iota" stack_frame_id=531} + %broadcast_in_dim.1654 = f32[1,41]{1,0} reshape(%iota.583), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/broadcast_in_dim" stack_frame_id=532} + %ge.547 = f32[1,41]{1,0} broadcast(%broadcast_in_dim.1654), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/ge" stack_frame_id=532} + %ge.548 = f32[41]{0} reshape(%ge.547), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/ge" stack_frame_id=532} + %ge.549 = f32[41,41]{1,0} broadcast(%ge.548), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/ge" stack_frame_id=532} + %ge.550 = pred[41,41]{1,0} compare(%ge.546, %ge.549), direction=GE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/ge" stack_frame_id=532} + %broadcast_in_dim.1655 = pred[1,41,41]{2,1,0} reshape(%ge.550), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/broadcast_in_dim" stack_frame_id=535} + %convert_element_type.1877 = s32[1,41,41]{2,1,0} convert(%broadcast_in_dim.1655), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/convert_element_type" stack_frame_id=551} + %iota.584 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/iota" stack_frame_id=538} + %broadcast_in_dim.1656 = s32[41,1]{1,0} reshape(%iota.584), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/broadcast_in_dim" stack_frame_id=536} + %lt.686 = s32[41,1]{1,0} broadcast(%broadcast_in_dim.1656), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/lt" stack_frame_id=543} + %lt.687 = s32[41]{0} reshape(%lt.686), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/lt" stack_frame_id=543} + %lt.688 = s32[41,41]{1,0} broadcast(%lt.687), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/lt" stack_frame_id=543} + %iota.585 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/iota" stack_frame_id=541} + %broadcast_in_dim.1657 = s32[1,41]{1,0} reshape(%iota.585), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/broadcast_in_dim" stack_frame_id=539} + %add.1186 = s32[1,41]{1,0} add(%broadcast_in_dim.1657, %broadcast.57), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/add" stack_frame_id=542} + %lt.689 = s32[1,41]{1,0} broadcast(%add.1186), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/lt" stack_frame_id=543} + %lt.690 = s32[41]{0} reshape(%lt.689), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/lt" stack_frame_id=543} + %lt.691 = s32[41,41]{1,0} broadcast(%lt.690), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/lt" stack_frame_id=543} + %lt.692 = pred[41,41]{1,0} compare(%lt.688, %lt.691), direction=LT, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/lt" stack_frame_id=543} + %convert_element_type.1876 = s32[41,41]{1,0} convert(%lt.692), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/convert_element_type" stack_frame_id=547} + %broadcast_in_dim.1658 = s32[1,41,41]{2,1,0} reshape(%convert_element_type.1876), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/broadcast_in_dim" stack_frame_id=551} + %min.104 = s32[1,41,41]{2,1,0} minimum(%convert_element_type.1877, %broadcast_in_dim.1658), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/min" stack_frame_id=551} + %broadcast_in_dim.1665 = s32[1,1,41,41]{3,2,1,0} reshape(%min.104), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/broadcast_in_dim" stack_frame_id=728} + %convert_element_type.1885 = pred[1,1,41,41]{3,2,1,0} compare(%broadcast_in_dim.1665, %broadcast.49), direction=NE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/convert_element_type" stack_frame_id=732} + %broadcast_in_dim.1666 = pred[1,1,1,41,41]{4,3,2,1,0} reshape(%convert_element_type.1885), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/broadcast_in_dim" stack_frame_id=34} + %and.325 = pred[1,1,1,41,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1666), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/vmap()/and" stack_frame_id=34} + %and.326 = pred[1,1,41,41]{3,2,1,0} reshape(%and.325), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/vmap()/and" stack_frame_id=34} + %and.327 = pred[1,1,32,41,41]{4,3,2,1,0} broadcast(%and.326), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/vmap()/and" stack_frame_id=34} + %state_1__83_.1 = bf16[4096,32,128]{2,1,0} parameter(86), metadata={op_name="state[1][83]"} + %dot_general.1149 = bf16[1,41,32,128]{3,2,1,0} dot(%convert_element_type.1880, %state_1__83_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=577} + %iota.587 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/iota" stack_frame_id=598} + %broadcast_in_dim.1661 = f32[1,41]{1,0} reshape(%iota.587), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/broadcast_in_dim" stack_frame_id=601} + %iota.586 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/iota" stack_frame_id=588} + %mul.3407 = f32[64]{0} multiply(%iota.586, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=588} + %div.1060 = f32[64]{0} divide(%mul.3407, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/div" stack_frame_id=589} + %neg.426 = f32[64]{0} negate(%div.1060), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/neg" stack_frame_id=590} + %pow.421 = f32[64]{0} power(%broadcast.51, %neg.426), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/pow" stack_frame_id=593} + %div.1061 = f32[64]{0} divide(%pow.421, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/div" stack_frame_id=594} + %dot_general.1150 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1661, %div.1061), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/bi,j->bij/dot_general" stack_frame_id=604} + %stack.1668 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1150), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/stack" stack_frame_id=607} + %stack.1669 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1150), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/stack" stack_frame_id=607} + %stack.1670 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1668, %stack.1669), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/stack" stack_frame_id=607} + %reshape.770 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1670), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/reshape"} + %cos.209 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.770), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/cos" stack_frame_id=611} + %convert_element_type.1881 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.209), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/convert_element_type" stack_frame_id=615} + %mul.3408 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1881), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=635} + %mul.3409 = bf16[1,41,128]{2,1,0} reshape(%mul.3408), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=635} + %mul.3410 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3409), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=635} + %mul.3411 = bf16[1,41,32,128]{3,2,1,0} multiply(%dot_general.1149, %mul.3410), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=635} + %split.419 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1149), slice={[0:1], [0:41], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/split" stack_frame_id=627} + %neg.427 = bf16[1,41,32,64]{3,2,1,0} negate(%split.419), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/neg" stack_frame_id=628} + %stack.1671 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%neg.427), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/stack" stack_frame_id=631} + %split.418 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1149), slice={[0:1], [0:41], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/split" stack_frame_id=627} + %stack.1672 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%split.418), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/stack" stack_frame_id=631} + %stack.1673 = bf16[1,41,32,2,64]{4,3,2,1,0} concatenate(%stack.1671, %stack.1672), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/stack" stack_frame_id=631} + %reshape.771 = bf16[1,41,32,128]{3,2,1,0} reshape(%stack.1673), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/reshape" stack_frame_id=634} + %sin.209 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.770), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/sin" stack_frame_id=619} + %convert_element_type.1882 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.209), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/convert_element_type" stack_frame_id=623} + %mul.3412 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1882), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=636} + %mul.3413 = bf16[1,41,128]{2,1,0} reshape(%mul.3412), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=636} + %mul.3414 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3413), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=636} + %mul.3415 = bf16[1,41,32,128]{3,2,1,0} multiply(%reshape.771, %mul.3414), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/mul" stack_frame_id=636} + %add.1188 = bf16[1,41,32,128]{3,2,1,0} add(%mul.3411, %mul.3415), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/rotary_embedding_9/add" stack_frame_id=637} + %reshape.776 = bf16[1,41,32,1,128]{4,3,2,1,0} reshape(%add.1188), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/reshape" stack_frame_id=34} + %broadcast_in_dim.1663 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%add.1189), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/broadcast_in_dim" stack_frame_id=722} + %reshape.774 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1663), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/reshape" stack_frame_id=722} + %dot_general.1154 = f32[1,32,41,1,41]{4,3,2,1,0} dot(%reshape.776, %reshape.774), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=34} + %mul.3425 = f32[1,32,41,1,41]{4,3,2,1,0} multiply(%dot_general.1154, %broadcast.48), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/vmap()/mul" stack_frame_id=34} + %vmap_jit__where__.104 = f32[1,1,32,41,41]{4,3,2,1,0} call(%and.327, %mul.3425, %constant.153), to_apply=%_where_0.5, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/vmap(jit(_where))" stack_frame_id=34} + %reduce_max.490 = f32[1,1,32,41]{3,2,1,0} reduce(%vmap_jit__where__.104, %constant.152), dimensions={4}, to_apply=%region_38.43, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/vmap()/reduce_max" stack_frame_id=34} + %max.114 = f32[1,1,32,41]{3,2,1,0} maximum(%reduce_max.490, %broadcast.47), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/vmap()/max" stack_frame_id=34} + %broadcast_in_dim.1667 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%max.114), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %sub.440 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1667), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/vmap()/sub" stack_frame_id=34} + %sub.441 = f32[1,1,32,41]{3,2,1,0} reshape(%sub.440), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/vmap()/sub" stack_frame_id=34} + %sub.442 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%sub.441), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/vmap()/sub" stack_frame_id=34} + %sub.443 = f32[1,1,32,41,41]{4,3,2,1,0} subtract(%vmap_jit__where__.104, %sub.442), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/vmap()/sub" stack_frame_id=34} + %exp.110 = f32[1,1,32,41,41]{4,3,2,1,0} exponential(%sub.443), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/vmap()/exp" stack_frame_id=34} + %reduce_sum.1477 = f32[1,1,32,41]{3,2,1,0} reduce(%exp.110, %constant.155), dimensions={4}, to_apply=%region_39.44, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/vmap()/reduce_sum" stack_frame_id=34} + %broadcast_in_dim.1668 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%reduce_sum.1477), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %div.1064 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1668), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/vmap()/div" stack_frame_id=34} + %div.1065 = f32[1,1,32,41]{3,2,1,0} reshape(%div.1064), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/vmap()/div" stack_frame_id=34} + %div.1066 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%div.1065), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/vmap()/div" stack_frame_id=34} + %div.1067 = f32[1,1,32,41,41]{4,3,2,1,0} divide(%exp.110, %div.1066), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/vmap()/div" stack_frame_id=34} + %convert_element_type.1886 = bf16[1,1,32,41,41]{4,3,2,1,0} convert(%div.1067), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/vmap()/convert_element_type" stack_frame_id=34} + %dot_general.1155 = bf16[1,32,128,1,41]{4,3,2,1,0} dot(%reshape.775, %convert_element_type.1886), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=34} + %transpose.108 = bf16[1,41,32,1,128]{1,3,4,2,0} transpose(%dot_general.1155), dimensions={0,4,1,3,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/vmap()/transpose;jit(compiled_generate_function)/transformer_layer_9/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %reshape.777 = bf16[1,41,32,128]{3,2,1,0} reshape(%transpose.108), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/reshape" stack_frame_id=34} + %state_1__86_.1 = bf16[32,128,4096]{2,1,0} parameter(89), metadata={op_name="state[1][86]"} + %dot_general.1156 = bf16[1,41,4096]{2,1,0} dot(%reshape.777, %state_1__86_.1), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=741} + %add.1190 = bf16[1,41,4096]{2,1,0} add(%dot_general.1156, %add.1185), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/add" stack_frame_id=742} + %convert_element_type.1887 = f32[1,41,4096]{2,1,0} convert(%add.1190), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/feedforward_layernorm/convert_element_type" stack_frame_id=746} + %pow.423 = f32[1,41,4096]{2,1,0} power(%convert_element_type.1887, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/feedforward_layernorm/pow" stack_frame_id=749} + %reduce_sum.1478 = f32[1,41]{1,0} reduce(%pow.423, %constant.155), dimensions={2}, to_apply=%region_40.45, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/feedforward_layernorm/reduce_sum" stack_frame_id=43} + %broadcast_in_dim.1669 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1478), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/feedforward_layernorm/broadcast_in_dim" stack_frame_id=43} + %div.1068 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1669, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/feedforward_layernorm/div" stack_frame_id=43} + %add.1191 = f32[1,41,1]{2,1,0} add(%div.1068, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/feedforward_layernorm/add" stack_frame_id=750} + %rsqrt.212 = f32[1,41,1]{2,1,0} rsqrt(%add.1191), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/feedforward_layernorm/rsqrt" stack_frame_id=753} + %mul.3426 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.212), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3427 = f32[1,41]{1,0} reshape(%mul.3426), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3428 = f32[1,41,4096]{2,1,0} broadcast(%mul.3427), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3429 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.1887, %mul.3428), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/feedforward_layernorm/mul" stack_frame_id=754} + %state_1__91_.1 = bf16[4096]{0} parameter(94), metadata={op_name="state[1][91]"} + %convert_element_type.1888 = f32[4096]{0} convert(%state_1__91_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/feedforward_layernorm/convert_element_type" stack_frame_id=755} + %broadcast_in_dim.1670 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1888), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/feedforward_layernorm/broadcast_in_dim" stack_frame_id=755} + %mul.3430 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1670), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3431 = f32[1,4096]{1,0} reshape(%mul.3430), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3432 = f32[1,41,4096]{2,1,0} broadcast(%mul.3431), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3433 = f32[1,41,4096]{2,1,0} multiply(%mul.3429, %mul.3432), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/feedforward_layernorm/mul" stack_frame_id=755} + %convert_element_type.1889 = bf16[1,41,4096]{2,1,0} convert(%mul.3433), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/feedforward_layernorm/convert_element_type" stack_frame_id=759} + %state_1__88_.1 = bf16[4096,14336]{1,0} parameter(91), metadata={op_name="state[1][88]"} + %dot_general.1158 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.1889, %state_1__88_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/feedforward_intermediate_dense/dot_general" stack_frame_id=788} + %state_1__89_.1 = bf16[4096,14336]{1,0} parameter(92), metadata={op_name="state[1][89]"} + %dot_general.1157 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.1889, %state_1__89_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/feedforward_gate_dense/dot_general" stack_frame_id=768} + %convert_element_type.1890 = f32[1,41,14336]{2,1,0} convert(%dot_general.1157), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/convert_element_type" stack_frame_id=772} + %jit_silu_.104 = f32[1,41,14336]{2,1,0} call(%convert_element_type.1890), to_apply=%silu.9, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/jit(silu)" stack_frame_id=775} + %convert_element_type.1891 = bf16[1,41,14336]{2,1,0} convert(%jit_silu_.104), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/convert_element_type" stack_frame_id=779} + %mul.3434 = bf16[1,41,14336]{2,1,0} multiply(%dot_general.1158, %convert_element_type.1891), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/mul" stack_frame_id=791} + %state_1__90_.1 = bf16[14336,4096]{1,0} parameter(93), metadata={op_name="state[1][90]"} + %dot_general.1159 = bf16[1,41,4096]{2,1,0} dot(%mul.3434, %state_1__90_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/feedforward_output_dense/dot_general" stack_frame_id=800} + %add.1192 = bf16[1,41,4096]{2,1,0} add(%dot_general.1159, %add.1190), metadata={op_name="jit(compiled_generate_function)/transformer_layer_9/add" stack_frame_id=801} + %convert_element_type.1894 = f32[1,41,4096]{2,1,0} convert(%add.1192), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention_layernorm/convert_element_type" stack_frame_id=555} + %pow.424 = f32[1,41,4096]{2,1,0} power(%convert_element_type.1894, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention_layernorm/pow" stack_frame_id=558} + %reduce_sum.1479 = f32[1,41]{1,0} reduce(%pow.424, %constant.155), dimensions={2}, to_apply=%region_41.46, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention_layernorm/reduce_sum" stack_frame_id=24} + %broadcast_in_dim.1677 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1479), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention_layernorm/broadcast_in_dim" stack_frame_id=24} + %div.1069 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1677, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention_layernorm/div" stack_frame_id=24} + %add.1194 = f32[1,41,1]{2,1,0} add(%div.1069, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention_layernorm/add" stack_frame_id=559} + %rsqrt.213 = f32[1,41,1]{2,1,0} rsqrt(%add.1194), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention_layernorm/rsqrt" stack_frame_id=562} + %mul.3435 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.213), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3436 = f32[1,41]{1,0} reshape(%mul.3435), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3437 = f32[1,41,4096]{2,1,0} broadcast(%mul.3436), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3438 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.1894, %mul.3437), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention_layernorm/mul" stack_frame_id=563} + %state_1__96_.1 = bf16[4096]{0} parameter(99), metadata={op_name="state[1][96]"} + %convert_element_type.1895 = f32[4096]{0} convert(%state_1__96_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention_layernorm/convert_element_type" stack_frame_id=564} + %broadcast_in_dim.1678 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1895), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention_layernorm/broadcast_in_dim" stack_frame_id=564} + %mul.3439 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1678), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3440 = f32[1,4096]{1,0} reshape(%mul.3439), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3441 = f32[1,41,4096]{2,1,0} broadcast(%mul.3440), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3442 = f32[1,41,4096]{2,1,0} multiply(%mul.3438, %mul.3441), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention_layernorm/mul" stack_frame_id=564} + %convert_element_type.1896 = bf16[1,41,4096]{2,1,0} convert(%mul.3442), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention_layernorm/convert_element_type" stack_frame_id=568} + %state_1__93_.1 = bf16[4096,8,128]{2,1,0} parameter(96), metadata={op_name="state[1][93]"} + %dot_general.1162 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.1896, %state_1__93_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=647} + %iota.597 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/iota" stack_frame_id=677} + %broadcast_in_dim.1680 = f32[1,41]{1,0} reshape(%iota.597), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/broadcast_in_dim" stack_frame_id=680} + %iota.596 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/iota" stack_frame_id=667} + %mul.3452 = f32[64]{0} multiply(%iota.596, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=667} + %div.1072 = f32[64]{0} divide(%mul.3452, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/div" stack_frame_id=668} + %neg.432 = f32[64]{0} negate(%div.1072), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/neg" stack_frame_id=669} + %pow.426 = f32[64]{0} power(%broadcast.51, %neg.432), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/pow" stack_frame_id=672} + %div.1073 = f32[64]{0} divide(%pow.426, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/div" stack_frame_id=673} + %dot_general.1164 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1680, %div.1073), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/bi,j->bij/dot_general" stack_frame_id=683} + %stack.1689 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1164), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/stack" stack_frame_id=686} + %stack.1690 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1164), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/stack" stack_frame_id=686} + %stack.1691 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1689, %stack.1690), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/stack" stack_frame_id=686} + %reshape.780 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1691), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/reshape"} + %cos.212 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.780), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/cos" stack_frame_id=690} + %convert_element_type.1899 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.212), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/convert_element_type" stack_frame_id=694} + %mul.3453 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1899), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=714} + %mul.3454 = bf16[1,41,128]{2,1,0} reshape(%mul.3453), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=714} + %mul.3455 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3454), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=714} + %mul.3456 = bf16[1,41,8,128]{3,2,1,0} multiply(%dot_general.1162, %mul.3455), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=714} + %split.425 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1162), slice={[0:1], [0:41], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/split" stack_frame_id=706} + %neg.433 = bf16[1,41,8,64]{3,2,1,0} negate(%split.425), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/neg" stack_frame_id=707} + %stack.1692 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%neg.433), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/stack" stack_frame_id=710} + %split.424 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1162), slice={[0:1], [0:41], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/split" stack_frame_id=706} + %stack.1693 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%split.424), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/stack" stack_frame_id=710} + %stack.1694 = bf16[1,41,8,2,64]{4,3,2,1,0} concatenate(%stack.1692, %stack.1693), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/stack" stack_frame_id=710} + %reshape.781 = bf16[1,41,8,128]{3,2,1,0} reshape(%stack.1694), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/reshape" stack_frame_id=713} + %sin.212 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.780), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/sin" stack_frame_id=698} + %convert_element_type.1900 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.212), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/convert_element_type" stack_frame_id=702} + %mul.3457 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1900), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=715} + %mul.3458 = bf16[1,41,128]{2,1,0} reshape(%mul.3457), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=715} + %mul.3459 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3458), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=715} + %mul.3460 = bf16[1,41,8,128]{3,2,1,0} multiply(%reshape.781, %mul.3459), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=715} + %add.1196 = bf16[1,41,8,128]{3,2,1,0} add(%mul.3456, %mul.3460), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/add" stack_frame_id=716} + %stack.1695 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%add.1196), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/stack" stack_frame_id=719} + %state_1__94_.1 = bf16[4096,8,128]{2,1,0} parameter(97), metadata={op_name="state[1][94]"} + %dot_general.1163 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.1896, %state_1__94_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=656} + %stack.1696 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dot_general.1163), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/stack" stack_frame_id=719} + %stack.1697 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1695, %stack.1696), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/stack" stack_frame_id=719} + %stack.2017 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1697), metadata={op_name="jit(compiled_generate_function)/stack" stack_frame_id=804} + %broadcast_in_dim.1682 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dot_general.1163), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/broadcast_in_dim" stack_frame_id=725} + %reshape.783 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1682), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/reshape" stack_frame_id=725} + %iota.590 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/iota" stack_frame_id=525} + %broadcast_in_dim.1671 = f32[41,1]{1,0} reshape(%iota.590), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/broadcast_in_dim" stack_frame_id=528} + %ge.551 = f32[41,1]{1,0} broadcast(%broadcast_in_dim.1671), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/ge" stack_frame_id=532} + %ge.552 = f32[41]{0} reshape(%ge.551), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/ge" stack_frame_id=532} + %ge.553 = f32[41,41]{1,0} broadcast(%ge.552), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/ge" stack_frame_id=532} + %iota.591 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/iota" stack_frame_id=531} + %broadcast_in_dim.1672 = f32[1,41]{1,0} reshape(%iota.591), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/broadcast_in_dim" stack_frame_id=532} + %ge.554 = f32[1,41]{1,0} broadcast(%broadcast_in_dim.1672), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/ge" stack_frame_id=532} + %ge.555 = f32[41]{0} reshape(%ge.554), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/ge" stack_frame_id=532} + %ge.556 = f32[41,41]{1,0} broadcast(%ge.555), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/ge" stack_frame_id=532} + %ge.557 = pred[41,41]{1,0} compare(%ge.553, %ge.556), direction=GE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/ge" stack_frame_id=532} + %broadcast_in_dim.1673 = pred[1,41,41]{2,1,0} reshape(%ge.557), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/broadcast_in_dim" stack_frame_id=535} + %convert_element_type.1893 = s32[1,41,41]{2,1,0} convert(%broadcast_in_dim.1673), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/convert_element_type" stack_frame_id=551} + %iota.592 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/iota" stack_frame_id=538} + %broadcast_in_dim.1674 = s32[41,1]{1,0} reshape(%iota.592), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/broadcast_in_dim" stack_frame_id=536} + %lt.693 = s32[41,1]{1,0} broadcast(%broadcast_in_dim.1674), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/lt" stack_frame_id=543} + %lt.694 = s32[41]{0} reshape(%lt.693), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/lt" stack_frame_id=543} + %lt.695 = s32[41,41]{1,0} broadcast(%lt.694), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/lt" stack_frame_id=543} + %iota.593 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/iota" stack_frame_id=541} + %broadcast_in_dim.1675 = s32[1,41]{1,0} reshape(%iota.593), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/broadcast_in_dim" stack_frame_id=539} + %add.1193 = s32[1,41]{1,0} add(%broadcast_in_dim.1675, %broadcast.57), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/add" stack_frame_id=542} + %lt.696 = s32[1,41]{1,0} broadcast(%add.1193), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/lt" stack_frame_id=543} + %lt.697 = s32[41]{0} reshape(%lt.696), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/lt" stack_frame_id=543} + %lt.698 = s32[41,41]{1,0} broadcast(%lt.697), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/lt" stack_frame_id=543} + %lt.699 = pred[41,41]{1,0} compare(%lt.695, %lt.698), direction=LT, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/lt" stack_frame_id=543} + %convert_element_type.1892 = s32[41,41]{1,0} convert(%lt.699), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/convert_element_type" stack_frame_id=547} + %broadcast_in_dim.1676 = s32[1,41,41]{2,1,0} reshape(%convert_element_type.1892), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/broadcast_in_dim" stack_frame_id=551} + %min.105 = s32[1,41,41]{2,1,0} minimum(%convert_element_type.1893, %broadcast_in_dim.1676), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/min" stack_frame_id=551} + %broadcast_in_dim.1683 = s32[1,1,41,41]{3,2,1,0} reshape(%min.105), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/broadcast_in_dim" stack_frame_id=728} + %convert_element_type.1901 = pred[1,1,41,41]{3,2,1,0} compare(%broadcast_in_dim.1683, %broadcast.49), direction=NE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/convert_element_type" stack_frame_id=732} + %broadcast_in_dim.1684 = pred[1,1,1,41,41]{4,3,2,1,0} reshape(%convert_element_type.1901), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/broadcast_in_dim" stack_frame_id=34} + %and.328 = pred[1,1,1,41,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1684), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/vmap()/and" stack_frame_id=34} + %and.329 = pred[1,1,41,41]{3,2,1,0} reshape(%and.328), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/vmap()/and" stack_frame_id=34} + %and.330 = pred[1,1,32,41,41]{4,3,2,1,0} broadcast(%and.329), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/vmap()/and" stack_frame_id=34} + %state_1__92_.1 = bf16[4096,32,128]{2,1,0} parameter(95), metadata={op_name="state[1][92]"} + %dot_general.1160 = bf16[1,41,32,128]{3,2,1,0} dot(%convert_element_type.1896, %state_1__92_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=577} + %iota.595 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/iota" stack_frame_id=598} + %broadcast_in_dim.1679 = f32[1,41]{1,0} reshape(%iota.595), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/broadcast_in_dim" stack_frame_id=601} + %iota.594 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/iota" stack_frame_id=588} + %mul.3443 = f32[64]{0} multiply(%iota.594, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=588} + %div.1070 = f32[64]{0} divide(%mul.3443, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/div" stack_frame_id=589} + %neg.430 = f32[64]{0} negate(%div.1070), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/neg" stack_frame_id=590} + %pow.425 = f32[64]{0} power(%broadcast.51, %neg.430), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/pow" stack_frame_id=593} + %div.1071 = f32[64]{0} divide(%pow.425, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/div" stack_frame_id=594} + %dot_general.1161 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1679, %div.1071), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/bi,j->bij/dot_general" stack_frame_id=604} + %stack.1683 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1161), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/stack" stack_frame_id=607} + %stack.1684 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1161), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/stack" stack_frame_id=607} + %stack.1685 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1683, %stack.1684), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/stack" stack_frame_id=607} + %reshape.778 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1685), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/reshape"} + %cos.211 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.778), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/cos" stack_frame_id=611} + %convert_element_type.1897 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.211), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/convert_element_type" stack_frame_id=615} + %mul.3444 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1897), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=635} + %mul.3445 = bf16[1,41,128]{2,1,0} reshape(%mul.3444), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=635} + %mul.3446 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3445), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=635} + %mul.3447 = bf16[1,41,32,128]{3,2,1,0} multiply(%dot_general.1160, %mul.3446), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=635} + %split.423 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1160), slice={[0:1], [0:41], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/split" stack_frame_id=627} + %neg.431 = bf16[1,41,32,64]{3,2,1,0} negate(%split.423), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/neg" stack_frame_id=628} + %stack.1686 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%neg.431), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/stack" stack_frame_id=631} + %split.422 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1160), slice={[0:1], [0:41], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/split" stack_frame_id=627} + %stack.1687 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%split.422), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/stack" stack_frame_id=631} + %stack.1688 = bf16[1,41,32,2,64]{4,3,2,1,0} concatenate(%stack.1686, %stack.1687), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/stack" stack_frame_id=631} + %reshape.779 = bf16[1,41,32,128]{3,2,1,0} reshape(%stack.1688), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/reshape" stack_frame_id=634} + %sin.211 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.778), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/sin" stack_frame_id=619} + %convert_element_type.1898 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.211), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/convert_element_type" stack_frame_id=623} + %mul.3448 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1898), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=636} + %mul.3449 = bf16[1,41,128]{2,1,0} reshape(%mul.3448), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=636} + %mul.3450 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3449), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=636} + %mul.3451 = bf16[1,41,32,128]{3,2,1,0} multiply(%reshape.779, %mul.3450), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/mul" stack_frame_id=636} + %add.1195 = bf16[1,41,32,128]{3,2,1,0} add(%mul.3447, %mul.3451), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/rotary_embedding_10/add" stack_frame_id=637} + %reshape.784 = bf16[1,41,32,1,128]{4,3,2,1,0} reshape(%add.1195), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/reshape" stack_frame_id=34} + %broadcast_in_dim.1681 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%add.1196), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/broadcast_in_dim" stack_frame_id=722} + %reshape.782 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1681), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/reshape" stack_frame_id=722} + %dot_general.1165 = f32[1,32,41,1,41]{4,3,2,1,0} dot(%reshape.784, %reshape.782), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=34} + %mul.3461 = f32[1,32,41,1,41]{4,3,2,1,0} multiply(%dot_general.1165, %broadcast.48), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/vmap()/mul" stack_frame_id=34} + %vmap_jit__where__.105 = f32[1,1,32,41,41]{4,3,2,1,0} call(%and.330, %mul.3461, %constant.153), to_apply=%_where_0.5, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/vmap(jit(_where))" stack_frame_id=34} + %reduce_max.491 = f32[1,1,32,41]{3,2,1,0} reduce(%vmap_jit__where__.105, %constant.152), dimensions={4}, to_apply=%region_42.47, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/vmap()/reduce_max" stack_frame_id=34} + %max.115 = f32[1,1,32,41]{3,2,1,0} maximum(%reduce_max.491, %broadcast.47), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/vmap()/max" stack_frame_id=34} + %broadcast_in_dim.1685 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%max.115), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %sub.444 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1685), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/vmap()/sub" stack_frame_id=34} + %sub.445 = f32[1,1,32,41]{3,2,1,0} reshape(%sub.444), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/vmap()/sub" stack_frame_id=34} + %sub.446 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%sub.445), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/vmap()/sub" stack_frame_id=34} + %sub.447 = f32[1,1,32,41,41]{4,3,2,1,0} subtract(%vmap_jit__where__.105, %sub.446), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/vmap()/sub" stack_frame_id=34} + %exp.111 = f32[1,1,32,41,41]{4,3,2,1,0} exponential(%sub.447), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/vmap()/exp" stack_frame_id=34} + %reduce_sum.1480 = f32[1,1,32,41]{3,2,1,0} reduce(%exp.111, %constant.155), dimensions={4}, to_apply=%region_43.48, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/vmap()/reduce_sum" stack_frame_id=34} + %broadcast_in_dim.1686 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%reduce_sum.1480), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %div.1074 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1686), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/vmap()/div" stack_frame_id=34} + %div.1075 = f32[1,1,32,41]{3,2,1,0} reshape(%div.1074), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/vmap()/div" stack_frame_id=34} + %div.1076 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%div.1075), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/vmap()/div" stack_frame_id=34} + %div.1077 = f32[1,1,32,41,41]{4,3,2,1,0} divide(%exp.111, %div.1076), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/vmap()/div" stack_frame_id=34} + %convert_element_type.1902 = bf16[1,1,32,41,41]{4,3,2,1,0} convert(%div.1077), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/vmap()/convert_element_type" stack_frame_id=34} + %dot_general.1166 = bf16[1,32,128,1,41]{4,3,2,1,0} dot(%reshape.783, %convert_element_type.1902), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=34} + %transpose.109 = bf16[1,41,32,1,128]{1,3,4,2,0} transpose(%dot_general.1166), dimensions={0,4,1,3,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/vmap()/transpose;jit(compiled_generate_function)/transformer_layer_10/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %reshape.785 = bf16[1,41,32,128]{3,2,1,0} reshape(%transpose.109), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/reshape" stack_frame_id=34} + %state_1__95_.1 = bf16[32,128,4096]{2,1,0} parameter(98), metadata={op_name="state[1][95]"} + %dot_general.1167 = bf16[1,41,4096]{2,1,0} dot(%reshape.785, %state_1__95_.1), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=741} + %add.1197 = bf16[1,41,4096]{2,1,0} add(%dot_general.1167, %add.1192), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/add" stack_frame_id=742} + %convert_element_type.1903 = f32[1,41,4096]{2,1,0} convert(%add.1197), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/feedforward_layernorm/convert_element_type" stack_frame_id=746} + %pow.427 = f32[1,41,4096]{2,1,0} power(%convert_element_type.1903, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/feedforward_layernorm/pow" stack_frame_id=749} + %reduce_sum.1481 = f32[1,41]{1,0} reduce(%pow.427, %constant.155), dimensions={2}, to_apply=%region_44.49, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/feedforward_layernorm/reduce_sum" stack_frame_id=43} + %broadcast_in_dim.1687 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1481), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/feedforward_layernorm/broadcast_in_dim" stack_frame_id=43} + %div.1078 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1687, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/feedforward_layernorm/div" stack_frame_id=43} + %add.1198 = f32[1,41,1]{2,1,0} add(%div.1078, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/feedforward_layernorm/add" stack_frame_id=750} + %rsqrt.214 = f32[1,41,1]{2,1,0} rsqrt(%add.1198), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/feedforward_layernorm/rsqrt" stack_frame_id=753} + %mul.3462 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.214), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3463 = f32[1,41]{1,0} reshape(%mul.3462), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3464 = f32[1,41,4096]{2,1,0} broadcast(%mul.3463), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3465 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.1903, %mul.3464), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/feedforward_layernorm/mul" stack_frame_id=754} + %state_1__100_.1 = bf16[4096]{0} parameter(103), metadata={op_name="state[1][100]"} + %convert_element_type.1904 = f32[4096]{0} convert(%state_1__100_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/feedforward_layernorm/convert_element_type" stack_frame_id=755} + %broadcast_in_dim.1688 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1904), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/feedforward_layernorm/broadcast_in_dim" stack_frame_id=755} + %mul.3466 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1688), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3467 = f32[1,4096]{1,0} reshape(%mul.3466), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3468 = f32[1,41,4096]{2,1,0} broadcast(%mul.3467), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3469 = f32[1,41,4096]{2,1,0} multiply(%mul.3465, %mul.3468), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/feedforward_layernorm/mul" stack_frame_id=755} + %convert_element_type.1905 = bf16[1,41,4096]{2,1,0} convert(%mul.3469), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/feedforward_layernorm/convert_element_type" stack_frame_id=759} + %state_1__97_.1 = bf16[4096,14336]{1,0} parameter(100), metadata={op_name="state[1][97]"} + %dot_general.1169 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.1905, %state_1__97_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/feedforward_intermediate_dense/dot_general" stack_frame_id=788} + %state_1__98_.1 = bf16[4096,14336]{1,0} parameter(101), metadata={op_name="state[1][98]"} + %dot_general.1168 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.1905, %state_1__98_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/feedforward_gate_dense/dot_general" stack_frame_id=768} + %convert_element_type.1906 = f32[1,41,14336]{2,1,0} convert(%dot_general.1168), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/convert_element_type" stack_frame_id=772} + %jit_silu_.105 = f32[1,41,14336]{2,1,0} call(%convert_element_type.1906), to_apply=%silu.9, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/jit(silu)" stack_frame_id=775} + %convert_element_type.1907 = bf16[1,41,14336]{2,1,0} convert(%jit_silu_.105), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/convert_element_type" stack_frame_id=779} + %mul.3470 = bf16[1,41,14336]{2,1,0} multiply(%dot_general.1169, %convert_element_type.1907), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/mul" stack_frame_id=791} + %state_1__99_.1 = bf16[14336,4096]{1,0} parameter(102), metadata={op_name="state[1][99]"} + %dot_general.1170 = bf16[1,41,4096]{2,1,0} dot(%mul.3470, %state_1__99_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/feedforward_output_dense/dot_general" stack_frame_id=800} + %add.1199 = bf16[1,41,4096]{2,1,0} add(%dot_general.1170, %add.1197), metadata={op_name="jit(compiled_generate_function)/transformer_layer_10/add" stack_frame_id=801} + %convert_element_type.1910 = f32[1,41,4096]{2,1,0} convert(%add.1199), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention_layernorm/convert_element_type" stack_frame_id=555} + %pow.428 = f32[1,41,4096]{2,1,0} power(%convert_element_type.1910, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention_layernorm/pow" stack_frame_id=558} + %reduce_sum.1482 = f32[1,41]{1,0} reduce(%pow.428, %constant.155), dimensions={2}, to_apply=%region_45.50, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention_layernorm/reduce_sum" stack_frame_id=24} + %broadcast_in_dim.1695 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1482), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention_layernorm/broadcast_in_dim" stack_frame_id=24} + %div.1079 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1695, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention_layernorm/div" stack_frame_id=24} + %add.1201 = f32[1,41,1]{2,1,0} add(%div.1079, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention_layernorm/add" stack_frame_id=559} + %rsqrt.215 = f32[1,41,1]{2,1,0} rsqrt(%add.1201), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention_layernorm/rsqrt" stack_frame_id=562} + %mul.3471 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.215), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3472 = f32[1,41]{1,0} reshape(%mul.3471), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3473 = f32[1,41,4096]{2,1,0} broadcast(%mul.3472), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3474 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.1910, %mul.3473), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention_layernorm/mul" stack_frame_id=563} + %state_1__105_.1 = bf16[4096]{0} parameter(108), metadata={op_name="state[1][105]"} + %convert_element_type.1911 = f32[4096]{0} convert(%state_1__105_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention_layernorm/convert_element_type" stack_frame_id=564} + %broadcast_in_dim.1696 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1911), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention_layernorm/broadcast_in_dim" stack_frame_id=564} + %mul.3475 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1696), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3476 = f32[1,4096]{1,0} reshape(%mul.3475), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3477 = f32[1,41,4096]{2,1,0} broadcast(%mul.3476), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3478 = f32[1,41,4096]{2,1,0} multiply(%mul.3474, %mul.3477), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention_layernorm/mul" stack_frame_id=564} + %convert_element_type.1912 = bf16[1,41,4096]{2,1,0} convert(%mul.3478), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention_layernorm/convert_element_type" stack_frame_id=568} + %state_1__102_.1 = bf16[4096,8,128]{2,1,0} parameter(105), metadata={op_name="state[1][102]"} + %dot_general.1173 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.1912, %state_1__102_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=647} + %iota.605 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/iota" stack_frame_id=677} + %broadcast_in_dim.1698 = f32[1,41]{1,0} reshape(%iota.605), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/broadcast_in_dim" stack_frame_id=680} + %iota.604 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/iota" stack_frame_id=667} + %mul.3488 = f32[64]{0} multiply(%iota.604, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=667} + %div.1082 = f32[64]{0} divide(%mul.3488, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/div" stack_frame_id=668} + %neg.436 = f32[64]{0} negate(%div.1082), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/neg" stack_frame_id=669} + %pow.430 = f32[64]{0} power(%broadcast.51, %neg.436), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/pow" stack_frame_id=672} + %div.1083 = f32[64]{0} divide(%pow.430, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/div" stack_frame_id=673} + %dot_general.1175 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1698, %div.1083), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/bi,j->bij/dot_general" stack_frame_id=683} + %stack.1704 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1175), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/stack" stack_frame_id=686} + %stack.1705 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1175), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/stack" stack_frame_id=686} + %stack.1706 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1704, %stack.1705), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/stack" stack_frame_id=686} + %reshape.788 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1706), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/reshape"} + %cos.214 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.788), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/cos" stack_frame_id=690} + %convert_element_type.1915 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.214), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/convert_element_type" stack_frame_id=694} + %mul.3489 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1915), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=714} + %mul.3490 = bf16[1,41,128]{2,1,0} reshape(%mul.3489), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=714} + %mul.3491 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3490), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=714} + %mul.3492 = bf16[1,41,8,128]{3,2,1,0} multiply(%dot_general.1173, %mul.3491), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=714} + %split.429 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1173), slice={[0:1], [0:41], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/split" stack_frame_id=706} + %neg.437 = bf16[1,41,8,64]{3,2,1,0} negate(%split.429), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/neg" stack_frame_id=707} + %stack.1707 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%neg.437), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/stack" stack_frame_id=710} + %split.428 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1173), slice={[0:1], [0:41], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/split" stack_frame_id=706} + %stack.1708 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%split.428), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/stack" stack_frame_id=710} + %stack.1709 = bf16[1,41,8,2,64]{4,3,2,1,0} concatenate(%stack.1707, %stack.1708), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/stack" stack_frame_id=710} + %reshape.789 = bf16[1,41,8,128]{3,2,1,0} reshape(%stack.1709), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/reshape" stack_frame_id=713} + %sin.214 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.788), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/sin" stack_frame_id=698} + %convert_element_type.1916 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.214), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/convert_element_type" stack_frame_id=702} + %mul.3493 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1916), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=715} + %mul.3494 = bf16[1,41,128]{2,1,0} reshape(%mul.3493), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=715} + %mul.3495 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3494), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=715} + %mul.3496 = bf16[1,41,8,128]{3,2,1,0} multiply(%reshape.789, %mul.3495), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=715} + %add.1203 = bf16[1,41,8,128]{3,2,1,0} add(%mul.3492, %mul.3496), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/add" stack_frame_id=716} + %stack.1710 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%add.1203), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/stack" stack_frame_id=719} + %state_1__103_.1 = bf16[4096,8,128]{2,1,0} parameter(106), metadata={op_name="state[1][103]"} + %dot_general.1174 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.1912, %state_1__103_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=656} + %stack.1711 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dot_general.1174), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/stack" stack_frame_id=719} + %stack.1712 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1710, %stack.1711), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/stack" stack_frame_id=719} + %stack.2018 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1712), metadata={op_name="jit(compiled_generate_function)/stack" stack_frame_id=804} + %broadcast_in_dim.1700 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dot_general.1174), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/broadcast_in_dim" stack_frame_id=725} + %reshape.791 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1700), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/reshape" stack_frame_id=725} + %iota.598 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/iota" stack_frame_id=525} + %broadcast_in_dim.1689 = f32[41,1]{1,0} reshape(%iota.598), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/broadcast_in_dim" stack_frame_id=528} + %ge.558 = f32[41,1]{1,0} broadcast(%broadcast_in_dim.1689), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/ge" stack_frame_id=532} + %ge.559 = f32[41]{0} reshape(%ge.558), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/ge" stack_frame_id=532} + %ge.560 = f32[41,41]{1,0} broadcast(%ge.559), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/ge" stack_frame_id=532} + %iota.599 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/iota" stack_frame_id=531} + %broadcast_in_dim.1690 = f32[1,41]{1,0} reshape(%iota.599), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/broadcast_in_dim" stack_frame_id=532} + %ge.561 = f32[1,41]{1,0} broadcast(%broadcast_in_dim.1690), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/ge" stack_frame_id=532} + %ge.562 = f32[41]{0} reshape(%ge.561), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/ge" stack_frame_id=532} + %ge.563 = f32[41,41]{1,0} broadcast(%ge.562), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/ge" stack_frame_id=532} + %ge.564 = pred[41,41]{1,0} compare(%ge.560, %ge.563), direction=GE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/ge" stack_frame_id=532} + %broadcast_in_dim.1691 = pred[1,41,41]{2,1,0} reshape(%ge.564), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/broadcast_in_dim" stack_frame_id=535} + %convert_element_type.1909 = s32[1,41,41]{2,1,0} convert(%broadcast_in_dim.1691), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/convert_element_type" stack_frame_id=551} + %iota.600 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/iota" stack_frame_id=538} + %broadcast_in_dim.1692 = s32[41,1]{1,0} reshape(%iota.600), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/broadcast_in_dim" stack_frame_id=536} + %lt.700 = s32[41,1]{1,0} broadcast(%broadcast_in_dim.1692), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/lt" stack_frame_id=543} + %lt.701 = s32[41]{0} reshape(%lt.700), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/lt" stack_frame_id=543} + %lt.702 = s32[41,41]{1,0} broadcast(%lt.701), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/lt" stack_frame_id=543} + %iota.601 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/iota" stack_frame_id=541} + %broadcast_in_dim.1693 = s32[1,41]{1,0} reshape(%iota.601), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/broadcast_in_dim" stack_frame_id=539} + %add.1200 = s32[1,41]{1,0} add(%broadcast_in_dim.1693, %broadcast.57), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/add" stack_frame_id=542} + %lt.703 = s32[1,41]{1,0} broadcast(%add.1200), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/lt" stack_frame_id=543} + %lt.704 = s32[41]{0} reshape(%lt.703), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/lt" stack_frame_id=543} + %lt.705 = s32[41,41]{1,0} broadcast(%lt.704), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/lt" stack_frame_id=543} + %lt.706 = pred[41,41]{1,0} compare(%lt.702, %lt.705), direction=LT, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/lt" stack_frame_id=543} + %convert_element_type.1908 = s32[41,41]{1,0} convert(%lt.706), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/convert_element_type" stack_frame_id=547} + %broadcast_in_dim.1694 = s32[1,41,41]{2,1,0} reshape(%convert_element_type.1908), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/broadcast_in_dim" stack_frame_id=551} + %min.106 = s32[1,41,41]{2,1,0} minimum(%convert_element_type.1909, %broadcast_in_dim.1694), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/min" stack_frame_id=551} + %broadcast_in_dim.1701 = s32[1,1,41,41]{3,2,1,0} reshape(%min.106), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/broadcast_in_dim" stack_frame_id=728} + %convert_element_type.1917 = pred[1,1,41,41]{3,2,1,0} compare(%broadcast_in_dim.1701, %broadcast.49), direction=NE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/convert_element_type" stack_frame_id=732} + %broadcast_in_dim.1702 = pred[1,1,1,41,41]{4,3,2,1,0} reshape(%convert_element_type.1917), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/broadcast_in_dim" stack_frame_id=34} + %and.331 = pred[1,1,1,41,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1702), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/vmap()/and" stack_frame_id=34} + %and.332 = pred[1,1,41,41]{3,2,1,0} reshape(%and.331), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/vmap()/and" stack_frame_id=34} + %and.333 = pred[1,1,32,41,41]{4,3,2,1,0} broadcast(%and.332), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/vmap()/and" stack_frame_id=34} + %state_1__101_.1 = bf16[4096,32,128]{2,1,0} parameter(104), metadata={op_name="state[1][101]"} + %dot_general.1171 = bf16[1,41,32,128]{3,2,1,0} dot(%convert_element_type.1912, %state_1__101_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=577} + %iota.603 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/iota" stack_frame_id=598} + %broadcast_in_dim.1697 = f32[1,41]{1,0} reshape(%iota.603), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/broadcast_in_dim" stack_frame_id=601} + %iota.602 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/iota" stack_frame_id=588} + %mul.3479 = f32[64]{0} multiply(%iota.602, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=588} + %div.1080 = f32[64]{0} divide(%mul.3479, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/div" stack_frame_id=589} + %neg.434 = f32[64]{0} negate(%div.1080), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/neg" stack_frame_id=590} + %pow.429 = f32[64]{0} power(%broadcast.51, %neg.434), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/pow" stack_frame_id=593} + %div.1081 = f32[64]{0} divide(%pow.429, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/div" stack_frame_id=594} + %dot_general.1172 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1697, %div.1081), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/bi,j->bij/dot_general" stack_frame_id=604} + %stack.1698 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1172), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/stack" stack_frame_id=607} + %stack.1699 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1172), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/stack" stack_frame_id=607} + %stack.1700 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1698, %stack.1699), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/stack" stack_frame_id=607} + %reshape.786 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1700), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/reshape"} + %cos.213 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.786), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/cos" stack_frame_id=611} + %convert_element_type.1913 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.213), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/convert_element_type" stack_frame_id=615} + %mul.3480 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1913), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=635} + %mul.3481 = bf16[1,41,128]{2,1,0} reshape(%mul.3480), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=635} + %mul.3482 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3481), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=635} + %mul.3483 = bf16[1,41,32,128]{3,2,1,0} multiply(%dot_general.1171, %mul.3482), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=635} + %split.427 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1171), slice={[0:1], [0:41], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/split" stack_frame_id=627} + %neg.435 = bf16[1,41,32,64]{3,2,1,0} negate(%split.427), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/neg" stack_frame_id=628} + %stack.1701 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%neg.435), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/stack" stack_frame_id=631} + %split.426 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1171), slice={[0:1], [0:41], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/split" stack_frame_id=627} + %stack.1702 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%split.426), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/stack" stack_frame_id=631} + %stack.1703 = bf16[1,41,32,2,64]{4,3,2,1,0} concatenate(%stack.1701, %stack.1702), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/stack" stack_frame_id=631} + %reshape.787 = bf16[1,41,32,128]{3,2,1,0} reshape(%stack.1703), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/reshape" stack_frame_id=634} + %sin.213 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.786), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/sin" stack_frame_id=619} + %convert_element_type.1914 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.213), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/convert_element_type" stack_frame_id=623} + %mul.3484 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1914), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=636} + %mul.3485 = bf16[1,41,128]{2,1,0} reshape(%mul.3484), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=636} + %mul.3486 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3485), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=636} + %mul.3487 = bf16[1,41,32,128]{3,2,1,0} multiply(%reshape.787, %mul.3486), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/mul" stack_frame_id=636} + %add.1202 = bf16[1,41,32,128]{3,2,1,0} add(%mul.3483, %mul.3487), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/rotary_embedding_11/add" stack_frame_id=637} + %reshape.792 = bf16[1,41,32,1,128]{4,3,2,1,0} reshape(%add.1202), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/reshape" stack_frame_id=34} + %broadcast_in_dim.1699 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%add.1203), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/broadcast_in_dim" stack_frame_id=722} + %reshape.790 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1699), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/reshape" stack_frame_id=722} + %dot_general.1176 = f32[1,32,41,1,41]{4,3,2,1,0} dot(%reshape.792, %reshape.790), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=34} + %mul.3497 = f32[1,32,41,1,41]{4,3,2,1,0} multiply(%dot_general.1176, %broadcast.48), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/vmap()/mul" stack_frame_id=34} + %vmap_jit__where__.106 = f32[1,1,32,41,41]{4,3,2,1,0} call(%and.333, %mul.3497, %constant.153), to_apply=%_where_0.5, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/vmap(jit(_where))" stack_frame_id=34} + %reduce_max.492 = f32[1,1,32,41]{3,2,1,0} reduce(%vmap_jit__where__.106, %constant.152), dimensions={4}, to_apply=%region_46.51, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/vmap()/reduce_max" stack_frame_id=34} + %max.116 = f32[1,1,32,41]{3,2,1,0} maximum(%reduce_max.492, %broadcast.47), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/vmap()/max" stack_frame_id=34} + %broadcast_in_dim.1703 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%max.116), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %sub.448 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1703), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/vmap()/sub" stack_frame_id=34} + %sub.449 = f32[1,1,32,41]{3,2,1,0} reshape(%sub.448), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/vmap()/sub" stack_frame_id=34} + %sub.450 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%sub.449), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/vmap()/sub" stack_frame_id=34} + %sub.451 = f32[1,1,32,41,41]{4,3,2,1,0} subtract(%vmap_jit__where__.106, %sub.450), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/vmap()/sub" stack_frame_id=34} + %exp.112 = f32[1,1,32,41,41]{4,3,2,1,0} exponential(%sub.451), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/vmap()/exp" stack_frame_id=34} + %reduce_sum.1483 = f32[1,1,32,41]{3,2,1,0} reduce(%exp.112, %constant.155), dimensions={4}, to_apply=%region_47.52, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/vmap()/reduce_sum" stack_frame_id=34} + %broadcast_in_dim.1704 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%reduce_sum.1483), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %div.1084 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1704), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/vmap()/div" stack_frame_id=34} + %div.1085 = f32[1,1,32,41]{3,2,1,0} reshape(%div.1084), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/vmap()/div" stack_frame_id=34} + %div.1086 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%div.1085), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/vmap()/div" stack_frame_id=34} + %div.1087 = f32[1,1,32,41,41]{4,3,2,1,0} divide(%exp.112, %div.1086), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/vmap()/div" stack_frame_id=34} + %convert_element_type.1918 = bf16[1,1,32,41,41]{4,3,2,1,0} convert(%div.1087), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/vmap()/convert_element_type" stack_frame_id=34} + %dot_general.1177 = bf16[1,32,128,1,41]{4,3,2,1,0} dot(%reshape.791, %convert_element_type.1918), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=34} + %transpose.110 = bf16[1,41,32,1,128]{1,3,4,2,0} transpose(%dot_general.1177), dimensions={0,4,1,3,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/vmap()/transpose;jit(compiled_generate_function)/transformer_layer_11/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %reshape.793 = bf16[1,41,32,128]{3,2,1,0} reshape(%transpose.110), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/reshape" stack_frame_id=34} + %state_1__104_.1 = bf16[32,128,4096]{2,1,0} parameter(107), metadata={op_name="state[1][104]"} + %dot_general.1178 = bf16[1,41,4096]{2,1,0} dot(%reshape.793, %state_1__104_.1), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=741} + %add.1204 = bf16[1,41,4096]{2,1,0} add(%dot_general.1178, %add.1199), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/add" stack_frame_id=742} + %convert_element_type.1919 = f32[1,41,4096]{2,1,0} convert(%add.1204), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/feedforward_layernorm/convert_element_type" stack_frame_id=746} + %pow.431 = f32[1,41,4096]{2,1,0} power(%convert_element_type.1919, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/feedforward_layernorm/pow" stack_frame_id=749} + %reduce_sum.1484 = f32[1,41]{1,0} reduce(%pow.431, %constant.155), dimensions={2}, to_apply=%region_48.53, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/feedforward_layernorm/reduce_sum" stack_frame_id=43} + %broadcast_in_dim.1705 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1484), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/feedforward_layernorm/broadcast_in_dim" stack_frame_id=43} + %div.1088 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1705, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/feedforward_layernorm/div" stack_frame_id=43} + %add.1205 = f32[1,41,1]{2,1,0} add(%div.1088, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/feedforward_layernorm/add" stack_frame_id=750} + %rsqrt.216 = f32[1,41,1]{2,1,0} rsqrt(%add.1205), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/feedforward_layernorm/rsqrt" stack_frame_id=753} + %mul.3498 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.216), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3499 = f32[1,41]{1,0} reshape(%mul.3498), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3500 = f32[1,41,4096]{2,1,0} broadcast(%mul.3499), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3501 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.1919, %mul.3500), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/feedforward_layernorm/mul" stack_frame_id=754} + %state_1__109_.1 = bf16[4096]{0} parameter(112), metadata={op_name="state[1][109]"} + %convert_element_type.1920 = f32[4096]{0} convert(%state_1__109_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/feedforward_layernorm/convert_element_type" stack_frame_id=755} + %broadcast_in_dim.1706 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1920), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/feedforward_layernorm/broadcast_in_dim" stack_frame_id=755} + %mul.3502 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1706), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3503 = f32[1,4096]{1,0} reshape(%mul.3502), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3504 = f32[1,41,4096]{2,1,0} broadcast(%mul.3503), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3505 = f32[1,41,4096]{2,1,0} multiply(%mul.3501, %mul.3504), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/feedforward_layernorm/mul" stack_frame_id=755} + %convert_element_type.1921 = bf16[1,41,4096]{2,1,0} convert(%mul.3505), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/feedforward_layernorm/convert_element_type" stack_frame_id=759} + %state_1__106_.1 = bf16[4096,14336]{1,0} parameter(109), metadata={op_name="state[1][106]"} + %dot_general.1180 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.1921, %state_1__106_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/feedforward_intermediate_dense/dot_general" stack_frame_id=788} + %state_1__107_.1 = bf16[4096,14336]{1,0} parameter(110), metadata={op_name="state[1][107]"} + %dot_general.1179 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.1921, %state_1__107_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/feedforward_gate_dense/dot_general" stack_frame_id=768} + %convert_element_type.1922 = f32[1,41,14336]{2,1,0} convert(%dot_general.1179), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/convert_element_type" stack_frame_id=772} + %jit_silu_.106 = f32[1,41,14336]{2,1,0} call(%convert_element_type.1922), to_apply=%silu.9, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/jit(silu)" stack_frame_id=775} + %convert_element_type.1923 = bf16[1,41,14336]{2,1,0} convert(%jit_silu_.106), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/convert_element_type" stack_frame_id=779} + %mul.3506 = bf16[1,41,14336]{2,1,0} multiply(%dot_general.1180, %convert_element_type.1923), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/mul" stack_frame_id=791} + %state_1__108_.1 = bf16[14336,4096]{1,0} parameter(111), metadata={op_name="state[1][108]"} + %dot_general.1181 = bf16[1,41,4096]{2,1,0} dot(%mul.3506, %state_1__108_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/feedforward_output_dense/dot_general" stack_frame_id=800} + %add.1206 = bf16[1,41,4096]{2,1,0} add(%dot_general.1181, %add.1204), metadata={op_name="jit(compiled_generate_function)/transformer_layer_11/add" stack_frame_id=801} + %convert_element_type.1926 = f32[1,41,4096]{2,1,0} convert(%add.1206), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention_layernorm/convert_element_type" stack_frame_id=555} + %pow.432 = f32[1,41,4096]{2,1,0} power(%convert_element_type.1926, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention_layernorm/pow" stack_frame_id=558} + %reduce_sum.1485 = f32[1,41]{1,0} reduce(%pow.432, %constant.155), dimensions={2}, to_apply=%region_49.54, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention_layernorm/reduce_sum" stack_frame_id=24} + %broadcast_in_dim.1713 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1485), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention_layernorm/broadcast_in_dim" stack_frame_id=24} + %div.1089 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1713, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention_layernorm/div" stack_frame_id=24} + %add.1208 = f32[1,41,1]{2,1,0} add(%div.1089, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention_layernorm/add" stack_frame_id=559} + %rsqrt.217 = f32[1,41,1]{2,1,0} rsqrt(%add.1208), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention_layernorm/rsqrt" stack_frame_id=562} + %mul.3507 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.217), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3508 = f32[1,41]{1,0} reshape(%mul.3507), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3509 = f32[1,41,4096]{2,1,0} broadcast(%mul.3508), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3510 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.1926, %mul.3509), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention_layernorm/mul" stack_frame_id=563} + %state_1__114_.1 = bf16[4096]{0} parameter(117), metadata={op_name="state[1][114]"} + %convert_element_type.1927 = f32[4096]{0} convert(%state_1__114_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention_layernorm/convert_element_type" stack_frame_id=564} + %broadcast_in_dim.1714 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1927), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention_layernorm/broadcast_in_dim" stack_frame_id=564} + %mul.3511 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1714), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3512 = f32[1,4096]{1,0} reshape(%mul.3511), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3513 = f32[1,41,4096]{2,1,0} broadcast(%mul.3512), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3514 = f32[1,41,4096]{2,1,0} multiply(%mul.3510, %mul.3513), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention_layernorm/mul" stack_frame_id=564} + %convert_element_type.1928 = bf16[1,41,4096]{2,1,0} convert(%mul.3514), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention_layernorm/convert_element_type" stack_frame_id=568} + %state_1__111_.1 = bf16[4096,8,128]{2,1,0} parameter(114), metadata={op_name="state[1][111]"} + %dot_general.1184 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.1928, %state_1__111_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=647} + %iota.613 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/iota" stack_frame_id=677} + %broadcast_in_dim.1716 = f32[1,41]{1,0} reshape(%iota.613), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/broadcast_in_dim" stack_frame_id=680} + %iota.612 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/iota" stack_frame_id=667} + %mul.3524 = f32[64]{0} multiply(%iota.612, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=667} + %div.1092 = f32[64]{0} divide(%mul.3524, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/div" stack_frame_id=668} + %neg.440 = f32[64]{0} negate(%div.1092), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/neg" stack_frame_id=669} + %pow.434 = f32[64]{0} power(%broadcast.51, %neg.440), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/pow" stack_frame_id=672} + %div.1093 = f32[64]{0} divide(%pow.434, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/div" stack_frame_id=673} + %dot_general.1186 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1716, %div.1093), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/bi,j->bij/dot_general" stack_frame_id=683} + %stack.1719 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1186), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/stack" stack_frame_id=686} + %stack.1720 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1186), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/stack" stack_frame_id=686} + %stack.1721 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1719, %stack.1720), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/stack" stack_frame_id=686} + %reshape.796 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1721), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/reshape"} + %cos.216 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.796), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/cos" stack_frame_id=690} + %convert_element_type.1931 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.216), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/convert_element_type" stack_frame_id=694} + %mul.3525 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1931), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=714} + %mul.3526 = bf16[1,41,128]{2,1,0} reshape(%mul.3525), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=714} + %mul.3527 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3526), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=714} + %mul.3528 = bf16[1,41,8,128]{3,2,1,0} multiply(%dot_general.1184, %mul.3527), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=714} + %split.433 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1184), slice={[0:1], [0:41], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/split" stack_frame_id=706} + %neg.441 = bf16[1,41,8,64]{3,2,1,0} negate(%split.433), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/neg" stack_frame_id=707} + %stack.1722 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%neg.441), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/stack" stack_frame_id=710} + %split.432 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1184), slice={[0:1], [0:41], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/split" stack_frame_id=706} + %stack.1723 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%split.432), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/stack" stack_frame_id=710} + %stack.1724 = bf16[1,41,8,2,64]{4,3,2,1,0} concatenate(%stack.1722, %stack.1723), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/stack" stack_frame_id=710} + %reshape.797 = bf16[1,41,8,128]{3,2,1,0} reshape(%stack.1724), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/reshape" stack_frame_id=713} + %sin.216 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.796), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/sin" stack_frame_id=698} + %convert_element_type.1932 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.216), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/convert_element_type" stack_frame_id=702} + %mul.3529 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1932), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=715} + %mul.3530 = bf16[1,41,128]{2,1,0} reshape(%mul.3529), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=715} + %mul.3531 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3530), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=715} + %mul.3532 = bf16[1,41,8,128]{3,2,1,0} multiply(%reshape.797, %mul.3531), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=715} + %add.1210 = bf16[1,41,8,128]{3,2,1,0} add(%mul.3528, %mul.3532), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/add" stack_frame_id=716} + %stack.1725 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%add.1210), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/stack" stack_frame_id=719} + %state_1__112_.1 = bf16[4096,8,128]{2,1,0} parameter(115), metadata={op_name="state[1][112]"} + %dot_general.1185 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.1928, %state_1__112_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=656} + %stack.1726 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dot_general.1185), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/stack" stack_frame_id=719} + %stack.1727 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1725, %stack.1726), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/stack" stack_frame_id=719} + %stack.2019 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1727), metadata={op_name="jit(compiled_generate_function)/stack" stack_frame_id=804} + %broadcast_in_dim.1718 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dot_general.1185), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/broadcast_in_dim" stack_frame_id=725} + %reshape.799 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1718), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/reshape" stack_frame_id=725} + %iota.606 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/iota" stack_frame_id=525} + %broadcast_in_dim.1707 = f32[41,1]{1,0} reshape(%iota.606), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/broadcast_in_dim" stack_frame_id=528} + %ge.565 = f32[41,1]{1,0} broadcast(%broadcast_in_dim.1707), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/ge" stack_frame_id=532} + %ge.566 = f32[41]{0} reshape(%ge.565), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/ge" stack_frame_id=532} + %ge.567 = f32[41,41]{1,0} broadcast(%ge.566), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/ge" stack_frame_id=532} + %iota.607 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/iota" stack_frame_id=531} + %broadcast_in_dim.1708 = f32[1,41]{1,0} reshape(%iota.607), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/broadcast_in_dim" stack_frame_id=532} + %ge.568 = f32[1,41]{1,0} broadcast(%broadcast_in_dim.1708), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/ge" stack_frame_id=532} + %ge.569 = f32[41]{0} reshape(%ge.568), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/ge" stack_frame_id=532} + %ge.570 = f32[41,41]{1,0} broadcast(%ge.569), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/ge" stack_frame_id=532} + %ge.571 = pred[41,41]{1,0} compare(%ge.567, %ge.570), direction=GE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/ge" stack_frame_id=532} + %broadcast_in_dim.1709 = pred[1,41,41]{2,1,0} reshape(%ge.571), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/broadcast_in_dim" stack_frame_id=535} + %convert_element_type.1925 = s32[1,41,41]{2,1,0} convert(%broadcast_in_dim.1709), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/convert_element_type" stack_frame_id=551} + %iota.608 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/iota" stack_frame_id=538} + %broadcast_in_dim.1710 = s32[41,1]{1,0} reshape(%iota.608), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/broadcast_in_dim" stack_frame_id=536} + %lt.707 = s32[41,1]{1,0} broadcast(%broadcast_in_dim.1710), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/lt" stack_frame_id=543} + %lt.708 = s32[41]{0} reshape(%lt.707), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/lt" stack_frame_id=543} + %lt.709 = s32[41,41]{1,0} broadcast(%lt.708), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/lt" stack_frame_id=543} + %iota.609 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/iota" stack_frame_id=541} + %broadcast_in_dim.1711 = s32[1,41]{1,0} reshape(%iota.609), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/broadcast_in_dim" stack_frame_id=539} + %add.1207 = s32[1,41]{1,0} add(%broadcast_in_dim.1711, %broadcast.57), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/add" stack_frame_id=542} + %lt.710 = s32[1,41]{1,0} broadcast(%add.1207), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/lt" stack_frame_id=543} + %lt.711 = s32[41]{0} reshape(%lt.710), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/lt" stack_frame_id=543} + %lt.712 = s32[41,41]{1,0} broadcast(%lt.711), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/lt" stack_frame_id=543} + %lt.713 = pred[41,41]{1,0} compare(%lt.709, %lt.712), direction=LT, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/lt" stack_frame_id=543} + %convert_element_type.1924 = s32[41,41]{1,0} convert(%lt.713), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/convert_element_type" stack_frame_id=547} + %broadcast_in_dim.1712 = s32[1,41,41]{2,1,0} reshape(%convert_element_type.1924), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/broadcast_in_dim" stack_frame_id=551} + %min.107 = s32[1,41,41]{2,1,0} minimum(%convert_element_type.1925, %broadcast_in_dim.1712), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/min" stack_frame_id=551} + %broadcast_in_dim.1719 = s32[1,1,41,41]{3,2,1,0} reshape(%min.107), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/broadcast_in_dim" stack_frame_id=728} + %convert_element_type.1933 = pred[1,1,41,41]{3,2,1,0} compare(%broadcast_in_dim.1719, %broadcast.49), direction=NE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/convert_element_type" stack_frame_id=732} + %broadcast_in_dim.1720 = pred[1,1,1,41,41]{4,3,2,1,0} reshape(%convert_element_type.1933), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/broadcast_in_dim" stack_frame_id=34} + %and.334 = pred[1,1,1,41,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1720), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/vmap()/and" stack_frame_id=34} + %and.335 = pred[1,1,41,41]{3,2,1,0} reshape(%and.334), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/vmap()/and" stack_frame_id=34} + %and.336 = pred[1,1,32,41,41]{4,3,2,1,0} broadcast(%and.335), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/vmap()/and" stack_frame_id=34} + %state_1__110_.1 = bf16[4096,32,128]{2,1,0} parameter(113), metadata={op_name="state[1][110]"} + %dot_general.1182 = bf16[1,41,32,128]{3,2,1,0} dot(%convert_element_type.1928, %state_1__110_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=577} + %iota.611 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/iota" stack_frame_id=598} + %broadcast_in_dim.1715 = f32[1,41]{1,0} reshape(%iota.611), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/broadcast_in_dim" stack_frame_id=601} + %iota.610 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/iota" stack_frame_id=588} + %mul.3515 = f32[64]{0} multiply(%iota.610, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=588} + %div.1090 = f32[64]{0} divide(%mul.3515, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/div" stack_frame_id=589} + %neg.438 = f32[64]{0} negate(%div.1090), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/neg" stack_frame_id=590} + %pow.433 = f32[64]{0} power(%broadcast.51, %neg.438), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/pow" stack_frame_id=593} + %div.1091 = f32[64]{0} divide(%pow.433, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/div" stack_frame_id=594} + %dot_general.1183 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1715, %div.1091), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/bi,j->bij/dot_general" stack_frame_id=604} + %stack.1713 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1183), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/stack" stack_frame_id=607} + %stack.1714 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1183), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/stack" stack_frame_id=607} + %stack.1715 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1713, %stack.1714), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/stack" stack_frame_id=607} + %reshape.794 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1715), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/reshape"} + %cos.215 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.794), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/cos" stack_frame_id=611} + %convert_element_type.1929 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.215), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/convert_element_type" stack_frame_id=615} + %mul.3516 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1929), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=635} + %mul.3517 = bf16[1,41,128]{2,1,0} reshape(%mul.3516), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=635} + %mul.3518 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3517), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=635} + %mul.3519 = bf16[1,41,32,128]{3,2,1,0} multiply(%dot_general.1182, %mul.3518), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=635} + %split.431 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1182), slice={[0:1], [0:41], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/split" stack_frame_id=627} + %neg.439 = bf16[1,41,32,64]{3,2,1,0} negate(%split.431), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/neg" stack_frame_id=628} + %stack.1716 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%neg.439), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/stack" stack_frame_id=631} + %split.430 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1182), slice={[0:1], [0:41], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/split" stack_frame_id=627} + %stack.1717 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%split.430), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/stack" stack_frame_id=631} + %stack.1718 = bf16[1,41,32,2,64]{4,3,2,1,0} concatenate(%stack.1716, %stack.1717), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/stack" stack_frame_id=631} + %reshape.795 = bf16[1,41,32,128]{3,2,1,0} reshape(%stack.1718), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/reshape" stack_frame_id=634} + %sin.215 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.794), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/sin" stack_frame_id=619} + %convert_element_type.1930 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.215), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/convert_element_type" stack_frame_id=623} + %mul.3520 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1930), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=636} + %mul.3521 = bf16[1,41,128]{2,1,0} reshape(%mul.3520), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=636} + %mul.3522 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3521), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=636} + %mul.3523 = bf16[1,41,32,128]{3,2,1,0} multiply(%reshape.795, %mul.3522), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/mul" stack_frame_id=636} + %add.1209 = bf16[1,41,32,128]{3,2,1,0} add(%mul.3519, %mul.3523), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/rotary_embedding_12/add" stack_frame_id=637} + %reshape.800 = bf16[1,41,32,1,128]{4,3,2,1,0} reshape(%add.1209), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/reshape" stack_frame_id=34} + %broadcast_in_dim.1717 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%add.1210), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/broadcast_in_dim" stack_frame_id=722} + %reshape.798 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1717), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/reshape" stack_frame_id=722} + %dot_general.1187 = f32[1,32,41,1,41]{4,3,2,1,0} dot(%reshape.800, %reshape.798), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=34} + %mul.3533 = f32[1,32,41,1,41]{4,3,2,1,0} multiply(%dot_general.1187, %broadcast.48), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/vmap()/mul" stack_frame_id=34} + %vmap_jit__where__.107 = f32[1,1,32,41,41]{4,3,2,1,0} call(%and.336, %mul.3533, %constant.153), to_apply=%_where_0.5, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/vmap(jit(_where))" stack_frame_id=34} + %reduce_max.493 = f32[1,1,32,41]{3,2,1,0} reduce(%vmap_jit__where__.107, %constant.152), dimensions={4}, to_apply=%region_50.55, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/vmap()/reduce_max" stack_frame_id=34} + %max.117 = f32[1,1,32,41]{3,2,1,0} maximum(%reduce_max.493, %broadcast.47), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/vmap()/max" stack_frame_id=34} + %broadcast_in_dim.1721 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%max.117), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %sub.452 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1721), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/vmap()/sub" stack_frame_id=34} + %sub.453 = f32[1,1,32,41]{3,2,1,0} reshape(%sub.452), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/vmap()/sub" stack_frame_id=34} + %sub.454 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%sub.453), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/vmap()/sub" stack_frame_id=34} + %sub.455 = f32[1,1,32,41,41]{4,3,2,1,0} subtract(%vmap_jit__where__.107, %sub.454), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/vmap()/sub" stack_frame_id=34} + %exp.113 = f32[1,1,32,41,41]{4,3,2,1,0} exponential(%sub.455), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/vmap()/exp" stack_frame_id=34} + %reduce_sum.1486 = f32[1,1,32,41]{3,2,1,0} reduce(%exp.113, %constant.155), dimensions={4}, to_apply=%region_51.56, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/vmap()/reduce_sum" stack_frame_id=34} + %broadcast_in_dim.1722 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%reduce_sum.1486), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %div.1094 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1722), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/vmap()/div" stack_frame_id=34} + %div.1095 = f32[1,1,32,41]{3,2,1,0} reshape(%div.1094), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/vmap()/div" stack_frame_id=34} + %div.1096 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%div.1095), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/vmap()/div" stack_frame_id=34} + %div.1097 = f32[1,1,32,41,41]{4,3,2,1,0} divide(%exp.113, %div.1096), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/vmap()/div" stack_frame_id=34} + %convert_element_type.1934 = bf16[1,1,32,41,41]{4,3,2,1,0} convert(%div.1097), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/vmap()/convert_element_type" stack_frame_id=34} + %dot_general.1188 = bf16[1,32,128,1,41]{4,3,2,1,0} dot(%reshape.799, %convert_element_type.1934), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=34} + %transpose.111 = bf16[1,41,32,1,128]{1,3,4,2,0} transpose(%dot_general.1188), dimensions={0,4,1,3,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/vmap()/transpose;jit(compiled_generate_function)/transformer_layer_12/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %reshape.801 = bf16[1,41,32,128]{3,2,1,0} reshape(%transpose.111), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/reshape" stack_frame_id=34} + %state_1__113_.1 = bf16[32,128,4096]{2,1,0} parameter(116), metadata={op_name="state[1][113]"} + %dot_general.1189 = bf16[1,41,4096]{2,1,0} dot(%reshape.801, %state_1__113_.1), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=741} + %add.1211 = bf16[1,41,4096]{2,1,0} add(%dot_general.1189, %add.1206), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/add" stack_frame_id=742} + %convert_element_type.1935 = f32[1,41,4096]{2,1,0} convert(%add.1211), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/feedforward_layernorm/convert_element_type" stack_frame_id=746} + %pow.435 = f32[1,41,4096]{2,1,0} power(%convert_element_type.1935, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/feedforward_layernorm/pow" stack_frame_id=749} + %reduce_sum.1487 = f32[1,41]{1,0} reduce(%pow.435, %constant.155), dimensions={2}, to_apply=%region_52.57, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/feedforward_layernorm/reduce_sum" stack_frame_id=43} + %broadcast_in_dim.1723 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1487), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/feedforward_layernorm/broadcast_in_dim" stack_frame_id=43} + %div.1098 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1723, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/feedforward_layernorm/div" stack_frame_id=43} + %add.1212 = f32[1,41,1]{2,1,0} add(%div.1098, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/feedforward_layernorm/add" stack_frame_id=750} + %rsqrt.218 = f32[1,41,1]{2,1,0} rsqrt(%add.1212), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/feedforward_layernorm/rsqrt" stack_frame_id=753} + %mul.3534 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.218), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3535 = f32[1,41]{1,0} reshape(%mul.3534), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3536 = f32[1,41,4096]{2,1,0} broadcast(%mul.3535), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3537 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.1935, %mul.3536), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/feedforward_layernorm/mul" stack_frame_id=754} + %state_1__118_.1 = bf16[4096]{0} parameter(121), metadata={op_name="state[1][118]"} + %convert_element_type.1936 = f32[4096]{0} convert(%state_1__118_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/feedforward_layernorm/convert_element_type" stack_frame_id=755} + %broadcast_in_dim.1724 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1936), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/feedforward_layernorm/broadcast_in_dim" stack_frame_id=755} + %mul.3538 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1724), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3539 = f32[1,4096]{1,0} reshape(%mul.3538), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3540 = f32[1,41,4096]{2,1,0} broadcast(%mul.3539), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3541 = f32[1,41,4096]{2,1,0} multiply(%mul.3537, %mul.3540), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/feedforward_layernorm/mul" stack_frame_id=755} + %convert_element_type.1937 = bf16[1,41,4096]{2,1,0} convert(%mul.3541), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/feedforward_layernorm/convert_element_type" stack_frame_id=759} + %state_1__115_.1 = bf16[4096,14336]{1,0} parameter(118), metadata={op_name="state[1][115]"} + %dot_general.1191 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.1937, %state_1__115_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/feedforward_intermediate_dense/dot_general" stack_frame_id=788} + %state_1__116_.1 = bf16[4096,14336]{1,0} parameter(119), metadata={op_name="state[1][116]"} + %dot_general.1190 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.1937, %state_1__116_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/feedforward_gate_dense/dot_general" stack_frame_id=768} + %convert_element_type.1938 = f32[1,41,14336]{2,1,0} convert(%dot_general.1190), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/convert_element_type" stack_frame_id=772} + %jit_silu_.107 = f32[1,41,14336]{2,1,0} call(%convert_element_type.1938), to_apply=%silu.9, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/jit(silu)" stack_frame_id=775} + %convert_element_type.1939 = bf16[1,41,14336]{2,1,0} convert(%jit_silu_.107), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/convert_element_type" stack_frame_id=779} + %mul.3542 = bf16[1,41,14336]{2,1,0} multiply(%dot_general.1191, %convert_element_type.1939), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/mul" stack_frame_id=791} + %state_1__117_.1 = bf16[14336,4096]{1,0} parameter(120), metadata={op_name="state[1][117]"} + %dot_general.1192 = bf16[1,41,4096]{2,1,0} dot(%mul.3542, %state_1__117_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/feedforward_output_dense/dot_general" stack_frame_id=800} + %add.1213 = bf16[1,41,4096]{2,1,0} add(%dot_general.1192, %add.1211), metadata={op_name="jit(compiled_generate_function)/transformer_layer_12/add" stack_frame_id=801} + %convert_element_type.1942 = f32[1,41,4096]{2,1,0} convert(%add.1213), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention_layernorm/convert_element_type" stack_frame_id=555} + %pow.436 = f32[1,41,4096]{2,1,0} power(%convert_element_type.1942, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention_layernorm/pow" stack_frame_id=558} + %reduce_sum.1488 = f32[1,41]{1,0} reduce(%pow.436, %constant.155), dimensions={2}, to_apply=%region_53.58, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention_layernorm/reduce_sum" stack_frame_id=24} + %broadcast_in_dim.1731 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1488), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention_layernorm/broadcast_in_dim" stack_frame_id=24} + %div.1099 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1731, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention_layernorm/div" stack_frame_id=24} + %add.1215 = f32[1,41,1]{2,1,0} add(%div.1099, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention_layernorm/add" stack_frame_id=559} + %rsqrt.219 = f32[1,41,1]{2,1,0} rsqrt(%add.1215), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention_layernorm/rsqrt" stack_frame_id=562} + %mul.3543 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.219), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3544 = f32[1,41]{1,0} reshape(%mul.3543), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3545 = f32[1,41,4096]{2,1,0} broadcast(%mul.3544), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3546 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.1942, %mul.3545), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention_layernorm/mul" stack_frame_id=563} + %state_1__123_.1 = bf16[4096]{0} parameter(126), metadata={op_name="state[1][123]"} + %convert_element_type.1943 = f32[4096]{0} convert(%state_1__123_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention_layernorm/convert_element_type" stack_frame_id=564} + %broadcast_in_dim.1732 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1943), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention_layernorm/broadcast_in_dim" stack_frame_id=564} + %mul.3547 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1732), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3548 = f32[1,4096]{1,0} reshape(%mul.3547), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3549 = f32[1,41,4096]{2,1,0} broadcast(%mul.3548), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3550 = f32[1,41,4096]{2,1,0} multiply(%mul.3546, %mul.3549), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention_layernorm/mul" stack_frame_id=564} + %convert_element_type.1944 = bf16[1,41,4096]{2,1,0} convert(%mul.3550), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention_layernorm/convert_element_type" stack_frame_id=568} + %state_1__120_.1 = bf16[4096,8,128]{2,1,0} parameter(123), metadata={op_name="state[1][120]"} + %dot_general.1195 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.1944, %state_1__120_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=647} + %iota.621 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/iota" stack_frame_id=677} + %broadcast_in_dim.1734 = f32[1,41]{1,0} reshape(%iota.621), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/broadcast_in_dim" stack_frame_id=680} + %iota.620 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/iota" stack_frame_id=667} + %mul.3560 = f32[64]{0} multiply(%iota.620, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=667} + %div.1102 = f32[64]{0} divide(%mul.3560, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/div" stack_frame_id=668} + %neg.444 = f32[64]{0} negate(%div.1102), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/neg" stack_frame_id=669} + %pow.438 = f32[64]{0} power(%broadcast.51, %neg.444), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/pow" stack_frame_id=672} + %div.1103 = f32[64]{0} divide(%pow.438, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/div" stack_frame_id=673} + %dot_general.1197 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1734, %div.1103), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/bi,j->bij/dot_general" stack_frame_id=683} + %stack.1734 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1197), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/stack" stack_frame_id=686} + %stack.1735 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1197), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/stack" stack_frame_id=686} + %stack.1736 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1734, %stack.1735), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/stack" stack_frame_id=686} + %reshape.804 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1736), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/reshape"} + %cos.218 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.804), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/cos" stack_frame_id=690} + %convert_element_type.1947 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.218), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/convert_element_type" stack_frame_id=694} + %mul.3561 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1947), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=714} + %mul.3562 = bf16[1,41,128]{2,1,0} reshape(%mul.3561), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=714} + %mul.3563 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3562), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=714} + %mul.3564 = bf16[1,41,8,128]{3,2,1,0} multiply(%dot_general.1195, %mul.3563), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=714} + %split.437 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1195), slice={[0:1], [0:41], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/split" stack_frame_id=706} + %neg.445 = bf16[1,41,8,64]{3,2,1,0} negate(%split.437), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/neg" stack_frame_id=707} + %stack.1737 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%neg.445), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/stack" stack_frame_id=710} + %split.436 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1195), slice={[0:1], [0:41], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/split" stack_frame_id=706} + %stack.1738 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%split.436), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/stack" stack_frame_id=710} + %stack.1739 = bf16[1,41,8,2,64]{4,3,2,1,0} concatenate(%stack.1737, %stack.1738), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/stack" stack_frame_id=710} + %reshape.805 = bf16[1,41,8,128]{3,2,1,0} reshape(%stack.1739), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/reshape" stack_frame_id=713} + %sin.218 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.804), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/sin" stack_frame_id=698} + %convert_element_type.1948 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.218), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/convert_element_type" stack_frame_id=702} + %mul.3565 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1948), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=715} + %mul.3566 = bf16[1,41,128]{2,1,0} reshape(%mul.3565), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=715} + %mul.3567 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3566), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=715} + %mul.3568 = bf16[1,41,8,128]{3,2,1,0} multiply(%reshape.805, %mul.3567), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=715} + %add.1217 = bf16[1,41,8,128]{3,2,1,0} add(%mul.3564, %mul.3568), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/add" stack_frame_id=716} + %stack.1740 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%add.1217), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/stack" stack_frame_id=719} + %state_1__121_.1 = bf16[4096,8,128]{2,1,0} parameter(124), metadata={op_name="state[1][121]"} + %dot_general.1196 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.1944, %state_1__121_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=656} + %stack.1741 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dot_general.1196), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/stack" stack_frame_id=719} + %stack.1742 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1740, %stack.1741), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/stack" stack_frame_id=719} + %stack.2020 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1742), metadata={op_name="jit(compiled_generate_function)/stack" stack_frame_id=804} + %broadcast_in_dim.1736 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dot_general.1196), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/broadcast_in_dim" stack_frame_id=725} + %reshape.807 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1736), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/reshape" stack_frame_id=725} + %iota.614 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/iota" stack_frame_id=525} + %broadcast_in_dim.1725 = f32[41,1]{1,0} reshape(%iota.614), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/broadcast_in_dim" stack_frame_id=528} + %ge.572 = f32[41,1]{1,0} broadcast(%broadcast_in_dim.1725), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/ge" stack_frame_id=532} + %ge.573 = f32[41]{0} reshape(%ge.572), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/ge" stack_frame_id=532} + %ge.574 = f32[41,41]{1,0} broadcast(%ge.573), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/ge" stack_frame_id=532} + %iota.615 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/iota" stack_frame_id=531} + %broadcast_in_dim.1726 = f32[1,41]{1,0} reshape(%iota.615), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/broadcast_in_dim" stack_frame_id=532} + %ge.575 = f32[1,41]{1,0} broadcast(%broadcast_in_dim.1726), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/ge" stack_frame_id=532} + %ge.576 = f32[41]{0} reshape(%ge.575), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/ge" stack_frame_id=532} + %ge.577 = f32[41,41]{1,0} broadcast(%ge.576), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/ge" stack_frame_id=532} + %ge.578 = pred[41,41]{1,0} compare(%ge.574, %ge.577), direction=GE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/ge" stack_frame_id=532} + %broadcast_in_dim.1727 = pred[1,41,41]{2,1,0} reshape(%ge.578), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/broadcast_in_dim" stack_frame_id=535} + %convert_element_type.1941 = s32[1,41,41]{2,1,0} convert(%broadcast_in_dim.1727), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/convert_element_type" stack_frame_id=551} + %iota.616 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/iota" stack_frame_id=538} + %broadcast_in_dim.1728 = s32[41,1]{1,0} reshape(%iota.616), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/broadcast_in_dim" stack_frame_id=536} + %lt.714 = s32[41,1]{1,0} broadcast(%broadcast_in_dim.1728), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/lt" stack_frame_id=543} + %lt.715 = s32[41]{0} reshape(%lt.714), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/lt" stack_frame_id=543} + %lt.716 = s32[41,41]{1,0} broadcast(%lt.715), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/lt" stack_frame_id=543} + %iota.617 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/iota" stack_frame_id=541} + %broadcast_in_dim.1729 = s32[1,41]{1,0} reshape(%iota.617), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/broadcast_in_dim" stack_frame_id=539} + %add.1214 = s32[1,41]{1,0} add(%broadcast_in_dim.1729, %broadcast.57), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/add" stack_frame_id=542} + %lt.717 = s32[1,41]{1,0} broadcast(%add.1214), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/lt" stack_frame_id=543} + %lt.718 = s32[41]{0} reshape(%lt.717), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/lt" stack_frame_id=543} + %lt.719 = s32[41,41]{1,0} broadcast(%lt.718), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/lt" stack_frame_id=543} + %lt.720 = pred[41,41]{1,0} compare(%lt.716, %lt.719), direction=LT, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/lt" stack_frame_id=543} + %convert_element_type.1940 = s32[41,41]{1,0} convert(%lt.720), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/convert_element_type" stack_frame_id=547} + %broadcast_in_dim.1730 = s32[1,41,41]{2,1,0} reshape(%convert_element_type.1940), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/broadcast_in_dim" stack_frame_id=551} + %min.108 = s32[1,41,41]{2,1,0} minimum(%convert_element_type.1941, %broadcast_in_dim.1730), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/min" stack_frame_id=551} + %broadcast_in_dim.1737 = s32[1,1,41,41]{3,2,1,0} reshape(%min.108), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/broadcast_in_dim" stack_frame_id=728} + %convert_element_type.1949 = pred[1,1,41,41]{3,2,1,0} compare(%broadcast_in_dim.1737, %broadcast.49), direction=NE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/convert_element_type" stack_frame_id=732} + %broadcast_in_dim.1738 = pred[1,1,1,41,41]{4,3,2,1,0} reshape(%convert_element_type.1949), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/broadcast_in_dim" stack_frame_id=34} + %and.337 = pred[1,1,1,41,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1738), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/vmap()/and" stack_frame_id=34} + %and.338 = pred[1,1,41,41]{3,2,1,0} reshape(%and.337), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/vmap()/and" stack_frame_id=34} + %and.339 = pred[1,1,32,41,41]{4,3,2,1,0} broadcast(%and.338), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/vmap()/and" stack_frame_id=34} + %state_1__119_.1 = bf16[4096,32,128]{2,1,0} parameter(122), metadata={op_name="state[1][119]"} + %dot_general.1193 = bf16[1,41,32,128]{3,2,1,0} dot(%convert_element_type.1944, %state_1__119_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=577} + %iota.619 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/iota" stack_frame_id=598} + %broadcast_in_dim.1733 = f32[1,41]{1,0} reshape(%iota.619), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/broadcast_in_dim" stack_frame_id=601} + %iota.618 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/iota" stack_frame_id=588} + %mul.3551 = f32[64]{0} multiply(%iota.618, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=588} + %div.1100 = f32[64]{0} divide(%mul.3551, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/div" stack_frame_id=589} + %neg.442 = f32[64]{0} negate(%div.1100), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/neg" stack_frame_id=590} + %pow.437 = f32[64]{0} power(%broadcast.51, %neg.442), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/pow" stack_frame_id=593} + %div.1101 = f32[64]{0} divide(%pow.437, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/div" stack_frame_id=594} + %dot_general.1194 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1733, %div.1101), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/bi,j->bij/dot_general" stack_frame_id=604} + %stack.1728 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1194), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/stack" stack_frame_id=607} + %stack.1729 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1194), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/stack" stack_frame_id=607} + %stack.1730 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1728, %stack.1729), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/stack" stack_frame_id=607} + %reshape.802 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1730), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/reshape"} + %cos.217 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.802), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/cos" stack_frame_id=611} + %convert_element_type.1945 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.217), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/convert_element_type" stack_frame_id=615} + %mul.3552 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1945), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=635} + %mul.3553 = bf16[1,41,128]{2,1,0} reshape(%mul.3552), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=635} + %mul.3554 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3553), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=635} + %mul.3555 = bf16[1,41,32,128]{3,2,1,0} multiply(%dot_general.1193, %mul.3554), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=635} + %split.435 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1193), slice={[0:1], [0:41], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/split" stack_frame_id=627} + %neg.443 = bf16[1,41,32,64]{3,2,1,0} negate(%split.435), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/neg" stack_frame_id=628} + %stack.1731 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%neg.443), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/stack" stack_frame_id=631} + %split.434 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1193), slice={[0:1], [0:41], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/split" stack_frame_id=627} + %stack.1732 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%split.434), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/stack" stack_frame_id=631} + %stack.1733 = bf16[1,41,32,2,64]{4,3,2,1,0} concatenate(%stack.1731, %stack.1732), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/stack" stack_frame_id=631} + %reshape.803 = bf16[1,41,32,128]{3,2,1,0} reshape(%stack.1733), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/reshape" stack_frame_id=634} + %sin.217 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.802), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/sin" stack_frame_id=619} + %convert_element_type.1946 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.217), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/convert_element_type" stack_frame_id=623} + %mul.3556 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1946), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=636} + %mul.3557 = bf16[1,41,128]{2,1,0} reshape(%mul.3556), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=636} + %mul.3558 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3557), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=636} + %mul.3559 = bf16[1,41,32,128]{3,2,1,0} multiply(%reshape.803, %mul.3558), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/mul" stack_frame_id=636} + %add.1216 = bf16[1,41,32,128]{3,2,1,0} add(%mul.3555, %mul.3559), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/rotary_embedding_13/add" stack_frame_id=637} + %reshape.808 = bf16[1,41,32,1,128]{4,3,2,1,0} reshape(%add.1216), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/reshape" stack_frame_id=34} + %broadcast_in_dim.1735 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%add.1217), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/broadcast_in_dim" stack_frame_id=722} + %reshape.806 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1735), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/reshape" stack_frame_id=722} + %dot_general.1198 = f32[1,32,41,1,41]{4,3,2,1,0} dot(%reshape.808, %reshape.806), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=34} + %mul.3569 = f32[1,32,41,1,41]{4,3,2,1,0} multiply(%dot_general.1198, %broadcast.48), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/vmap()/mul" stack_frame_id=34} + %vmap_jit__where__.108 = f32[1,1,32,41,41]{4,3,2,1,0} call(%and.339, %mul.3569, %constant.153), to_apply=%_where_0.5, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/vmap(jit(_where))" stack_frame_id=34} + %reduce_max.494 = f32[1,1,32,41]{3,2,1,0} reduce(%vmap_jit__where__.108, %constant.152), dimensions={4}, to_apply=%region_54.59, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/vmap()/reduce_max" stack_frame_id=34} + %max.118 = f32[1,1,32,41]{3,2,1,0} maximum(%reduce_max.494, %broadcast.47), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/vmap()/max" stack_frame_id=34} + %broadcast_in_dim.1739 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%max.118), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %sub.456 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1739), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/vmap()/sub" stack_frame_id=34} + %sub.457 = f32[1,1,32,41]{3,2,1,0} reshape(%sub.456), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/vmap()/sub" stack_frame_id=34} + %sub.458 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%sub.457), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/vmap()/sub" stack_frame_id=34} + %sub.459 = f32[1,1,32,41,41]{4,3,2,1,0} subtract(%vmap_jit__where__.108, %sub.458), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/vmap()/sub" stack_frame_id=34} + %exp.114 = f32[1,1,32,41,41]{4,3,2,1,0} exponential(%sub.459), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/vmap()/exp" stack_frame_id=34} + %reduce_sum.1489 = f32[1,1,32,41]{3,2,1,0} reduce(%exp.114, %constant.155), dimensions={4}, to_apply=%region_55.60, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/vmap()/reduce_sum" stack_frame_id=34} + %broadcast_in_dim.1740 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%reduce_sum.1489), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %div.1104 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1740), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/vmap()/div" stack_frame_id=34} + %div.1105 = f32[1,1,32,41]{3,2,1,0} reshape(%div.1104), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/vmap()/div" stack_frame_id=34} + %div.1106 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%div.1105), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/vmap()/div" stack_frame_id=34} + %div.1107 = f32[1,1,32,41,41]{4,3,2,1,0} divide(%exp.114, %div.1106), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/vmap()/div" stack_frame_id=34} + %convert_element_type.1950 = bf16[1,1,32,41,41]{4,3,2,1,0} convert(%div.1107), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/vmap()/convert_element_type" stack_frame_id=34} + %dot_general.1199 = bf16[1,32,128,1,41]{4,3,2,1,0} dot(%reshape.807, %convert_element_type.1950), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=34} + %transpose.112 = bf16[1,41,32,1,128]{1,3,4,2,0} transpose(%dot_general.1199), dimensions={0,4,1,3,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/vmap()/transpose;jit(compiled_generate_function)/transformer_layer_13/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %reshape.809 = bf16[1,41,32,128]{3,2,1,0} reshape(%transpose.112), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/reshape" stack_frame_id=34} + %state_1__122_.1 = bf16[32,128,4096]{2,1,0} parameter(125), metadata={op_name="state[1][122]"} + %dot_general.1200 = bf16[1,41,4096]{2,1,0} dot(%reshape.809, %state_1__122_.1), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=741} + %add.1218 = bf16[1,41,4096]{2,1,0} add(%dot_general.1200, %add.1213), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/add" stack_frame_id=742} + %convert_element_type.1951 = f32[1,41,4096]{2,1,0} convert(%add.1218), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/feedforward_layernorm/convert_element_type" stack_frame_id=746} + %pow.439 = f32[1,41,4096]{2,1,0} power(%convert_element_type.1951, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/feedforward_layernorm/pow" stack_frame_id=749} + %reduce_sum.1490 = f32[1,41]{1,0} reduce(%pow.439, %constant.155), dimensions={2}, to_apply=%region_56.61, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/feedforward_layernorm/reduce_sum" stack_frame_id=43} + %broadcast_in_dim.1741 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1490), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/feedforward_layernorm/broadcast_in_dim" stack_frame_id=43} + %div.1108 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1741, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/feedforward_layernorm/div" stack_frame_id=43} + %add.1219 = f32[1,41,1]{2,1,0} add(%div.1108, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/feedforward_layernorm/add" stack_frame_id=750} + %rsqrt.220 = f32[1,41,1]{2,1,0} rsqrt(%add.1219), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/feedforward_layernorm/rsqrt" stack_frame_id=753} + %mul.3570 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.220), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3571 = f32[1,41]{1,0} reshape(%mul.3570), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3572 = f32[1,41,4096]{2,1,0} broadcast(%mul.3571), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3573 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.1951, %mul.3572), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/feedforward_layernorm/mul" stack_frame_id=754} + %state_1__127_.1 = bf16[4096]{0} parameter(130), metadata={op_name="state[1][127]"} + %convert_element_type.1952 = f32[4096]{0} convert(%state_1__127_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/feedforward_layernorm/convert_element_type" stack_frame_id=755} + %broadcast_in_dim.1742 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1952), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/feedforward_layernorm/broadcast_in_dim" stack_frame_id=755} + %mul.3574 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1742), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3575 = f32[1,4096]{1,0} reshape(%mul.3574), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3576 = f32[1,41,4096]{2,1,0} broadcast(%mul.3575), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3577 = f32[1,41,4096]{2,1,0} multiply(%mul.3573, %mul.3576), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/feedforward_layernorm/mul" stack_frame_id=755} + %convert_element_type.1953 = bf16[1,41,4096]{2,1,0} convert(%mul.3577), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/feedforward_layernorm/convert_element_type" stack_frame_id=759} + %state_1__124_.1 = bf16[4096,14336]{1,0} parameter(127), metadata={op_name="state[1][124]"} + %dot_general.1202 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.1953, %state_1__124_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/feedforward_intermediate_dense/dot_general" stack_frame_id=788} + %state_1__125_.1 = bf16[4096,14336]{1,0} parameter(128), metadata={op_name="state[1][125]"} + %dot_general.1201 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.1953, %state_1__125_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/feedforward_gate_dense/dot_general" stack_frame_id=768} + %convert_element_type.1954 = f32[1,41,14336]{2,1,0} convert(%dot_general.1201), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/convert_element_type" stack_frame_id=772} + %jit_silu_.108 = f32[1,41,14336]{2,1,0} call(%convert_element_type.1954), to_apply=%silu.9, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/jit(silu)" stack_frame_id=775} + %convert_element_type.1955 = bf16[1,41,14336]{2,1,0} convert(%jit_silu_.108), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/convert_element_type" stack_frame_id=779} + %mul.3578 = bf16[1,41,14336]{2,1,0} multiply(%dot_general.1202, %convert_element_type.1955), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/mul" stack_frame_id=791} + %state_1__126_.1 = bf16[14336,4096]{1,0} parameter(129), metadata={op_name="state[1][126]"} + %dot_general.1203 = bf16[1,41,4096]{2,1,0} dot(%mul.3578, %state_1__126_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/feedforward_output_dense/dot_general" stack_frame_id=800} + %add.1220 = bf16[1,41,4096]{2,1,0} add(%dot_general.1203, %add.1218), metadata={op_name="jit(compiled_generate_function)/transformer_layer_13/add" stack_frame_id=801} + %convert_element_type.1958 = f32[1,41,4096]{2,1,0} convert(%add.1220), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention_layernorm/convert_element_type" stack_frame_id=555} + %pow.440 = f32[1,41,4096]{2,1,0} power(%convert_element_type.1958, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention_layernorm/pow" stack_frame_id=558} + %reduce_sum.1491 = f32[1,41]{1,0} reduce(%pow.440, %constant.155), dimensions={2}, to_apply=%region_57.62, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention_layernorm/reduce_sum" stack_frame_id=24} + %broadcast_in_dim.1749 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1491), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention_layernorm/broadcast_in_dim" stack_frame_id=24} + %div.1109 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1749, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention_layernorm/div" stack_frame_id=24} + %add.1222 = f32[1,41,1]{2,1,0} add(%div.1109, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention_layernorm/add" stack_frame_id=559} + %rsqrt.221 = f32[1,41,1]{2,1,0} rsqrt(%add.1222), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention_layernorm/rsqrt" stack_frame_id=562} + %mul.3579 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.221), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3580 = f32[1,41]{1,0} reshape(%mul.3579), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3581 = f32[1,41,4096]{2,1,0} broadcast(%mul.3580), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3582 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.1958, %mul.3581), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention_layernorm/mul" stack_frame_id=563} + %state_1__132_.1 = bf16[4096]{0} parameter(135), metadata={op_name="state[1][132]"} + %convert_element_type.1959 = f32[4096]{0} convert(%state_1__132_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention_layernorm/convert_element_type" stack_frame_id=564} + %broadcast_in_dim.1750 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1959), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention_layernorm/broadcast_in_dim" stack_frame_id=564} + %mul.3583 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1750), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3584 = f32[1,4096]{1,0} reshape(%mul.3583), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3585 = f32[1,41,4096]{2,1,0} broadcast(%mul.3584), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3586 = f32[1,41,4096]{2,1,0} multiply(%mul.3582, %mul.3585), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention_layernorm/mul" stack_frame_id=564} + %convert_element_type.1960 = bf16[1,41,4096]{2,1,0} convert(%mul.3586), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention_layernorm/convert_element_type" stack_frame_id=568} + %state_1__129_.1 = bf16[4096,8,128]{2,1,0} parameter(132), metadata={op_name="state[1][129]"} + %dot_general.1206 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.1960, %state_1__129_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=647} + %iota.629 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/iota" stack_frame_id=677} + %broadcast_in_dim.1752 = f32[1,41]{1,0} reshape(%iota.629), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/broadcast_in_dim" stack_frame_id=680} + %iota.628 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/iota" stack_frame_id=667} + %mul.3596 = f32[64]{0} multiply(%iota.628, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=667} + %div.1112 = f32[64]{0} divide(%mul.3596, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/div" stack_frame_id=668} + %neg.448 = f32[64]{0} negate(%div.1112), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/neg" stack_frame_id=669} + %pow.442 = f32[64]{0} power(%broadcast.51, %neg.448), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/pow" stack_frame_id=672} + %div.1113 = f32[64]{0} divide(%pow.442, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/div" stack_frame_id=673} + %dot_general.1208 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1752, %div.1113), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/bi,j->bij/dot_general" stack_frame_id=683} + %stack.1749 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1208), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/stack" stack_frame_id=686} + %stack.1750 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1208), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/stack" stack_frame_id=686} + %stack.1751 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1749, %stack.1750), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/stack" stack_frame_id=686} + %reshape.812 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1751), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/reshape"} + %cos.220 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.812), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/cos" stack_frame_id=690} + %convert_element_type.1963 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.220), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/convert_element_type" stack_frame_id=694} + %mul.3597 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1963), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=714} + %mul.3598 = bf16[1,41,128]{2,1,0} reshape(%mul.3597), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=714} + %mul.3599 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3598), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=714} + %mul.3600 = bf16[1,41,8,128]{3,2,1,0} multiply(%dot_general.1206, %mul.3599), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=714} + %split.441 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1206), slice={[0:1], [0:41], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/split" stack_frame_id=706} + %neg.449 = bf16[1,41,8,64]{3,2,1,0} negate(%split.441), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/neg" stack_frame_id=707} + %stack.1752 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%neg.449), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/stack" stack_frame_id=710} + %split.440 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1206), slice={[0:1], [0:41], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/split" stack_frame_id=706} + %stack.1753 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%split.440), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/stack" stack_frame_id=710} + %stack.1754 = bf16[1,41,8,2,64]{4,3,2,1,0} concatenate(%stack.1752, %stack.1753), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/stack" stack_frame_id=710} + %reshape.813 = bf16[1,41,8,128]{3,2,1,0} reshape(%stack.1754), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/reshape" stack_frame_id=713} + %sin.220 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.812), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/sin" stack_frame_id=698} + %convert_element_type.1964 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.220), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/convert_element_type" stack_frame_id=702} + %mul.3601 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1964), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=715} + %mul.3602 = bf16[1,41,128]{2,1,0} reshape(%mul.3601), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=715} + %mul.3603 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3602), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=715} + %mul.3604 = bf16[1,41,8,128]{3,2,1,0} multiply(%reshape.813, %mul.3603), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=715} + %add.1224 = bf16[1,41,8,128]{3,2,1,0} add(%mul.3600, %mul.3604), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/add" stack_frame_id=716} + %stack.1755 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%add.1224), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/stack" stack_frame_id=719} + %state_1__130_.1 = bf16[4096,8,128]{2,1,0} parameter(133), metadata={op_name="state[1][130]"} + %dot_general.1207 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.1960, %state_1__130_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=656} + %stack.1756 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dot_general.1207), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/stack" stack_frame_id=719} + %stack.1757 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1755, %stack.1756), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/stack" stack_frame_id=719} + %stack.2021 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1757), metadata={op_name="jit(compiled_generate_function)/stack" stack_frame_id=804} + %broadcast_in_dim.1754 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dot_general.1207), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/broadcast_in_dim" stack_frame_id=725} + %reshape.815 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1754), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/reshape" stack_frame_id=725} + %iota.622 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/iota" stack_frame_id=525} + %broadcast_in_dim.1743 = f32[41,1]{1,0} reshape(%iota.622), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/broadcast_in_dim" stack_frame_id=528} + %ge.579 = f32[41,1]{1,0} broadcast(%broadcast_in_dim.1743), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/ge" stack_frame_id=532} + %ge.580 = f32[41]{0} reshape(%ge.579), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/ge" stack_frame_id=532} + %ge.581 = f32[41,41]{1,0} broadcast(%ge.580), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/ge" stack_frame_id=532} + %iota.623 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/iota" stack_frame_id=531} + %broadcast_in_dim.1744 = f32[1,41]{1,0} reshape(%iota.623), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/broadcast_in_dim" stack_frame_id=532} + %ge.582 = f32[1,41]{1,0} broadcast(%broadcast_in_dim.1744), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/ge" stack_frame_id=532} + %ge.583 = f32[41]{0} reshape(%ge.582), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/ge" stack_frame_id=532} + %ge.584 = f32[41,41]{1,0} broadcast(%ge.583), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/ge" stack_frame_id=532} + %ge.585 = pred[41,41]{1,0} compare(%ge.581, %ge.584), direction=GE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/ge" stack_frame_id=532} + %broadcast_in_dim.1745 = pred[1,41,41]{2,1,0} reshape(%ge.585), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/broadcast_in_dim" stack_frame_id=535} + %convert_element_type.1957 = s32[1,41,41]{2,1,0} convert(%broadcast_in_dim.1745), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/convert_element_type" stack_frame_id=551} + %iota.624 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/iota" stack_frame_id=538} + %broadcast_in_dim.1746 = s32[41,1]{1,0} reshape(%iota.624), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/broadcast_in_dim" stack_frame_id=536} + %lt.721 = s32[41,1]{1,0} broadcast(%broadcast_in_dim.1746), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/lt" stack_frame_id=543} + %lt.722 = s32[41]{0} reshape(%lt.721), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/lt" stack_frame_id=543} + %lt.723 = s32[41,41]{1,0} broadcast(%lt.722), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/lt" stack_frame_id=543} + %iota.625 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/iota" stack_frame_id=541} + %broadcast_in_dim.1747 = s32[1,41]{1,0} reshape(%iota.625), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/broadcast_in_dim" stack_frame_id=539} + %add.1221 = s32[1,41]{1,0} add(%broadcast_in_dim.1747, %broadcast.57), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/add" stack_frame_id=542} + %lt.724 = s32[1,41]{1,0} broadcast(%add.1221), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/lt" stack_frame_id=543} + %lt.725 = s32[41]{0} reshape(%lt.724), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/lt" stack_frame_id=543} + %lt.726 = s32[41,41]{1,0} broadcast(%lt.725), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/lt" stack_frame_id=543} + %lt.727 = pred[41,41]{1,0} compare(%lt.723, %lt.726), direction=LT, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/lt" stack_frame_id=543} + %convert_element_type.1956 = s32[41,41]{1,0} convert(%lt.727), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/convert_element_type" stack_frame_id=547} + %broadcast_in_dim.1748 = s32[1,41,41]{2,1,0} reshape(%convert_element_type.1956), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/broadcast_in_dim" stack_frame_id=551} + %min.109 = s32[1,41,41]{2,1,0} minimum(%convert_element_type.1957, %broadcast_in_dim.1748), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/min" stack_frame_id=551} + %broadcast_in_dim.1755 = s32[1,1,41,41]{3,2,1,0} reshape(%min.109), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/broadcast_in_dim" stack_frame_id=728} + %convert_element_type.1965 = pred[1,1,41,41]{3,2,1,0} compare(%broadcast_in_dim.1755, %broadcast.49), direction=NE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/convert_element_type" stack_frame_id=732} + %broadcast_in_dim.1756 = pred[1,1,1,41,41]{4,3,2,1,0} reshape(%convert_element_type.1965), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/broadcast_in_dim" stack_frame_id=34} + %and.340 = pred[1,1,1,41,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1756), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/vmap()/and" stack_frame_id=34} + %and.341 = pred[1,1,41,41]{3,2,1,0} reshape(%and.340), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/vmap()/and" stack_frame_id=34} + %and.342 = pred[1,1,32,41,41]{4,3,2,1,0} broadcast(%and.341), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/vmap()/and" stack_frame_id=34} + %state_1__128_.1 = bf16[4096,32,128]{2,1,0} parameter(131), metadata={op_name="state[1][128]"} + %dot_general.1204 = bf16[1,41,32,128]{3,2,1,0} dot(%convert_element_type.1960, %state_1__128_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=577} + %iota.627 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/iota" stack_frame_id=598} + %broadcast_in_dim.1751 = f32[1,41]{1,0} reshape(%iota.627), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/broadcast_in_dim" stack_frame_id=601} + %iota.626 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/iota" stack_frame_id=588} + %mul.3587 = f32[64]{0} multiply(%iota.626, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=588} + %div.1110 = f32[64]{0} divide(%mul.3587, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/div" stack_frame_id=589} + %neg.446 = f32[64]{0} negate(%div.1110), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/neg" stack_frame_id=590} + %pow.441 = f32[64]{0} power(%broadcast.51, %neg.446), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/pow" stack_frame_id=593} + %div.1111 = f32[64]{0} divide(%pow.441, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/div" stack_frame_id=594} + %dot_general.1205 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1751, %div.1111), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/bi,j->bij/dot_general" stack_frame_id=604} + %stack.1743 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1205), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/stack" stack_frame_id=607} + %stack.1744 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1205), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/stack" stack_frame_id=607} + %stack.1745 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1743, %stack.1744), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/stack" stack_frame_id=607} + %reshape.810 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1745), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/reshape"} + %cos.219 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.810), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/cos" stack_frame_id=611} + %convert_element_type.1961 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.219), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/convert_element_type" stack_frame_id=615} + %mul.3588 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1961), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=635} + %mul.3589 = bf16[1,41,128]{2,1,0} reshape(%mul.3588), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=635} + %mul.3590 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3589), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=635} + %mul.3591 = bf16[1,41,32,128]{3,2,1,0} multiply(%dot_general.1204, %mul.3590), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=635} + %split.439 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1204), slice={[0:1], [0:41], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/split" stack_frame_id=627} + %neg.447 = bf16[1,41,32,64]{3,2,1,0} negate(%split.439), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/neg" stack_frame_id=628} + %stack.1746 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%neg.447), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/stack" stack_frame_id=631} + %split.438 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1204), slice={[0:1], [0:41], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/split" stack_frame_id=627} + %stack.1747 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%split.438), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/stack" stack_frame_id=631} + %stack.1748 = bf16[1,41,32,2,64]{4,3,2,1,0} concatenate(%stack.1746, %stack.1747), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/stack" stack_frame_id=631} + %reshape.811 = bf16[1,41,32,128]{3,2,1,0} reshape(%stack.1748), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/reshape" stack_frame_id=634} + %sin.219 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.810), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/sin" stack_frame_id=619} + %convert_element_type.1962 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.219), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/convert_element_type" stack_frame_id=623} + %mul.3592 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1962), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=636} + %mul.3593 = bf16[1,41,128]{2,1,0} reshape(%mul.3592), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=636} + %mul.3594 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3593), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=636} + %mul.3595 = bf16[1,41,32,128]{3,2,1,0} multiply(%reshape.811, %mul.3594), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/mul" stack_frame_id=636} + %add.1223 = bf16[1,41,32,128]{3,2,1,0} add(%mul.3591, %mul.3595), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/rotary_embedding_14/add" stack_frame_id=637} + %reshape.816 = bf16[1,41,32,1,128]{4,3,2,1,0} reshape(%add.1223), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/reshape" stack_frame_id=34} + %broadcast_in_dim.1753 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%add.1224), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/broadcast_in_dim" stack_frame_id=722} + %reshape.814 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1753), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/reshape" stack_frame_id=722} + %dot_general.1209 = f32[1,32,41,1,41]{4,3,2,1,0} dot(%reshape.816, %reshape.814), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=34} + %mul.3605 = f32[1,32,41,1,41]{4,3,2,1,0} multiply(%dot_general.1209, %broadcast.48), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/vmap()/mul" stack_frame_id=34} + %vmap_jit__where__.109 = f32[1,1,32,41,41]{4,3,2,1,0} call(%and.342, %mul.3605, %constant.153), to_apply=%_where_0.5, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/vmap(jit(_where))" stack_frame_id=34} + %reduce_max.495 = f32[1,1,32,41]{3,2,1,0} reduce(%vmap_jit__where__.109, %constant.152), dimensions={4}, to_apply=%region_58.63, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/vmap()/reduce_max" stack_frame_id=34} + %max.119 = f32[1,1,32,41]{3,2,1,0} maximum(%reduce_max.495, %broadcast.47), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/vmap()/max" stack_frame_id=34} + %broadcast_in_dim.1757 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%max.119), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %sub.460 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1757), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/vmap()/sub" stack_frame_id=34} + %sub.461 = f32[1,1,32,41]{3,2,1,0} reshape(%sub.460), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/vmap()/sub" stack_frame_id=34} + %sub.462 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%sub.461), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/vmap()/sub" stack_frame_id=34} + %sub.463 = f32[1,1,32,41,41]{4,3,2,1,0} subtract(%vmap_jit__where__.109, %sub.462), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/vmap()/sub" stack_frame_id=34} + %exp.115 = f32[1,1,32,41,41]{4,3,2,1,0} exponential(%sub.463), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/vmap()/exp" stack_frame_id=34} + %reduce_sum.1492 = f32[1,1,32,41]{3,2,1,0} reduce(%exp.115, %constant.155), dimensions={4}, to_apply=%region_59.64, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/vmap()/reduce_sum" stack_frame_id=34} + %broadcast_in_dim.1758 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%reduce_sum.1492), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %div.1114 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1758), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/vmap()/div" stack_frame_id=34} + %div.1115 = f32[1,1,32,41]{3,2,1,0} reshape(%div.1114), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/vmap()/div" stack_frame_id=34} + %div.1116 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%div.1115), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/vmap()/div" stack_frame_id=34} + %div.1117 = f32[1,1,32,41,41]{4,3,2,1,0} divide(%exp.115, %div.1116), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/vmap()/div" stack_frame_id=34} + %convert_element_type.1966 = bf16[1,1,32,41,41]{4,3,2,1,0} convert(%div.1117), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/vmap()/convert_element_type" stack_frame_id=34} + %dot_general.1210 = bf16[1,32,128,1,41]{4,3,2,1,0} dot(%reshape.815, %convert_element_type.1966), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=34} + %transpose.113 = bf16[1,41,32,1,128]{1,3,4,2,0} transpose(%dot_general.1210), dimensions={0,4,1,3,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/vmap()/transpose;jit(compiled_generate_function)/transformer_layer_14/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %reshape.817 = bf16[1,41,32,128]{3,2,1,0} reshape(%transpose.113), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/reshape" stack_frame_id=34} + %state_1__131_.1 = bf16[32,128,4096]{2,1,0} parameter(134), metadata={op_name="state[1][131]"} + %dot_general.1211 = bf16[1,41,4096]{2,1,0} dot(%reshape.817, %state_1__131_.1), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=741} + %add.1225 = bf16[1,41,4096]{2,1,0} add(%dot_general.1211, %add.1220), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/add" stack_frame_id=742} + %convert_element_type.1967 = f32[1,41,4096]{2,1,0} convert(%add.1225), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/feedforward_layernorm/convert_element_type" stack_frame_id=746} + %pow.443 = f32[1,41,4096]{2,1,0} power(%convert_element_type.1967, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/feedforward_layernorm/pow" stack_frame_id=749} + %reduce_sum.1493 = f32[1,41]{1,0} reduce(%pow.443, %constant.155), dimensions={2}, to_apply=%region_60.65, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/feedforward_layernorm/reduce_sum" stack_frame_id=43} + %broadcast_in_dim.1759 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1493), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/feedforward_layernorm/broadcast_in_dim" stack_frame_id=43} + %div.1118 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1759, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/feedforward_layernorm/div" stack_frame_id=43} + %add.1226 = f32[1,41,1]{2,1,0} add(%div.1118, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/feedforward_layernorm/add" stack_frame_id=750} + %rsqrt.222 = f32[1,41,1]{2,1,0} rsqrt(%add.1226), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/feedforward_layernorm/rsqrt" stack_frame_id=753} + %mul.3606 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.222), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3607 = f32[1,41]{1,0} reshape(%mul.3606), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3608 = f32[1,41,4096]{2,1,0} broadcast(%mul.3607), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3609 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.1967, %mul.3608), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/feedforward_layernorm/mul" stack_frame_id=754} + %state_1__136_.1 = bf16[4096]{0} parameter(139), metadata={op_name="state[1][136]"} + %convert_element_type.1968 = f32[4096]{0} convert(%state_1__136_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/feedforward_layernorm/convert_element_type" stack_frame_id=755} + %broadcast_in_dim.1760 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1968), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/feedforward_layernorm/broadcast_in_dim" stack_frame_id=755} + %mul.3610 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1760), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3611 = f32[1,4096]{1,0} reshape(%mul.3610), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3612 = f32[1,41,4096]{2,1,0} broadcast(%mul.3611), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3613 = f32[1,41,4096]{2,1,0} multiply(%mul.3609, %mul.3612), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/feedforward_layernorm/mul" stack_frame_id=755} + %convert_element_type.1969 = bf16[1,41,4096]{2,1,0} convert(%mul.3613), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/feedforward_layernorm/convert_element_type" stack_frame_id=759} + %state_1__133_.1 = bf16[4096,14336]{1,0} parameter(136), metadata={op_name="state[1][133]"} + %dot_general.1213 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.1969, %state_1__133_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/feedforward_intermediate_dense/dot_general" stack_frame_id=788} + %state_1__134_.1 = bf16[4096,14336]{1,0} parameter(137), metadata={op_name="state[1][134]"} + %dot_general.1212 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.1969, %state_1__134_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/feedforward_gate_dense/dot_general" stack_frame_id=768} + %convert_element_type.1970 = f32[1,41,14336]{2,1,0} convert(%dot_general.1212), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/convert_element_type" stack_frame_id=772} + %jit_silu_.109 = f32[1,41,14336]{2,1,0} call(%convert_element_type.1970), to_apply=%silu.9, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/jit(silu)" stack_frame_id=775} + %convert_element_type.1971 = bf16[1,41,14336]{2,1,0} convert(%jit_silu_.109), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/convert_element_type" stack_frame_id=779} + %mul.3614 = bf16[1,41,14336]{2,1,0} multiply(%dot_general.1213, %convert_element_type.1971), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/mul" stack_frame_id=791} + %state_1__135_.1 = bf16[14336,4096]{1,0} parameter(138), metadata={op_name="state[1][135]"} + %dot_general.1214 = bf16[1,41,4096]{2,1,0} dot(%mul.3614, %state_1__135_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/feedforward_output_dense/dot_general" stack_frame_id=800} + %add.1227 = bf16[1,41,4096]{2,1,0} add(%dot_general.1214, %add.1225), metadata={op_name="jit(compiled_generate_function)/transformer_layer_14/add" stack_frame_id=801} + %convert_element_type.1974 = f32[1,41,4096]{2,1,0} convert(%add.1227), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention_layernorm/convert_element_type" stack_frame_id=555} + %pow.444 = f32[1,41,4096]{2,1,0} power(%convert_element_type.1974, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention_layernorm/pow" stack_frame_id=558} + %reduce_sum.1494 = f32[1,41]{1,0} reduce(%pow.444, %constant.155), dimensions={2}, to_apply=%region_61.66, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention_layernorm/reduce_sum" stack_frame_id=24} + %broadcast_in_dim.1767 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1494), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention_layernorm/broadcast_in_dim" stack_frame_id=24} + %div.1119 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1767, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention_layernorm/div" stack_frame_id=24} + %add.1229 = f32[1,41,1]{2,1,0} add(%div.1119, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention_layernorm/add" stack_frame_id=559} + %rsqrt.223 = f32[1,41,1]{2,1,0} rsqrt(%add.1229), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention_layernorm/rsqrt" stack_frame_id=562} + %mul.3615 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.223), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3616 = f32[1,41]{1,0} reshape(%mul.3615), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3617 = f32[1,41,4096]{2,1,0} broadcast(%mul.3616), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3618 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.1974, %mul.3617), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention_layernorm/mul" stack_frame_id=563} + %state_1__141_.1 = bf16[4096]{0} parameter(144), metadata={op_name="state[1][141]"} + %convert_element_type.1975 = f32[4096]{0} convert(%state_1__141_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention_layernorm/convert_element_type" stack_frame_id=564} + %broadcast_in_dim.1768 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1975), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention_layernorm/broadcast_in_dim" stack_frame_id=564} + %mul.3619 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1768), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3620 = f32[1,4096]{1,0} reshape(%mul.3619), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3621 = f32[1,41,4096]{2,1,0} broadcast(%mul.3620), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3622 = f32[1,41,4096]{2,1,0} multiply(%mul.3618, %mul.3621), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention_layernorm/mul" stack_frame_id=564} + %convert_element_type.1976 = bf16[1,41,4096]{2,1,0} convert(%mul.3622), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention_layernorm/convert_element_type" stack_frame_id=568} + %state_1__138_.1 = bf16[4096,8,128]{2,1,0} parameter(141), metadata={op_name="state[1][138]"} + %dot_general.1217 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.1976, %state_1__138_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=647} + %iota.637 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/iota" stack_frame_id=677} + %broadcast_in_dim.1770 = f32[1,41]{1,0} reshape(%iota.637), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/broadcast_in_dim" stack_frame_id=680} + %iota.636 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/iota" stack_frame_id=667} + %mul.3632 = f32[64]{0} multiply(%iota.636, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=667} + %div.1122 = f32[64]{0} divide(%mul.3632, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/div" stack_frame_id=668} + %neg.452 = f32[64]{0} negate(%div.1122), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/neg" stack_frame_id=669} + %pow.446 = f32[64]{0} power(%broadcast.51, %neg.452), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/pow" stack_frame_id=672} + %div.1123 = f32[64]{0} divide(%pow.446, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/div" stack_frame_id=673} + %dot_general.1219 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1770, %div.1123), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/bi,j->bij/dot_general" stack_frame_id=683} + %stack.1764 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1219), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/stack" stack_frame_id=686} + %stack.1765 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1219), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/stack" stack_frame_id=686} + %stack.1766 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1764, %stack.1765), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/stack" stack_frame_id=686} + %reshape.820 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1766), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/reshape"} + %cos.222 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.820), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/cos" stack_frame_id=690} + %convert_element_type.1979 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.222), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/convert_element_type" stack_frame_id=694} + %mul.3633 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1979), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=714} + %mul.3634 = bf16[1,41,128]{2,1,0} reshape(%mul.3633), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=714} + %mul.3635 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3634), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=714} + %mul.3636 = bf16[1,41,8,128]{3,2,1,0} multiply(%dot_general.1217, %mul.3635), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=714} + %split.445 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1217), slice={[0:1], [0:41], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/split" stack_frame_id=706} + %neg.453 = bf16[1,41,8,64]{3,2,1,0} negate(%split.445), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/neg" stack_frame_id=707} + %stack.1767 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%neg.453), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/stack" stack_frame_id=710} + %split.444 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1217), slice={[0:1], [0:41], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/split" stack_frame_id=706} + %stack.1768 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%split.444), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/stack" stack_frame_id=710} + %stack.1769 = bf16[1,41,8,2,64]{4,3,2,1,0} concatenate(%stack.1767, %stack.1768), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/stack" stack_frame_id=710} + %reshape.821 = bf16[1,41,8,128]{3,2,1,0} reshape(%stack.1769), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/reshape" stack_frame_id=713} + %sin.222 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.820), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/sin" stack_frame_id=698} + %convert_element_type.1980 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.222), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/convert_element_type" stack_frame_id=702} + %mul.3637 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1980), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=715} + %mul.3638 = bf16[1,41,128]{2,1,0} reshape(%mul.3637), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=715} + %mul.3639 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3638), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=715} + %mul.3640 = bf16[1,41,8,128]{3,2,1,0} multiply(%reshape.821, %mul.3639), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=715} + %add.1231 = bf16[1,41,8,128]{3,2,1,0} add(%mul.3636, %mul.3640), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/add" stack_frame_id=716} + %stack.1770 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%add.1231), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/stack" stack_frame_id=719} + %state_1__139_.1 = bf16[4096,8,128]{2,1,0} parameter(142), metadata={op_name="state[1][139]"} + %dot_general.1218 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.1976, %state_1__139_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=656} + %stack.1771 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dot_general.1218), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/stack" stack_frame_id=719} + %stack.1772 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1770, %stack.1771), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/stack" stack_frame_id=719} + %stack.2022 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1772), metadata={op_name="jit(compiled_generate_function)/stack" stack_frame_id=804} + %broadcast_in_dim.1772 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dot_general.1218), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/broadcast_in_dim" stack_frame_id=725} + %reshape.823 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1772), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/reshape" stack_frame_id=725} + %iota.630 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/iota" stack_frame_id=525} + %broadcast_in_dim.1761 = f32[41,1]{1,0} reshape(%iota.630), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/broadcast_in_dim" stack_frame_id=528} + %ge.586 = f32[41,1]{1,0} broadcast(%broadcast_in_dim.1761), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/ge" stack_frame_id=532} + %ge.587 = f32[41]{0} reshape(%ge.586), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/ge" stack_frame_id=532} + %ge.588 = f32[41,41]{1,0} broadcast(%ge.587), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/ge" stack_frame_id=532} + %iota.631 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/iota" stack_frame_id=531} + %broadcast_in_dim.1762 = f32[1,41]{1,0} reshape(%iota.631), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/broadcast_in_dim" stack_frame_id=532} + %ge.589 = f32[1,41]{1,0} broadcast(%broadcast_in_dim.1762), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/ge" stack_frame_id=532} + %ge.590 = f32[41]{0} reshape(%ge.589), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/ge" stack_frame_id=532} + %ge.591 = f32[41,41]{1,0} broadcast(%ge.590), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/ge" stack_frame_id=532} + %ge.592 = pred[41,41]{1,0} compare(%ge.588, %ge.591), direction=GE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/ge" stack_frame_id=532} + %broadcast_in_dim.1763 = pred[1,41,41]{2,1,0} reshape(%ge.592), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/broadcast_in_dim" stack_frame_id=535} + %convert_element_type.1973 = s32[1,41,41]{2,1,0} convert(%broadcast_in_dim.1763), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/convert_element_type" stack_frame_id=551} + %iota.632 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/iota" stack_frame_id=538} + %broadcast_in_dim.1764 = s32[41,1]{1,0} reshape(%iota.632), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/broadcast_in_dim" stack_frame_id=536} + %lt.728 = s32[41,1]{1,0} broadcast(%broadcast_in_dim.1764), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/lt" stack_frame_id=543} + %lt.729 = s32[41]{0} reshape(%lt.728), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/lt" stack_frame_id=543} + %lt.730 = s32[41,41]{1,0} broadcast(%lt.729), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/lt" stack_frame_id=543} + %iota.633 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/iota" stack_frame_id=541} + %broadcast_in_dim.1765 = s32[1,41]{1,0} reshape(%iota.633), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/broadcast_in_dim" stack_frame_id=539} + %add.1228 = s32[1,41]{1,0} add(%broadcast_in_dim.1765, %broadcast.57), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/add" stack_frame_id=542} + %lt.731 = s32[1,41]{1,0} broadcast(%add.1228), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/lt" stack_frame_id=543} + %lt.732 = s32[41]{0} reshape(%lt.731), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/lt" stack_frame_id=543} + %lt.733 = s32[41,41]{1,0} broadcast(%lt.732), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/lt" stack_frame_id=543} + %lt.734 = pred[41,41]{1,0} compare(%lt.730, %lt.733), direction=LT, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/lt" stack_frame_id=543} + %convert_element_type.1972 = s32[41,41]{1,0} convert(%lt.734), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/convert_element_type" stack_frame_id=547} + %broadcast_in_dim.1766 = s32[1,41,41]{2,1,0} reshape(%convert_element_type.1972), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/broadcast_in_dim" stack_frame_id=551} + %min.110 = s32[1,41,41]{2,1,0} minimum(%convert_element_type.1973, %broadcast_in_dim.1766), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/min" stack_frame_id=551} + %broadcast_in_dim.1773 = s32[1,1,41,41]{3,2,1,0} reshape(%min.110), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/broadcast_in_dim" stack_frame_id=728} + %convert_element_type.1981 = pred[1,1,41,41]{3,2,1,0} compare(%broadcast_in_dim.1773, %broadcast.49), direction=NE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/convert_element_type" stack_frame_id=732} + %broadcast_in_dim.1774 = pred[1,1,1,41,41]{4,3,2,1,0} reshape(%convert_element_type.1981), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/broadcast_in_dim" stack_frame_id=34} + %and.343 = pred[1,1,1,41,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1774), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/vmap()/and" stack_frame_id=34} + %and.344 = pred[1,1,41,41]{3,2,1,0} reshape(%and.343), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/vmap()/and" stack_frame_id=34} + %and.345 = pred[1,1,32,41,41]{4,3,2,1,0} broadcast(%and.344), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/vmap()/and" stack_frame_id=34} + %state_1__137_.1 = bf16[4096,32,128]{2,1,0} parameter(140), metadata={op_name="state[1][137]"} + %dot_general.1215 = bf16[1,41,32,128]{3,2,1,0} dot(%convert_element_type.1976, %state_1__137_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=577} + %iota.635 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/iota" stack_frame_id=598} + %broadcast_in_dim.1769 = f32[1,41]{1,0} reshape(%iota.635), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/broadcast_in_dim" stack_frame_id=601} + %iota.634 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/iota" stack_frame_id=588} + %mul.3623 = f32[64]{0} multiply(%iota.634, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=588} + %div.1120 = f32[64]{0} divide(%mul.3623, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/div" stack_frame_id=589} + %neg.450 = f32[64]{0} negate(%div.1120), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/neg" stack_frame_id=590} + %pow.445 = f32[64]{0} power(%broadcast.51, %neg.450), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/pow" stack_frame_id=593} + %div.1121 = f32[64]{0} divide(%pow.445, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/div" stack_frame_id=594} + %dot_general.1216 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1769, %div.1121), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/bi,j->bij/dot_general" stack_frame_id=604} + %stack.1758 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1216), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/stack" stack_frame_id=607} + %stack.1759 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1216), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/stack" stack_frame_id=607} + %stack.1760 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1758, %stack.1759), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/stack" stack_frame_id=607} + %reshape.818 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1760), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/reshape"} + %cos.221 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.818), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/cos" stack_frame_id=611} + %convert_element_type.1977 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.221), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/convert_element_type" stack_frame_id=615} + %mul.3624 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1977), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=635} + %mul.3625 = bf16[1,41,128]{2,1,0} reshape(%mul.3624), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=635} + %mul.3626 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3625), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=635} + %mul.3627 = bf16[1,41,32,128]{3,2,1,0} multiply(%dot_general.1215, %mul.3626), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=635} + %split.443 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1215), slice={[0:1], [0:41], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/split" stack_frame_id=627} + %neg.451 = bf16[1,41,32,64]{3,2,1,0} negate(%split.443), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/neg" stack_frame_id=628} + %stack.1761 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%neg.451), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/stack" stack_frame_id=631} + %split.442 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1215), slice={[0:1], [0:41], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/split" stack_frame_id=627} + %stack.1762 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%split.442), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/stack" stack_frame_id=631} + %stack.1763 = bf16[1,41,32,2,64]{4,3,2,1,0} concatenate(%stack.1761, %stack.1762), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/stack" stack_frame_id=631} + %reshape.819 = bf16[1,41,32,128]{3,2,1,0} reshape(%stack.1763), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/reshape" stack_frame_id=634} + %sin.221 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.818), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/sin" stack_frame_id=619} + %convert_element_type.1978 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.221), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/convert_element_type" stack_frame_id=623} + %mul.3628 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1978), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=636} + %mul.3629 = bf16[1,41,128]{2,1,0} reshape(%mul.3628), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=636} + %mul.3630 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3629), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=636} + %mul.3631 = bf16[1,41,32,128]{3,2,1,0} multiply(%reshape.819, %mul.3630), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/mul" stack_frame_id=636} + %add.1230 = bf16[1,41,32,128]{3,2,1,0} add(%mul.3627, %mul.3631), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/rotary_embedding_15/add" stack_frame_id=637} + %reshape.824 = bf16[1,41,32,1,128]{4,3,2,1,0} reshape(%add.1230), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/reshape" stack_frame_id=34} + %broadcast_in_dim.1771 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%add.1231), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/broadcast_in_dim" stack_frame_id=722} + %reshape.822 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1771), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/reshape" stack_frame_id=722} + %dot_general.1220 = f32[1,32,41,1,41]{4,3,2,1,0} dot(%reshape.824, %reshape.822), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=34} + %mul.3641 = f32[1,32,41,1,41]{4,3,2,1,0} multiply(%dot_general.1220, %broadcast.48), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/vmap()/mul" stack_frame_id=34} + %vmap_jit__where__.110 = f32[1,1,32,41,41]{4,3,2,1,0} call(%and.345, %mul.3641, %constant.153), to_apply=%_where_0.5, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/vmap(jit(_where))" stack_frame_id=34} + %reduce_max.496 = f32[1,1,32,41]{3,2,1,0} reduce(%vmap_jit__where__.110, %constant.152), dimensions={4}, to_apply=%region_62.67, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/vmap()/reduce_max" stack_frame_id=34} + %max.120 = f32[1,1,32,41]{3,2,1,0} maximum(%reduce_max.496, %broadcast.47), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/vmap()/max" stack_frame_id=34} + %broadcast_in_dim.1775 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%max.120), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %sub.464 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1775), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/vmap()/sub" stack_frame_id=34} + %sub.465 = f32[1,1,32,41]{3,2,1,0} reshape(%sub.464), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/vmap()/sub" stack_frame_id=34} + %sub.466 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%sub.465), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/vmap()/sub" stack_frame_id=34} + %sub.467 = f32[1,1,32,41,41]{4,3,2,1,0} subtract(%vmap_jit__where__.110, %sub.466), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/vmap()/sub" stack_frame_id=34} + %exp.116 = f32[1,1,32,41,41]{4,3,2,1,0} exponential(%sub.467), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/vmap()/exp" stack_frame_id=34} + %reduce_sum.1495 = f32[1,1,32,41]{3,2,1,0} reduce(%exp.116, %constant.155), dimensions={4}, to_apply=%region_63.68, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/vmap()/reduce_sum" stack_frame_id=34} + %broadcast_in_dim.1776 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%reduce_sum.1495), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %div.1124 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1776), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/vmap()/div" stack_frame_id=34} + %div.1125 = f32[1,1,32,41]{3,2,1,0} reshape(%div.1124), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/vmap()/div" stack_frame_id=34} + %div.1126 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%div.1125), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/vmap()/div" stack_frame_id=34} + %div.1127 = f32[1,1,32,41,41]{4,3,2,1,0} divide(%exp.116, %div.1126), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/vmap()/div" stack_frame_id=34} + %convert_element_type.1982 = bf16[1,1,32,41,41]{4,3,2,1,0} convert(%div.1127), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/vmap()/convert_element_type" stack_frame_id=34} + %dot_general.1221 = bf16[1,32,128,1,41]{4,3,2,1,0} dot(%reshape.823, %convert_element_type.1982), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=34} + %transpose.114 = bf16[1,41,32,1,128]{1,3,4,2,0} transpose(%dot_general.1221), dimensions={0,4,1,3,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/vmap()/transpose;jit(compiled_generate_function)/transformer_layer_15/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %reshape.825 = bf16[1,41,32,128]{3,2,1,0} reshape(%transpose.114), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/reshape" stack_frame_id=34} + %state_1__140_.1 = bf16[32,128,4096]{2,1,0} parameter(143), metadata={op_name="state[1][140]"} + %dot_general.1222 = bf16[1,41,4096]{2,1,0} dot(%reshape.825, %state_1__140_.1), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=741} + %add.1232 = bf16[1,41,4096]{2,1,0} add(%dot_general.1222, %add.1227), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/add" stack_frame_id=742} + %convert_element_type.1983 = f32[1,41,4096]{2,1,0} convert(%add.1232), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/feedforward_layernorm/convert_element_type" stack_frame_id=746} + %pow.447 = f32[1,41,4096]{2,1,0} power(%convert_element_type.1983, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/feedforward_layernorm/pow" stack_frame_id=749} + %reduce_sum.1496 = f32[1,41]{1,0} reduce(%pow.447, %constant.155), dimensions={2}, to_apply=%region_64.69, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/feedforward_layernorm/reduce_sum" stack_frame_id=43} + %broadcast_in_dim.1777 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1496), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/feedforward_layernorm/broadcast_in_dim" stack_frame_id=43} + %div.1128 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1777, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/feedforward_layernorm/div" stack_frame_id=43} + %add.1233 = f32[1,41,1]{2,1,0} add(%div.1128, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/feedforward_layernorm/add" stack_frame_id=750} + %rsqrt.224 = f32[1,41,1]{2,1,0} rsqrt(%add.1233), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/feedforward_layernorm/rsqrt" stack_frame_id=753} + %mul.3642 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.224), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3643 = f32[1,41]{1,0} reshape(%mul.3642), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3644 = f32[1,41,4096]{2,1,0} broadcast(%mul.3643), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3645 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.1983, %mul.3644), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/feedforward_layernorm/mul" stack_frame_id=754} + %state_1__145_.1 = bf16[4096]{0} parameter(148), metadata={op_name="state[1][145]"} + %convert_element_type.1984 = f32[4096]{0} convert(%state_1__145_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/feedforward_layernorm/convert_element_type" stack_frame_id=755} + %broadcast_in_dim.1778 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1984), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/feedforward_layernorm/broadcast_in_dim" stack_frame_id=755} + %mul.3646 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1778), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3647 = f32[1,4096]{1,0} reshape(%mul.3646), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3648 = f32[1,41,4096]{2,1,0} broadcast(%mul.3647), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3649 = f32[1,41,4096]{2,1,0} multiply(%mul.3645, %mul.3648), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/feedforward_layernorm/mul" stack_frame_id=755} + %convert_element_type.1985 = bf16[1,41,4096]{2,1,0} convert(%mul.3649), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/feedforward_layernorm/convert_element_type" stack_frame_id=759} + %state_1__142_.1 = bf16[4096,14336]{1,0} parameter(145), metadata={op_name="state[1][142]"} + %dot_general.1224 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.1985, %state_1__142_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/feedforward_intermediate_dense/dot_general" stack_frame_id=788} + %state_1__143_.1 = bf16[4096,14336]{1,0} parameter(146), metadata={op_name="state[1][143]"} + %dot_general.1223 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.1985, %state_1__143_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/feedforward_gate_dense/dot_general" stack_frame_id=768} + %convert_element_type.1986 = f32[1,41,14336]{2,1,0} convert(%dot_general.1223), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/convert_element_type" stack_frame_id=772} + %jit_silu_.110 = f32[1,41,14336]{2,1,0} call(%convert_element_type.1986), to_apply=%silu.9, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/jit(silu)" stack_frame_id=775} + %convert_element_type.1987 = bf16[1,41,14336]{2,1,0} convert(%jit_silu_.110), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/convert_element_type" stack_frame_id=779} + %mul.3650 = bf16[1,41,14336]{2,1,0} multiply(%dot_general.1224, %convert_element_type.1987), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/mul" stack_frame_id=791} + %state_1__144_.1 = bf16[14336,4096]{1,0} parameter(147), metadata={op_name="state[1][144]"} + %dot_general.1225 = bf16[1,41,4096]{2,1,0} dot(%mul.3650, %state_1__144_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/feedforward_output_dense/dot_general" stack_frame_id=800} + %add.1234 = bf16[1,41,4096]{2,1,0} add(%dot_general.1225, %add.1232), metadata={op_name="jit(compiled_generate_function)/transformer_layer_15/add" stack_frame_id=801} + %convert_element_type.1990 = f32[1,41,4096]{2,1,0} convert(%add.1234), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention_layernorm/convert_element_type" stack_frame_id=555} + %pow.448 = f32[1,41,4096]{2,1,0} power(%convert_element_type.1990, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention_layernorm/pow" stack_frame_id=558} + %reduce_sum.1497 = f32[1,41]{1,0} reduce(%pow.448, %constant.155), dimensions={2}, to_apply=%region_65.70, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention_layernorm/reduce_sum" stack_frame_id=24} + %broadcast_in_dim.1785 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1497), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention_layernorm/broadcast_in_dim" stack_frame_id=24} + %div.1129 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1785, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention_layernorm/div" stack_frame_id=24} + %add.1236 = f32[1,41,1]{2,1,0} add(%div.1129, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention_layernorm/add" stack_frame_id=559} + %rsqrt.225 = f32[1,41,1]{2,1,0} rsqrt(%add.1236), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention_layernorm/rsqrt" stack_frame_id=562} + %mul.3651 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.225), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3652 = f32[1,41]{1,0} reshape(%mul.3651), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3653 = f32[1,41,4096]{2,1,0} broadcast(%mul.3652), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3654 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.1990, %mul.3653), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention_layernorm/mul" stack_frame_id=563} + %state_1__150_.1 = bf16[4096]{0} parameter(153), metadata={op_name="state[1][150]"} + %convert_element_type.1991 = f32[4096]{0} convert(%state_1__150_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention_layernorm/convert_element_type" stack_frame_id=564} + %broadcast_in_dim.1786 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.1991), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention_layernorm/broadcast_in_dim" stack_frame_id=564} + %mul.3655 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1786), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3656 = f32[1,4096]{1,0} reshape(%mul.3655), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3657 = f32[1,41,4096]{2,1,0} broadcast(%mul.3656), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3658 = f32[1,41,4096]{2,1,0} multiply(%mul.3654, %mul.3657), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention_layernorm/mul" stack_frame_id=564} + %convert_element_type.1992 = bf16[1,41,4096]{2,1,0} convert(%mul.3658), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention_layernorm/convert_element_type" stack_frame_id=568} + %state_1__147_.1 = bf16[4096,8,128]{2,1,0} parameter(150), metadata={op_name="state[1][147]"} + %dot_general.1228 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.1992, %state_1__147_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=647} + %iota.645 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/iota" stack_frame_id=677} + %broadcast_in_dim.1788 = f32[1,41]{1,0} reshape(%iota.645), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/broadcast_in_dim" stack_frame_id=680} + %iota.644 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/iota" stack_frame_id=667} + %mul.3668 = f32[64]{0} multiply(%iota.644, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=667} + %div.1132 = f32[64]{0} divide(%mul.3668, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/div" stack_frame_id=668} + %neg.456 = f32[64]{0} negate(%div.1132), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/neg" stack_frame_id=669} + %pow.450 = f32[64]{0} power(%broadcast.51, %neg.456), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/pow" stack_frame_id=672} + %div.1133 = f32[64]{0} divide(%pow.450, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/div" stack_frame_id=673} + %dot_general.1230 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1788, %div.1133), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/bi,j->bij/dot_general" stack_frame_id=683} + %stack.1779 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1230), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/stack" stack_frame_id=686} + %stack.1780 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1230), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/stack" stack_frame_id=686} + %stack.1781 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1779, %stack.1780), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/stack" stack_frame_id=686} + %reshape.828 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1781), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/reshape"} + %cos.224 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.828), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/cos" stack_frame_id=690} + %convert_element_type.1995 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.224), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/convert_element_type" stack_frame_id=694} + %mul.3669 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1995), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=714} + %mul.3670 = bf16[1,41,128]{2,1,0} reshape(%mul.3669), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=714} + %mul.3671 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3670), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=714} + %mul.3672 = bf16[1,41,8,128]{3,2,1,0} multiply(%dot_general.1228, %mul.3671), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=714} + %split.449 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1228), slice={[0:1], [0:41], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/split" stack_frame_id=706} + %neg.457 = bf16[1,41,8,64]{3,2,1,0} negate(%split.449), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/neg" stack_frame_id=707} + %stack.1782 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%neg.457), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/stack" stack_frame_id=710} + %split.448 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1228), slice={[0:1], [0:41], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/split" stack_frame_id=706} + %stack.1783 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%split.448), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/stack" stack_frame_id=710} + %stack.1784 = bf16[1,41,8,2,64]{4,3,2,1,0} concatenate(%stack.1782, %stack.1783), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/stack" stack_frame_id=710} + %reshape.829 = bf16[1,41,8,128]{3,2,1,0} reshape(%stack.1784), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/reshape" stack_frame_id=713} + %sin.224 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.828), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/sin" stack_frame_id=698} + %convert_element_type.1996 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.224), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/convert_element_type" stack_frame_id=702} + %mul.3673 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1996), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=715} + %mul.3674 = bf16[1,41,128]{2,1,0} reshape(%mul.3673), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=715} + %mul.3675 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3674), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=715} + %mul.3676 = bf16[1,41,8,128]{3,2,1,0} multiply(%reshape.829, %mul.3675), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=715} + %add.1238 = bf16[1,41,8,128]{3,2,1,0} add(%mul.3672, %mul.3676), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/add" stack_frame_id=716} + %stack.1785 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%add.1238), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/stack" stack_frame_id=719} + %state_1__148_.1 = bf16[4096,8,128]{2,1,0} parameter(151), metadata={op_name="state[1][148]"} + %dot_general.1229 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.1992, %state_1__148_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=656} + %stack.1786 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dot_general.1229), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/stack" stack_frame_id=719} + %stack.1787 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1785, %stack.1786), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/stack" stack_frame_id=719} + %stack.2023 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1787), metadata={op_name="jit(compiled_generate_function)/stack" stack_frame_id=804} + %broadcast_in_dim.1790 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dot_general.1229), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/broadcast_in_dim" stack_frame_id=725} + %reshape.831 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1790), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/reshape" stack_frame_id=725} + %iota.638 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/iota" stack_frame_id=525} + %broadcast_in_dim.1779 = f32[41,1]{1,0} reshape(%iota.638), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/broadcast_in_dim" stack_frame_id=528} + %ge.593 = f32[41,1]{1,0} broadcast(%broadcast_in_dim.1779), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/ge" stack_frame_id=532} + %ge.594 = f32[41]{0} reshape(%ge.593), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/ge" stack_frame_id=532} + %ge.595 = f32[41,41]{1,0} broadcast(%ge.594), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/ge" stack_frame_id=532} + %iota.639 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/iota" stack_frame_id=531} + %broadcast_in_dim.1780 = f32[1,41]{1,0} reshape(%iota.639), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/broadcast_in_dim" stack_frame_id=532} + %ge.596 = f32[1,41]{1,0} broadcast(%broadcast_in_dim.1780), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/ge" stack_frame_id=532} + %ge.597 = f32[41]{0} reshape(%ge.596), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/ge" stack_frame_id=532} + %ge.598 = f32[41,41]{1,0} broadcast(%ge.597), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/ge" stack_frame_id=532} + %ge.599 = pred[41,41]{1,0} compare(%ge.595, %ge.598), direction=GE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/ge" stack_frame_id=532} + %broadcast_in_dim.1781 = pred[1,41,41]{2,1,0} reshape(%ge.599), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/broadcast_in_dim" stack_frame_id=535} + %convert_element_type.1989 = s32[1,41,41]{2,1,0} convert(%broadcast_in_dim.1781), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/convert_element_type" stack_frame_id=551} + %iota.640 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/iota" stack_frame_id=538} + %broadcast_in_dim.1782 = s32[41,1]{1,0} reshape(%iota.640), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/broadcast_in_dim" stack_frame_id=536} + %lt.735 = s32[41,1]{1,0} broadcast(%broadcast_in_dim.1782), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/lt" stack_frame_id=543} + %lt.736 = s32[41]{0} reshape(%lt.735), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/lt" stack_frame_id=543} + %lt.737 = s32[41,41]{1,0} broadcast(%lt.736), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/lt" stack_frame_id=543} + %iota.641 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/iota" stack_frame_id=541} + %broadcast_in_dim.1783 = s32[1,41]{1,0} reshape(%iota.641), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/broadcast_in_dim" stack_frame_id=539} + %add.1235 = s32[1,41]{1,0} add(%broadcast_in_dim.1783, %broadcast.57), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/add" stack_frame_id=542} + %lt.738 = s32[1,41]{1,0} broadcast(%add.1235), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/lt" stack_frame_id=543} + %lt.739 = s32[41]{0} reshape(%lt.738), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/lt" stack_frame_id=543} + %lt.740 = s32[41,41]{1,0} broadcast(%lt.739), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/lt" stack_frame_id=543} + %lt.741 = pred[41,41]{1,0} compare(%lt.737, %lt.740), direction=LT, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/lt" stack_frame_id=543} + %convert_element_type.1988 = s32[41,41]{1,0} convert(%lt.741), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/convert_element_type" stack_frame_id=547} + %broadcast_in_dim.1784 = s32[1,41,41]{2,1,0} reshape(%convert_element_type.1988), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/broadcast_in_dim" stack_frame_id=551} + %min.111 = s32[1,41,41]{2,1,0} minimum(%convert_element_type.1989, %broadcast_in_dim.1784), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/min" stack_frame_id=551} + %broadcast_in_dim.1791 = s32[1,1,41,41]{3,2,1,0} reshape(%min.111), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/broadcast_in_dim" stack_frame_id=728} + %convert_element_type.1997 = pred[1,1,41,41]{3,2,1,0} compare(%broadcast_in_dim.1791, %broadcast.49), direction=NE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/convert_element_type" stack_frame_id=732} + %broadcast_in_dim.1792 = pred[1,1,1,41,41]{4,3,2,1,0} reshape(%convert_element_type.1997), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/broadcast_in_dim" stack_frame_id=34} + %and.346 = pred[1,1,1,41,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1792), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/vmap()/and" stack_frame_id=34} + %and.347 = pred[1,1,41,41]{3,2,1,0} reshape(%and.346), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/vmap()/and" stack_frame_id=34} + %and.348 = pred[1,1,32,41,41]{4,3,2,1,0} broadcast(%and.347), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/vmap()/and" stack_frame_id=34} + %state_1__146_.1 = bf16[4096,32,128]{2,1,0} parameter(149), metadata={op_name="state[1][146]"} + %dot_general.1226 = bf16[1,41,32,128]{3,2,1,0} dot(%convert_element_type.1992, %state_1__146_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=577} + %iota.643 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/iota" stack_frame_id=598} + %broadcast_in_dim.1787 = f32[1,41]{1,0} reshape(%iota.643), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/broadcast_in_dim" stack_frame_id=601} + %iota.642 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/iota" stack_frame_id=588} + %mul.3659 = f32[64]{0} multiply(%iota.642, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=588} + %div.1130 = f32[64]{0} divide(%mul.3659, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/div" stack_frame_id=589} + %neg.454 = f32[64]{0} negate(%div.1130), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/neg" stack_frame_id=590} + %pow.449 = f32[64]{0} power(%broadcast.51, %neg.454), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/pow" stack_frame_id=593} + %div.1131 = f32[64]{0} divide(%pow.449, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/div" stack_frame_id=594} + %dot_general.1227 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1787, %div.1131), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/bi,j->bij/dot_general" stack_frame_id=604} + %stack.1773 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1227), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/stack" stack_frame_id=607} + %stack.1774 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1227), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/stack" stack_frame_id=607} + %stack.1775 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1773, %stack.1774), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/stack" stack_frame_id=607} + %reshape.826 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1775), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/reshape"} + %cos.223 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.826), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/cos" stack_frame_id=611} + %convert_element_type.1993 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.223), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/convert_element_type" stack_frame_id=615} + %mul.3660 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1993), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=635} + %mul.3661 = bf16[1,41,128]{2,1,0} reshape(%mul.3660), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=635} + %mul.3662 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3661), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=635} + %mul.3663 = bf16[1,41,32,128]{3,2,1,0} multiply(%dot_general.1226, %mul.3662), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=635} + %split.447 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1226), slice={[0:1], [0:41], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/split" stack_frame_id=627} + %neg.455 = bf16[1,41,32,64]{3,2,1,0} negate(%split.447), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/neg" stack_frame_id=628} + %stack.1776 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%neg.455), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/stack" stack_frame_id=631} + %split.446 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1226), slice={[0:1], [0:41], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/split" stack_frame_id=627} + %stack.1777 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%split.446), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/stack" stack_frame_id=631} + %stack.1778 = bf16[1,41,32,2,64]{4,3,2,1,0} concatenate(%stack.1776, %stack.1777), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/stack" stack_frame_id=631} + %reshape.827 = bf16[1,41,32,128]{3,2,1,0} reshape(%stack.1778), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/reshape" stack_frame_id=634} + %sin.223 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.826), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/sin" stack_frame_id=619} + %convert_element_type.1994 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.223), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/convert_element_type" stack_frame_id=623} + %mul.3664 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.1994), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=636} + %mul.3665 = bf16[1,41,128]{2,1,0} reshape(%mul.3664), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=636} + %mul.3666 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3665), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=636} + %mul.3667 = bf16[1,41,32,128]{3,2,1,0} multiply(%reshape.827, %mul.3666), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/mul" stack_frame_id=636} + %add.1237 = bf16[1,41,32,128]{3,2,1,0} add(%mul.3663, %mul.3667), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/rotary_embedding_16/add" stack_frame_id=637} + %reshape.832 = bf16[1,41,32,1,128]{4,3,2,1,0} reshape(%add.1237), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/reshape" stack_frame_id=34} + %broadcast_in_dim.1789 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%add.1238), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/broadcast_in_dim" stack_frame_id=722} + %reshape.830 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1789), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/reshape" stack_frame_id=722} + %dot_general.1231 = f32[1,32,41,1,41]{4,3,2,1,0} dot(%reshape.832, %reshape.830), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=34} + %mul.3677 = f32[1,32,41,1,41]{4,3,2,1,0} multiply(%dot_general.1231, %broadcast.48), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/vmap()/mul" stack_frame_id=34} + %vmap_jit__where__.111 = f32[1,1,32,41,41]{4,3,2,1,0} call(%and.348, %mul.3677, %constant.153), to_apply=%_where_0.5, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/vmap(jit(_where))" stack_frame_id=34} + %reduce_max.497 = f32[1,1,32,41]{3,2,1,0} reduce(%vmap_jit__where__.111, %constant.152), dimensions={4}, to_apply=%region_66.71, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/vmap()/reduce_max" stack_frame_id=34} + %max.121 = f32[1,1,32,41]{3,2,1,0} maximum(%reduce_max.497, %broadcast.47), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/vmap()/max" stack_frame_id=34} + %broadcast_in_dim.1793 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%max.121), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %sub.468 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1793), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/vmap()/sub" stack_frame_id=34} + %sub.469 = f32[1,1,32,41]{3,2,1,0} reshape(%sub.468), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/vmap()/sub" stack_frame_id=34} + %sub.470 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%sub.469), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/vmap()/sub" stack_frame_id=34} + %sub.471 = f32[1,1,32,41,41]{4,3,2,1,0} subtract(%vmap_jit__where__.111, %sub.470), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/vmap()/sub" stack_frame_id=34} + %exp.117 = f32[1,1,32,41,41]{4,3,2,1,0} exponential(%sub.471), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/vmap()/exp" stack_frame_id=34} + %reduce_sum.1498 = f32[1,1,32,41]{3,2,1,0} reduce(%exp.117, %constant.155), dimensions={4}, to_apply=%region_67.72, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/vmap()/reduce_sum" stack_frame_id=34} + %broadcast_in_dim.1794 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%reduce_sum.1498), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %div.1134 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1794), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/vmap()/div" stack_frame_id=34} + %div.1135 = f32[1,1,32,41]{3,2,1,0} reshape(%div.1134), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/vmap()/div" stack_frame_id=34} + %div.1136 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%div.1135), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/vmap()/div" stack_frame_id=34} + %div.1137 = f32[1,1,32,41,41]{4,3,2,1,0} divide(%exp.117, %div.1136), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/vmap()/div" stack_frame_id=34} + %convert_element_type.1998 = bf16[1,1,32,41,41]{4,3,2,1,0} convert(%div.1137), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/vmap()/convert_element_type" stack_frame_id=34} + %dot_general.1232 = bf16[1,32,128,1,41]{4,3,2,1,0} dot(%reshape.831, %convert_element_type.1998), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=34} + %transpose.115 = bf16[1,41,32,1,128]{1,3,4,2,0} transpose(%dot_general.1232), dimensions={0,4,1,3,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/vmap()/transpose;jit(compiled_generate_function)/transformer_layer_16/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %reshape.833 = bf16[1,41,32,128]{3,2,1,0} reshape(%transpose.115), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/reshape" stack_frame_id=34} + %state_1__149_.1 = bf16[32,128,4096]{2,1,0} parameter(152), metadata={op_name="state[1][149]"} + %dot_general.1233 = bf16[1,41,4096]{2,1,0} dot(%reshape.833, %state_1__149_.1), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=741} + %add.1239 = bf16[1,41,4096]{2,1,0} add(%dot_general.1233, %add.1234), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/add" stack_frame_id=742} + %convert_element_type.1999 = f32[1,41,4096]{2,1,0} convert(%add.1239), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/feedforward_layernorm/convert_element_type" stack_frame_id=746} + %pow.451 = f32[1,41,4096]{2,1,0} power(%convert_element_type.1999, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/feedforward_layernorm/pow" stack_frame_id=749} + %reduce_sum.1499 = f32[1,41]{1,0} reduce(%pow.451, %constant.155), dimensions={2}, to_apply=%region_68.73, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/feedforward_layernorm/reduce_sum" stack_frame_id=43} + %broadcast_in_dim.1795 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1499), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/feedforward_layernorm/broadcast_in_dim" stack_frame_id=43} + %div.1138 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1795, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/feedforward_layernorm/div" stack_frame_id=43} + %add.1240 = f32[1,41,1]{2,1,0} add(%div.1138, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/feedforward_layernorm/add" stack_frame_id=750} + %rsqrt.226 = f32[1,41,1]{2,1,0} rsqrt(%add.1240), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/feedforward_layernorm/rsqrt" stack_frame_id=753} + %mul.3678 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.226), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3679 = f32[1,41]{1,0} reshape(%mul.3678), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3680 = f32[1,41,4096]{2,1,0} broadcast(%mul.3679), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3681 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.1999, %mul.3680), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/feedforward_layernorm/mul" stack_frame_id=754} + %state_1__154_.1 = bf16[4096]{0} parameter(157), metadata={op_name="state[1][154]"} + %convert_element_type.2000 = f32[4096]{0} convert(%state_1__154_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/feedforward_layernorm/convert_element_type" stack_frame_id=755} + %broadcast_in_dim.1796 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.2000), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/feedforward_layernorm/broadcast_in_dim" stack_frame_id=755} + %mul.3682 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1796), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3683 = f32[1,4096]{1,0} reshape(%mul.3682), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3684 = f32[1,41,4096]{2,1,0} broadcast(%mul.3683), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3685 = f32[1,41,4096]{2,1,0} multiply(%mul.3681, %mul.3684), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/feedforward_layernorm/mul" stack_frame_id=755} + %convert_element_type.2001 = bf16[1,41,4096]{2,1,0} convert(%mul.3685), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/feedforward_layernorm/convert_element_type" stack_frame_id=759} + %state_1__151_.1 = bf16[4096,14336]{1,0} parameter(154), metadata={op_name="state[1][151]"} + %dot_general.1235 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.2001, %state_1__151_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/feedforward_intermediate_dense/dot_general" stack_frame_id=788} + %state_1__152_.1 = bf16[4096,14336]{1,0} parameter(155), metadata={op_name="state[1][152]"} + %dot_general.1234 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.2001, %state_1__152_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/feedforward_gate_dense/dot_general" stack_frame_id=768} + %convert_element_type.2002 = f32[1,41,14336]{2,1,0} convert(%dot_general.1234), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/convert_element_type" stack_frame_id=772} + %jit_silu_.111 = f32[1,41,14336]{2,1,0} call(%convert_element_type.2002), to_apply=%silu.9, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/jit(silu)" stack_frame_id=775} + %convert_element_type.2003 = bf16[1,41,14336]{2,1,0} convert(%jit_silu_.111), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/convert_element_type" stack_frame_id=779} + %mul.3686 = bf16[1,41,14336]{2,1,0} multiply(%dot_general.1235, %convert_element_type.2003), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/mul" stack_frame_id=791} + %state_1__153_.1 = bf16[14336,4096]{1,0} parameter(156), metadata={op_name="state[1][153]"} + %dot_general.1236 = bf16[1,41,4096]{2,1,0} dot(%mul.3686, %state_1__153_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/feedforward_output_dense/dot_general" stack_frame_id=800} + %add.1241 = bf16[1,41,4096]{2,1,0} add(%dot_general.1236, %add.1239), metadata={op_name="jit(compiled_generate_function)/transformer_layer_16/add" stack_frame_id=801} + %convert_element_type.2006 = f32[1,41,4096]{2,1,0} convert(%add.1241), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention_layernorm/convert_element_type" stack_frame_id=555} + %pow.452 = f32[1,41,4096]{2,1,0} power(%convert_element_type.2006, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention_layernorm/pow" stack_frame_id=558} + %reduce_sum.1500 = f32[1,41]{1,0} reduce(%pow.452, %constant.155), dimensions={2}, to_apply=%region_69.74, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention_layernorm/reduce_sum" stack_frame_id=24} + %broadcast_in_dim.1803 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1500), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention_layernorm/broadcast_in_dim" stack_frame_id=24} + %div.1139 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1803, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention_layernorm/div" stack_frame_id=24} + %add.1243 = f32[1,41,1]{2,1,0} add(%div.1139, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention_layernorm/add" stack_frame_id=559} + %rsqrt.227 = f32[1,41,1]{2,1,0} rsqrt(%add.1243), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention_layernorm/rsqrt" stack_frame_id=562} + %mul.3687 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.227), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3688 = f32[1,41]{1,0} reshape(%mul.3687), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3689 = f32[1,41,4096]{2,1,0} broadcast(%mul.3688), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3690 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.2006, %mul.3689), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention_layernorm/mul" stack_frame_id=563} + %state_1__159_.1 = bf16[4096]{0} parameter(162), metadata={op_name="state[1][159]"} + %convert_element_type.2007 = f32[4096]{0} convert(%state_1__159_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention_layernorm/convert_element_type" stack_frame_id=564} + %broadcast_in_dim.1804 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.2007), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention_layernorm/broadcast_in_dim" stack_frame_id=564} + %mul.3691 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1804), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3692 = f32[1,4096]{1,0} reshape(%mul.3691), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3693 = f32[1,41,4096]{2,1,0} broadcast(%mul.3692), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3694 = f32[1,41,4096]{2,1,0} multiply(%mul.3690, %mul.3693), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention_layernorm/mul" stack_frame_id=564} + %convert_element_type.2008 = bf16[1,41,4096]{2,1,0} convert(%mul.3694), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention_layernorm/convert_element_type" stack_frame_id=568} + %state_1__156_.1 = bf16[4096,8,128]{2,1,0} parameter(159), metadata={op_name="state[1][156]"} + %dot_general.1239 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.2008, %state_1__156_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=647} + %iota.653 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/iota" stack_frame_id=677} + %broadcast_in_dim.1806 = f32[1,41]{1,0} reshape(%iota.653), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/broadcast_in_dim" stack_frame_id=680} + %iota.652 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/iota" stack_frame_id=667} + %mul.3704 = f32[64]{0} multiply(%iota.652, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=667} + %div.1142 = f32[64]{0} divide(%mul.3704, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/div" stack_frame_id=668} + %neg.460 = f32[64]{0} negate(%div.1142), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/neg" stack_frame_id=669} + %pow.454 = f32[64]{0} power(%broadcast.51, %neg.460), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/pow" stack_frame_id=672} + %div.1143 = f32[64]{0} divide(%pow.454, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/div" stack_frame_id=673} + %dot_general.1241 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1806, %div.1143), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/bi,j->bij/dot_general" stack_frame_id=683} + %stack.1794 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1241), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/stack" stack_frame_id=686} + %stack.1795 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1241), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/stack" stack_frame_id=686} + %stack.1796 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1794, %stack.1795), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/stack" stack_frame_id=686} + %reshape.836 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1796), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/reshape"} + %cos.226 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.836), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/cos" stack_frame_id=690} + %convert_element_type.2011 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.226), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/convert_element_type" stack_frame_id=694} + %mul.3705 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2011), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=714} + %mul.3706 = bf16[1,41,128]{2,1,0} reshape(%mul.3705), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=714} + %mul.3707 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3706), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=714} + %mul.3708 = bf16[1,41,8,128]{3,2,1,0} multiply(%dot_general.1239, %mul.3707), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=714} + %split.453 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1239), slice={[0:1], [0:41], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/split" stack_frame_id=706} + %neg.461 = bf16[1,41,8,64]{3,2,1,0} negate(%split.453), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/neg" stack_frame_id=707} + %stack.1797 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%neg.461), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/stack" stack_frame_id=710} + %split.452 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1239), slice={[0:1], [0:41], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/split" stack_frame_id=706} + %stack.1798 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%split.452), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/stack" stack_frame_id=710} + %stack.1799 = bf16[1,41,8,2,64]{4,3,2,1,0} concatenate(%stack.1797, %stack.1798), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/stack" stack_frame_id=710} + %reshape.837 = bf16[1,41,8,128]{3,2,1,0} reshape(%stack.1799), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/reshape" stack_frame_id=713} + %sin.226 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.836), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/sin" stack_frame_id=698} + %convert_element_type.2012 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.226), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/convert_element_type" stack_frame_id=702} + %mul.3709 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2012), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=715} + %mul.3710 = bf16[1,41,128]{2,1,0} reshape(%mul.3709), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=715} + %mul.3711 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3710), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=715} + %mul.3712 = bf16[1,41,8,128]{3,2,1,0} multiply(%reshape.837, %mul.3711), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=715} + %add.1245 = bf16[1,41,8,128]{3,2,1,0} add(%mul.3708, %mul.3712), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/add" stack_frame_id=716} + %stack.1800 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%add.1245), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/stack" stack_frame_id=719} + %state_1__157_.1 = bf16[4096,8,128]{2,1,0} parameter(160), metadata={op_name="state[1][157]"} + %dot_general.1240 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.2008, %state_1__157_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=656} + %stack.1801 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dot_general.1240), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/stack" stack_frame_id=719} + %stack.1802 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1800, %stack.1801), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/stack" stack_frame_id=719} + %stack.2024 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1802), metadata={op_name="jit(compiled_generate_function)/stack" stack_frame_id=804} + %broadcast_in_dim.1808 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dot_general.1240), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/broadcast_in_dim" stack_frame_id=725} + %reshape.839 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1808), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/reshape" stack_frame_id=725} + %iota.646 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/iota" stack_frame_id=525} + %broadcast_in_dim.1797 = f32[41,1]{1,0} reshape(%iota.646), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/broadcast_in_dim" stack_frame_id=528} + %ge.600 = f32[41,1]{1,0} broadcast(%broadcast_in_dim.1797), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/ge" stack_frame_id=532} + %ge.601 = f32[41]{0} reshape(%ge.600), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/ge" stack_frame_id=532} + %ge.602 = f32[41,41]{1,0} broadcast(%ge.601), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/ge" stack_frame_id=532} + %iota.647 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/iota" stack_frame_id=531} + %broadcast_in_dim.1798 = f32[1,41]{1,0} reshape(%iota.647), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/broadcast_in_dim" stack_frame_id=532} + %ge.603 = f32[1,41]{1,0} broadcast(%broadcast_in_dim.1798), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/ge" stack_frame_id=532} + %ge.604 = f32[41]{0} reshape(%ge.603), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/ge" stack_frame_id=532} + %ge.605 = f32[41,41]{1,0} broadcast(%ge.604), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/ge" stack_frame_id=532} + %ge.606 = pred[41,41]{1,0} compare(%ge.602, %ge.605), direction=GE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/ge" stack_frame_id=532} + %broadcast_in_dim.1799 = pred[1,41,41]{2,1,0} reshape(%ge.606), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/broadcast_in_dim" stack_frame_id=535} + %convert_element_type.2005 = s32[1,41,41]{2,1,0} convert(%broadcast_in_dim.1799), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/convert_element_type" stack_frame_id=551} + %iota.648 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/iota" stack_frame_id=538} + %broadcast_in_dim.1800 = s32[41,1]{1,0} reshape(%iota.648), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/broadcast_in_dim" stack_frame_id=536} + %lt.742 = s32[41,1]{1,0} broadcast(%broadcast_in_dim.1800), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/lt" stack_frame_id=543} + %lt.743 = s32[41]{0} reshape(%lt.742), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/lt" stack_frame_id=543} + %lt.744 = s32[41,41]{1,0} broadcast(%lt.743), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/lt" stack_frame_id=543} + %iota.649 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/iota" stack_frame_id=541} + %broadcast_in_dim.1801 = s32[1,41]{1,0} reshape(%iota.649), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/broadcast_in_dim" stack_frame_id=539} + %add.1242 = s32[1,41]{1,0} add(%broadcast_in_dim.1801, %broadcast.57), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/add" stack_frame_id=542} + %lt.745 = s32[1,41]{1,0} broadcast(%add.1242), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/lt" stack_frame_id=543} + %lt.746 = s32[41]{0} reshape(%lt.745), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/lt" stack_frame_id=543} + %lt.747 = s32[41,41]{1,0} broadcast(%lt.746), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/lt" stack_frame_id=543} + %lt.748 = pred[41,41]{1,0} compare(%lt.744, %lt.747), direction=LT, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/lt" stack_frame_id=543} + %convert_element_type.2004 = s32[41,41]{1,0} convert(%lt.748), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/convert_element_type" stack_frame_id=547} + %broadcast_in_dim.1802 = s32[1,41,41]{2,1,0} reshape(%convert_element_type.2004), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/broadcast_in_dim" stack_frame_id=551} + %min.112 = s32[1,41,41]{2,1,0} minimum(%convert_element_type.2005, %broadcast_in_dim.1802), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/min" stack_frame_id=551} + %broadcast_in_dim.1809 = s32[1,1,41,41]{3,2,1,0} reshape(%min.112), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/broadcast_in_dim" stack_frame_id=728} + %convert_element_type.2013 = pred[1,1,41,41]{3,2,1,0} compare(%broadcast_in_dim.1809, %broadcast.49), direction=NE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/convert_element_type" stack_frame_id=732} + %broadcast_in_dim.1810 = pred[1,1,1,41,41]{4,3,2,1,0} reshape(%convert_element_type.2013), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/broadcast_in_dim" stack_frame_id=34} + %and.349 = pred[1,1,1,41,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1810), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/vmap()/and" stack_frame_id=34} + %and.350 = pred[1,1,41,41]{3,2,1,0} reshape(%and.349), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/vmap()/and" stack_frame_id=34} + %and.351 = pred[1,1,32,41,41]{4,3,2,1,0} broadcast(%and.350), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/vmap()/and" stack_frame_id=34} + %state_1__155_.1 = bf16[4096,32,128]{2,1,0} parameter(158), metadata={op_name="state[1][155]"} + %dot_general.1237 = bf16[1,41,32,128]{3,2,1,0} dot(%convert_element_type.2008, %state_1__155_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=577} + %iota.651 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/iota" stack_frame_id=598} + %broadcast_in_dim.1805 = f32[1,41]{1,0} reshape(%iota.651), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/broadcast_in_dim" stack_frame_id=601} + %iota.650 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/iota" stack_frame_id=588} + %mul.3695 = f32[64]{0} multiply(%iota.650, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=588} + %div.1140 = f32[64]{0} divide(%mul.3695, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/div" stack_frame_id=589} + %neg.458 = f32[64]{0} negate(%div.1140), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/neg" stack_frame_id=590} + %pow.453 = f32[64]{0} power(%broadcast.51, %neg.458), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/pow" stack_frame_id=593} + %div.1141 = f32[64]{0} divide(%pow.453, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/div" stack_frame_id=594} + %dot_general.1238 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1805, %div.1141), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/bi,j->bij/dot_general" stack_frame_id=604} + %stack.1788 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1238), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/stack" stack_frame_id=607} + %stack.1789 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1238), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/stack" stack_frame_id=607} + %stack.1790 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1788, %stack.1789), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/stack" stack_frame_id=607} + %reshape.834 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1790), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/reshape"} + %cos.225 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.834), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/cos" stack_frame_id=611} + %convert_element_type.2009 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.225), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/convert_element_type" stack_frame_id=615} + %mul.3696 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2009), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=635} + %mul.3697 = bf16[1,41,128]{2,1,0} reshape(%mul.3696), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=635} + %mul.3698 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3697), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=635} + %mul.3699 = bf16[1,41,32,128]{3,2,1,0} multiply(%dot_general.1237, %mul.3698), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=635} + %split.451 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1237), slice={[0:1], [0:41], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/split" stack_frame_id=627} + %neg.459 = bf16[1,41,32,64]{3,2,1,0} negate(%split.451), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/neg" stack_frame_id=628} + %stack.1791 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%neg.459), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/stack" stack_frame_id=631} + %split.450 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1237), slice={[0:1], [0:41], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/split" stack_frame_id=627} + %stack.1792 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%split.450), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/stack" stack_frame_id=631} + %stack.1793 = bf16[1,41,32,2,64]{4,3,2,1,0} concatenate(%stack.1791, %stack.1792), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/stack" stack_frame_id=631} + %reshape.835 = bf16[1,41,32,128]{3,2,1,0} reshape(%stack.1793), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/reshape" stack_frame_id=634} + %sin.225 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.834), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/sin" stack_frame_id=619} + %convert_element_type.2010 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.225), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/convert_element_type" stack_frame_id=623} + %mul.3700 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2010), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=636} + %mul.3701 = bf16[1,41,128]{2,1,0} reshape(%mul.3700), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=636} + %mul.3702 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3701), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=636} + %mul.3703 = bf16[1,41,32,128]{3,2,1,0} multiply(%reshape.835, %mul.3702), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/mul" stack_frame_id=636} + %add.1244 = bf16[1,41,32,128]{3,2,1,0} add(%mul.3699, %mul.3703), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/rotary_embedding_17/add" stack_frame_id=637} + %reshape.840 = bf16[1,41,32,1,128]{4,3,2,1,0} reshape(%add.1244), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/reshape" stack_frame_id=34} + %broadcast_in_dim.1807 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%add.1245), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/broadcast_in_dim" stack_frame_id=722} + %reshape.838 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1807), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/reshape" stack_frame_id=722} + %dot_general.1242 = f32[1,32,41,1,41]{4,3,2,1,0} dot(%reshape.840, %reshape.838), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=34} + %mul.3713 = f32[1,32,41,1,41]{4,3,2,1,0} multiply(%dot_general.1242, %broadcast.48), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/vmap()/mul" stack_frame_id=34} + %vmap_jit__where__.112 = f32[1,1,32,41,41]{4,3,2,1,0} call(%and.351, %mul.3713, %constant.153), to_apply=%_where_0.5, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/vmap(jit(_where))" stack_frame_id=34} + %reduce_max.498 = f32[1,1,32,41]{3,2,1,0} reduce(%vmap_jit__where__.112, %constant.152), dimensions={4}, to_apply=%region_70.75, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/vmap()/reduce_max" stack_frame_id=34} + %max.122 = f32[1,1,32,41]{3,2,1,0} maximum(%reduce_max.498, %broadcast.47), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/vmap()/max" stack_frame_id=34} + %broadcast_in_dim.1811 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%max.122), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %sub.472 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1811), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/vmap()/sub" stack_frame_id=34} + %sub.473 = f32[1,1,32,41]{3,2,1,0} reshape(%sub.472), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/vmap()/sub" stack_frame_id=34} + %sub.474 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%sub.473), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/vmap()/sub" stack_frame_id=34} + %sub.475 = f32[1,1,32,41,41]{4,3,2,1,0} subtract(%vmap_jit__where__.112, %sub.474), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/vmap()/sub" stack_frame_id=34} + %exp.118 = f32[1,1,32,41,41]{4,3,2,1,0} exponential(%sub.475), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/vmap()/exp" stack_frame_id=34} + %reduce_sum.1501 = f32[1,1,32,41]{3,2,1,0} reduce(%exp.118, %constant.155), dimensions={4}, to_apply=%region_71.76, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/vmap()/reduce_sum" stack_frame_id=34} + %broadcast_in_dim.1812 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%reduce_sum.1501), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %div.1144 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1812), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/vmap()/div" stack_frame_id=34} + %div.1145 = f32[1,1,32,41]{3,2,1,0} reshape(%div.1144), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/vmap()/div" stack_frame_id=34} + %div.1146 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%div.1145), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/vmap()/div" stack_frame_id=34} + %div.1147 = f32[1,1,32,41,41]{4,3,2,1,0} divide(%exp.118, %div.1146), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/vmap()/div" stack_frame_id=34} + %convert_element_type.2014 = bf16[1,1,32,41,41]{4,3,2,1,0} convert(%div.1147), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/vmap()/convert_element_type" stack_frame_id=34} + %dot_general.1243 = bf16[1,32,128,1,41]{4,3,2,1,0} dot(%reshape.839, %convert_element_type.2014), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=34} + %transpose.116 = bf16[1,41,32,1,128]{1,3,4,2,0} transpose(%dot_general.1243), dimensions={0,4,1,3,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/vmap()/transpose;jit(compiled_generate_function)/transformer_layer_17/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %reshape.841 = bf16[1,41,32,128]{3,2,1,0} reshape(%transpose.116), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/reshape" stack_frame_id=34} + %state_1__158_.1 = bf16[32,128,4096]{2,1,0} parameter(161), metadata={op_name="state[1][158]"} + %dot_general.1244 = bf16[1,41,4096]{2,1,0} dot(%reshape.841, %state_1__158_.1), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=741} + %add.1246 = bf16[1,41,4096]{2,1,0} add(%dot_general.1244, %add.1241), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/add" stack_frame_id=742} + %convert_element_type.2015 = f32[1,41,4096]{2,1,0} convert(%add.1246), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/feedforward_layernorm/convert_element_type" stack_frame_id=746} + %pow.455 = f32[1,41,4096]{2,1,0} power(%convert_element_type.2015, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/feedforward_layernorm/pow" stack_frame_id=749} + %reduce_sum.1502 = f32[1,41]{1,0} reduce(%pow.455, %constant.155), dimensions={2}, to_apply=%region_72.77, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/feedforward_layernorm/reduce_sum" stack_frame_id=43} + %broadcast_in_dim.1813 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1502), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/feedforward_layernorm/broadcast_in_dim" stack_frame_id=43} + %div.1148 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1813, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/feedforward_layernorm/div" stack_frame_id=43} + %add.1247 = f32[1,41,1]{2,1,0} add(%div.1148, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/feedforward_layernorm/add" stack_frame_id=750} + %rsqrt.228 = f32[1,41,1]{2,1,0} rsqrt(%add.1247), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/feedforward_layernorm/rsqrt" stack_frame_id=753} + %mul.3714 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.228), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3715 = f32[1,41]{1,0} reshape(%mul.3714), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3716 = f32[1,41,4096]{2,1,0} broadcast(%mul.3715), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3717 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.2015, %mul.3716), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/feedforward_layernorm/mul" stack_frame_id=754} + %state_1__163_.1 = bf16[4096]{0} parameter(166), metadata={op_name="state[1][163]"} + %convert_element_type.2016 = f32[4096]{0} convert(%state_1__163_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/feedforward_layernorm/convert_element_type" stack_frame_id=755} + %broadcast_in_dim.1814 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.2016), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/feedforward_layernorm/broadcast_in_dim" stack_frame_id=755} + %mul.3718 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1814), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3719 = f32[1,4096]{1,0} reshape(%mul.3718), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3720 = f32[1,41,4096]{2,1,0} broadcast(%mul.3719), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3721 = f32[1,41,4096]{2,1,0} multiply(%mul.3717, %mul.3720), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/feedforward_layernorm/mul" stack_frame_id=755} + %convert_element_type.2017 = bf16[1,41,4096]{2,1,0} convert(%mul.3721), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/feedforward_layernorm/convert_element_type" stack_frame_id=759} + %state_1__160_.1 = bf16[4096,14336]{1,0} parameter(163), metadata={op_name="state[1][160]"} + %dot_general.1246 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.2017, %state_1__160_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/feedforward_intermediate_dense/dot_general" stack_frame_id=788} + %state_1__161_.1 = bf16[4096,14336]{1,0} parameter(164), metadata={op_name="state[1][161]"} + %dot_general.1245 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.2017, %state_1__161_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/feedforward_gate_dense/dot_general" stack_frame_id=768} + %convert_element_type.2018 = f32[1,41,14336]{2,1,0} convert(%dot_general.1245), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/convert_element_type" stack_frame_id=772} + %jit_silu_.112 = f32[1,41,14336]{2,1,0} call(%convert_element_type.2018), to_apply=%silu.9, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/jit(silu)" stack_frame_id=775} + %convert_element_type.2019 = bf16[1,41,14336]{2,1,0} convert(%jit_silu_.112), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/convert_element_type" stack_frame_id=779} + %mul.3722 = bf16[1,41,14336]{2,1,0} multiply(%dot_general.1246, %convert_element_type.2019), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/mul" stack_frame_id=791} + %state_1__162_.1 = bf16[14336,4096]{1,0} parameter(165), metadata={op_name="state[1][162]"} + %dot_general.1247 = bf16[1,41,4096]{2,1,0} dot(%mul.3722, %state_1__162_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/feedforward_output_dense/dot_general" stack_frame_id=800} + %add.1248 = bf16[1,41,4096]{2,1,0} add(%dot_general.1247, %add.1246), metadata={op_name="jit(compiled_generate_function)/transformer_layer_17/add" stack_frame_id=801} + %convert_element_type.2022 = f32[1,41,4096]{2,1,0} convert(%add.1248), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention_layernorm/convert_element_type" stack_frame_id=555} + %pow.456 = f32[1,41,4096]{2,1,0} power(%convert_element_type.2022, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention_layernorm/pow" stack_frame_id=558} + %reduce_sum.1503 = f32[1,41]{1,0} reduce(%pow.456, %constant.155), dimensions={2}, to_apply=%region_73.78, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention_layernorm/reduce_sum" stack_frame_id=24} + %broadcast_in_dim.1821 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1503), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention_layernorm/broadcast_in_dim" stack_frame_id=24} + %div.1149 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1821, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention_layernorm/div" stack_frame_id=24} + %add.1250 = f32[1,41,1]{2,1,0} add(%div.1149, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention_layernorm/add" stack_frame_id=559} + %rsqrt.229 = f32[1,41,1]{2,1,0} rsqrt(%add.1250), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention_layernorm/rsqrt" stack_frame_id=562} + %mul.3723 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.229), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3724 = f32[1,41]{1,0} reshape(%mul.3723), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3725 = f32[1,41,4096]{2,1,0} broadcast(%mul.3724), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3726 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.2022, %mul.3725), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention_layernorm/mul" stack_frame_id=563} + %state_1__168_.1 = bf16[4096]{0} parameter(171), metadata={op_name="state[1][168]"} + %convert_element_type.2023 = f32[4096]{0} convert(%state_1__168_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention_layernorm/convert_element_type" stack_frame_id=564} + %broadcast_in_dim.1822 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.2023), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention_layernorm/broadcast_in_dim" stack_frame_id=564} + %mul.3727 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1822), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3728 = f32[1,4096]{1,0} reshape(%mul.3727), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3729 = f32[1,41,4096]{2,1,0} broadcast(%mul.3728), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3730 = f32[1,41,4096]{2,1,0} multiply(%mul.3726, %mul.3729), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention_layernorm/mul" stack_frame_id=564} + %convert_element_type.2024 = bf16[1,41,4096]{2,1,0} convert(%mul.3730), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention_layernorm/convert_element_type" stack_frame_id=568} + %state_1__165_.1 = bf16[4096,8,128]{2,1,0} parameter(168), metadata={op_name="state[1][165]"} + %dot_general.1250 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.2024, %state_1__165_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=647} + %iota.661 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/iota" stack_frame_id=677} + %broadcast_in_dim.1824 = f32[1,41]{1,0} reshape(%iota.661), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/broadcast_in_dim" stack_frame_id=680} + %iota.660 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/iota" stack_frame_id=667} + %mul.3740 = f32[64]{0} multiply(%iota.660, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=667} + %div.1152 = f32[64]{0} divide(%mul.3740, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/div" stack_frame_id=668} + %neg.464 = f32[64]{0} negate(%div.1152), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/neg" stack_frame_id=669} + %pow.458 = f32[64]{0} power(%broadcast.51, %neg.464), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/pow" stack_frame_id=672} + %div.1153 = f32[64]{0} divide(%pow.458, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/div" stack_frame_id=673} + %dot_general.1252 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1824, %div.1153), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/bi,j->bij/dot_general" stack_frame_id=683} + %stack.1809 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1252), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/stack" stack_frame_id=686} + %stack.1810 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1252), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/stack" stack_frame_id=686} + %stack.1811 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1809, %stack.1810), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/stack" stack_frame_id=686} + %reshape.844 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1811), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/reshape"} + %cos.228 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.844), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/cos" stack_frame_id=690} + %convert_element_type.2027 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.228), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/convert_element_type" stack_frame_id=694} + %mul.3741 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2027), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=714} + %mul.3742 = bf16[1,41,128]{2,1,0} reshape(%mul.3741), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=714} + %mul.3743 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3742), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=714} + %mul.3744 = bf16[1,41,8,128]{3,2,1,0} multiply(%dot_general.1250, %mul.3743), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=714} + %split.457 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1250), slice={[0:1], [0:41], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/split" stack_frame_id=706} + %neg.465 = bf16[1,41,8,64]{3,2,1,0} negate(%split.457), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/neg" stack_frame_id=707} + %stack.1812 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%neg.465), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/stack" stack_frame_id=710} + %split.456 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1250), slice={[0:1], [0:41], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/split" stack_frame_id=706} + %stack.1813 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%split.456), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/stack" stack_frame_id=710} + %stack.1814 = bf16[1,41,8,2,64]{4,3,2,1,0} concatenate(%stack.1812, %stack.1813), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/stack" stack_frame_id=710} + %reshape.845 = bf16[1,41,8,128]{3,2,1,0} reshape(%stack.1814), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/reshape" stack_frame_id=713} + %sin.228 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.844), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/sin" stack_frame_id=698} + %convert_element_type.2028 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.228), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/convert_element_type" stack_frame_id=702} + %mul.3745 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2028), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=715} + %mul.3746 = bf16[1,41,128]{2,1,0} reshape(%mul.3745), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=715} + %mul.3747 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3746), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=715} + %mul.3748 = bf16[1,41,8,128]{3,2,1,0} multiply(%reshape.845, %mul.3747), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=715} + %add.1252 = bf16[1,41,8,128]{3,2,1,0} add(%mul.3744, %mul.3748), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/add" stack_frame_id=716} + %stack.1815 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%add.1252), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/stack" stack_frame_id=719} + %state_1__166_.1 = bf16[4096,8,128]{2,1,0} parameter(169), metadata={op_name="state[1][166]"} + %dot_general.1251 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.2024, %state_1__166_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=656} + %stack.1816 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dot_general.1251), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/stack" stack_frame_id=719} + %stack.1817 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1815, %stack.1816), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/stack" stack_frame_id=719} + %stack.2025 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1817), metadata={op_name="jit(compiled_generate_function)/stack" stack_frame_id=804} + %broadcast_in_dim.1826 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dot_general.1251), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/broadcast_in_dim" stack_frame_id=725} + %reshape.847 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1826), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/reshape" stack_frame_id=725} + %iota.654 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/iota" stack_frame_id=525} + %broadcast_in_dim.1815 = f32[41,1]{1,0} reshape(%iota.654), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/broadcast_in_dim" stack_frame_id=528} + %ge.607 = f32[41,1]{1,0} broadcast(%broadcast_in_dim.1815), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/ge" stack_frame_id=532} + %ge.608 = f32[41]{0} reshape(%ge.607), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/ge" stack_frame_id=532} + %ge.609 = f32[41,41]{1,0} broadcast(%ge.608), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/ge" stack_frame_id=532} + %iota.655 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/iota" stack_frame_id=531} + %broadcast_in_dim.1816 = f32[1,41]{1,0} reshape(%iota.655), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/broadcast_in_dim" stack_frame_id=532} + %ge.610 = f32[1,41]{1,0} broadcast(%broadcast_in_dim.1816), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/ge" stack_frame_id=532} + %ge.611 = f32[41]{0} reshape(%ge.610), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/ge" stack_frame_id=532} + %ge.612 = f32[41,41]{1,0} broadcast(%ge.611), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/ge" stack_frame_id=532} + %ge.613 = pred[41,41]{1,0} compare(%ge.609, %ge.612), direction=GE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/ge" stack_frame_id=532} + %broadcast_in_dim.1817 = pred[1,41,41]{2,1,0} reshape(%ge.613), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/broadcast_in_dim" stack_frame_id=535} + %convert_element_type.2021 = s32[1,41,41]{2,1,0} convert(%broadcast_in_dim.1817), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/convert_element_type" stack_frame_id=551} + %iota.656 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/iota" stack_frame_id=538} + %broadcast_in_dim.1818 = s32[41,1]{1,0} reshape(%iota.656), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/broadcast_in_dim" stack_frame_id=536} + %lt.749 = s32[41,1]{1,0} broadcast(%broadcast_in_dim.1818), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/lt" stack_frame_id=543} + %lt.750 = s32[41]{0} reshape(%lt.749), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/lt" stack_frame_id=543} + %lt.751 = s32[41,41]{1,0} broadcast(%lt.750), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/lt" stack_frame_id=543} + %iota.657 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/iota" stack_frame_id=541} + %broadcast_in_dim.1819 = s32[1,41]{1,0} reshape(%iota.657), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/broadcast_in_dim" stack_frame_id=539} + %add.1249 = s32[1,41]{1,0} add(%broadcast_in_dim.1819, %broadcast.57), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/add" stack_frame_id=542} + %lt.752 = s32[1,41]{1,0} broadcast(%add.1249), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/lt" stack_frame_id=543} + %lt.753 = s32[41]{0} reshape(%lt.752), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/lt" stack_frame_id=543} + %lt.754 = s32[41,41]{1,0} broadcast(%lt.753), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/lt" stack_frame_id=543} + %lt.755 = pred[41,41]{1,0} compare(%lt.751, %lt.754), direction=LT, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/lt" stack_frame_id=543} + %convert_element_type.2020 = s32[41,41]{1,0} convert(%lt.755), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/convert_element_type" stack_frame_id=547} + %broadcast_in_dim.1820 = s32[1,41,41]{2,1,0} reshape(%convert_element_type.2020), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/broadcast_in_dim" stack_frame_id=551} + %min.113 = s32[1,41,41]{2,1,0} minimum(%convert_element_type.2021, %broadcast_in_dim.1820), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/min" stack_frame_id=551} + %broadcast_in_dim.1827 = s32[1,1,41,41]{3,2,1,0} reshape(%min.113), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/broadcast_in_dim" stack_frame_id=728} + %convert_element_type.2029 = pred[1,1,41,41]{3,2,1,0} compare(%broadcast_in_dim.1827, %broadcast.49), direction=NE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/convert_element_type" stack_frame_id=732} + %broadcast_in_dim.1828 = pred[1,1,1,41,41]{4,3,2,1,0} reshape(%convert_element_type.2029), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/broadcast_in_dim" stack_frame_id=34} + %and.352 = pred[1,1,1,41,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1828), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/vmap()/and" stack_frame_id=34} + %and.353 = pred[1,1,41,41]{3,2,1,0} reshape(%and.352), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/vmap()/and" stack_frame_id=34} + %and.354 = pred[1,1,32,41,41]{4,3,2,1,0} broadcast(%and.353), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/vmap()/and" stack_frame_id=34} + %state_1__164_.1 = bf16[4096,32,128]{2,1,0} parameter(167), metadata={op_name="state[1][164]"} + %dot_general.1248 = bf16[1,41,32,128]{3,2,1,0} dot(%convert_element_type.2024, %state_1__164_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=577} + %iota.659 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/iota" stack_frame_id=598} + %broadcast_in_dim.1823 = f32[1,41]{1,0} reshape(%iota.659), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/broadcast_in_dim" stack_frame_id=601} + %iota.658 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/iota" stack_frame_id=588} + %mul.3731 = f32[64]{0} multiply(%iota.658, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=588} + %div.1150 = f32[64]{0} divide(%mul.3731, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/div" stack_frame_id=589} + %neg.462 = f32[64]{0} negate(%div.1150), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/neg" stack_frame_id=590} + %pow.457 = f32[64]{0} power(%broadcast.51, %neg.462), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/pow" stack_frame_id=593} + %div.1151 = f32[64]{0} divide(%pow.457, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/div" stack_frame_id=594} + %dot_general.1249 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1823, %div.1151), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/bi,j->bij/dot_general" stack_frame_id=604} + %stack.1803 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1249), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/stack" stack_frame_id=607} + %stack.1804 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1249), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/stack" stack_frame_id=607} + %stack.1805 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1803, %stack.1804), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/stack" stack_frame_id=607} + %reshape.842 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1805), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/reshape"} + %cos.227 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.842), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/cos" stack_frame_id=611} + %convert_element_type.2025 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.227), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/convert_element_type" stack_frame_id=615} + %mul.3732 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2025), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=635} + %mul.3733 = bf16[1,41,128]{2,1,0} reshape(%mul.3732), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=635} + %mul.3734 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3733), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=635} + %mul.3735 = bf16[1,41,32,128]{3,2,1,0} multiply(%dot_general.1248, %mul.3734), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=635} + %split.455 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1248), slice={[0:1], [0:41], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/split" stack_frame_id=627} + %neg.463 = bf16[1,41,32,64]{3,2,1,0} negate(%split.455), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/neg" stack_frame_id=628} + %stack.1806 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%neg.463), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/stack" stack_frame_id=631} + %split.454 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1248), slice={[0:1], [0:41], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/split" stack_frame_id=627} + %stack.1807 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%split.454), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/stack" stack_frame_id=631} + %stack.1808 = bf16[1,41,32,2,64]{4,3,2,1,0} concatenate(%stack.1806, %stack.1807), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/stack" stack_frame_id=631} + %reshape.843 = bf16[1,41,32,128]{3,2,1,0} reshape(%stack.1808), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/reshape" stack_frame_id=634} + %sin.227 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.842), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/sin" stack_frame_id=619} + %convert_element_type.2026 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.227), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/convert_element_type" stack_frame_id=623} + %mul.3736 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2026), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=636} + %mul.3737 = bf16[1,41,128]{2,1,0} reshape(%mul.3736), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=636} + %mul.3738 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3737), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=636} + %mul.3739 = bf16[1,41,32,128]{3,2,1,0} multiply(%reshape.843, %mul.3738), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/mul" stack_frame_id=636} + %add.1251 = bf16[1,41,32,128]{3,2,1,0} add(%mul.3735, %mul.3739), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/rotary_embedding_18/add" stack_frame_id=637} + %reshape.848 = bf16[1,41,32,1,128]{4,3,2,1,0} reshape(%add.1251), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/reshape" stack_frame_id=34} + %broadcast_in_dim.1825 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%add.1252), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/broadcast_in_dim" stack_frame_id=722} + %reshape.846 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1825), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/reshape" stack_frame_id=722} + %dot_general.1253 = f32[1,32,41,1,41]{4,3,2,1,0} dot(%reshape.848, %reshape.846), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=34} + %mul.3749 = f32[1,32,41,1,41]{4,3,2,1,0} multiply(%dot_general.1253, %broadcast.48), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/vmap()/mul" stack_frame_id=34} + %vmap_jit__where__.113 = f32[1,1,32,41,41]{4,3,2,1,0} call(%and.354, %mul.3749, %constant.153), to_apply=%_where_0.5, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/vmap(jit(_where))" stack_frame_id=34} + %reduce_max.499 = f32[1,1,32,41]{3,2,1,0} reduce(%vmap_jit__where__.113, %constant.152), dimensions={4}, to_apply=%region_74.79, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/vmap()/reduce_max" stack_frame_id=34} + %max.123 = f32[1,1,32,41]{3,2,1,0} maximum(%reduce_max.499, %broadcast.47), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/vmap()/max" stack_frame_id=34} + %broadcast_in_dim.1829 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%max.123), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %sub.476 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1829), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/vmap()/sub" stack_frame_id=34} + %sub.477 = f32[1,1,32,41]{3,2,1,0} reshape(%sub.476), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/vmap()/sub" stack_frame_id=34} + %sub.478 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%sub.477), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/vmap()/sub" stack_frame_id=34} + %sub.479 = f32[1,1,32,41,41]{4,3,2,1,0} subtract(%vmap_jit__where__.113, %sub.478), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/vmap()/sub" stack_frame_id=34} + %exp.119 = f32[1,1,32,41,41]{4,3,2,1,0} exponential(%sub.479), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/vmap()/exp" stack_frame_id=34} + %reduce_sum.1504 = f32[1,1,32,41]{3,2,1,0} reduce(%exp.119, %constant.155), dimensions={4}, to_apply=%region_75.80, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/vmap()/reduce_sum" stack_frame_id=34} + %broadcast_in_dim.1830 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%reduce_sum.1504), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %div.1154 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1830), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/vmap()/div" stack_frame_id=34} + %div.1155 = f32[1,1,32,41]{3,2,1,0} reshape(%div.1154), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/vmap()/div" stack_frame_id=34} + %div.1156 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%div.1155), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/vmap()/div" stack_frame_id=34} + %div.1157 = f32[1,1,32,41,41]{4,3,2,1,0} divide(%exp.119, %div.1156), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/vmap()/div" stack_frame_id=34} + %convert_element_type.2030 = bf16[1,1,32,41,41]{4,3,2,1,0} convert(%div.1157), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/vmap()/convert_element_type" stack_frame_id=34} + %dot_general.1254 = bf16[1,32,128,1,41]{4,3,2,1,0} dot(%reshape.847, %convert_element_type.2030), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=34} + %transpose.117 = bf16[1,41,32,1,128]{1,3,4,2,0} transpose(%dot_general.1254), dimensions={0,4,1,3,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/vmap()/transpose;jit(compiled_generate_function)/transformer_layer_18/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %reshape.849 = bf16[1,41,32,128]{3,2,1,0} reshape(%transpose.117), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/reshape" stack_frame_id=34} + %state_1__167_.1 = bf16[32,128,4096]{2,1,0} parameter(170), metadata={op_name="state[1][167]"} + %dot_general.1255 = bf16[1,41,4096]{2,1,0} dot(%reshape.849, %state_1__167_.1), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=741} + %add.1253 = bf16[1,41,4096]{2,1,0} add(%dot_general.1255, %add.1248), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/add" stack_frame_id=742} + %convert_element_type.2031 = f32[1,41,4096]{2,1,0} convert(%add.1253), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/feedforward_layernorm/convert_element_type" stack_frame_id=746} + %pow.459 = f32[1,41,4096]{2,1,0} power(%convert_element_type.2031, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/feedforward_layernorm/pow" stack_frame_id=749} + %reduce_sum.1505 = f32[1,41]{1,0} reduce(%pow.459, %constant.155), dimensions={2}, to_apply=%region_76.81, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/feedforward_layernorm/reduce_sum" stack_frame_id=43} + %broadcast_in_dim.1831 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1505), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/feedforward_layernorm/broadcast_in_dim" stack_frame_id=43} + %div.1158 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1831, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/feedforward_layernorm/div" stack_frame_id=43} + %add.1254 = f32[1,41,1]{2,1,0} add(%div.1158, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/feedforward_layernorm/add" stack_frame_id=750} + %rsqrt.230 = f32[1,41,1]{2,1,0} rsqrt(%add.1254), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/feedforward_layernorm/rsqrt" stack_frame_id=753} + %mul.3750 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.230), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3751 = f32[1,41]{1,0} reshape(%mul.3750), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3752 = f32[1,41,4096]{2,1,0} broadcast(%mul.3751), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3753 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.2031, %mul.3752), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/feedforward_layernorm/mul" stack_frame_id=754} + %state_1__172_.1 = bf16[4096]{0} parameter(175), metadata={op_name="state[1][172]"} + %convert_element_type.2032 = f32[4096]{0} convert(%state_1__172_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/feedforward_layernorm/convert_element_type" stack_frame_id=755} + %broadcast_in_dim.1832 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.2032), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/feedforward_layernorm/broadcast_in_dim" stack_frame_id=755} + %mul.3754 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1832), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3755 = f32[1,4096]{1,0} reshape(%mul.3754), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3756 = f32[1,41,4096]{2,1,0} broadcast(%mul.3755), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3757 = f32[1,41,4096]{2,1,0} multiply(%mul.3753, %mul.3756), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/feedforward_layernorm/mul" stack_frame_id=755} + %convert_element_type.2033 = bf16[1,41,4096]{2,1,0} convert(%mul.3757), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/feedforward_layernorm/convert_element_type" stack_frame_id=759} + %state_1__169_.1 = bf16[4096,14336]{1,0} parameter(172), metadata={op_name="state[1][169]"} + %dot_general.1257 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.2033, %state_1__169_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/feedforward_intermediate_dense/dot_general" stack_frame_id=788} + %state_1__170_.1 = bf16[4096,14336]{1,0} parameter(173), metadata={op_name="state[1][170]"} + %dot_general.1256 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.2033, %state_1__170_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/feedforward_gate_dense/dot_general" stack_frame_id=768} + %convert_element_type.2034 = f32[1,41,14336]{2,1,0} convert(%dot_general.1256), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/convert_element_type" stack_frame_id=772} + %jit_silu_.113 = f32[1,41,14336]{2,1,0} call(%convert_element_type.2034), to_apply=%silu.9, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/jit(silu)" stack_frame_id=775} + %convert_element_type.2035 = bf16[1,41,14336]{2,1,0} convert(%jit_silu_.113), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/convert_element_type" stack_frame_id=779} + %mul.3758 = bf16[1,41,14336]{2,1,0} multiply(%dot_general.1257, %convert_element_type.2035), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/mul" stack_frame_id=791} + %state_1__171_.1 = bf16[14336,4096]{1,0} parameter(174), metadata={op_name="state[1][171]"} + %dot_general.1258 = bf16[1,41,4096]{2,1,0} dot(%mul.3758, %state_1__171_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/feedforward_output_dense/dot_general" stack_frame_id=800} + %add.1255 = bf16[1,41,4096]{2,1,0} add(%dot_general.1258, %add.1253), metadata={op_name="jit(compiled_generate_function)/transformer_layer_18/add" stack_frame_id=801} + %convert_element_type.2038 = f32[1,41,4096]{2,1,0} convert(%add.1255), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention_layernorm/convert_element_type" stack_frame_id=555} + %pow.460 = f32[1,41,4096]{2,1,0} power(%convert_element_type.2038, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention_layernorm/pow" stack_frame_id=558} + %reduce_sum.1506 = f32[1,41]{1,0} reduce(%pow.460, %constant.155), dimensions={2}, to_apply=%region_77.82, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention_layernorm/reduce_sum" stack_frame_id=24} + %broadcast_in_dim.1839 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1506), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention_layernorm/broadcast_in_dim" stack_frame_id=24} + %div.1159 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1839, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention_layernorm/div" stack_frame_id=24} + %add.1257 = f32[1,41,1]{2,1,0} add(%div.1159, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention_layernorm/add" stack_frame_id=559} + %rsqrt.231 = f32[1,41,1]{2,1,0} rsqrt(%add.1257), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention_layernorm/rsqrt" stack_frame_id=562} + %mul.3759 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.231), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3760 = f32[1,41]{1,0} reshape(%mul.3759), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3761 = f32[1,41,4096]{2,1,0} broadcast(%mul.3760), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3762 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.2038, %mul.3761), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention_layernorm/mul" stack_frame_id=563} + %state_1__177_.1 = bf16[4096]{0} parameter(180), metadata={op_name="state[1][177]"} + %convert_element_type.2039 = f32[4096]{0} convert(%state_1__177_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention_layernorm/convert_element_type" stack_frame_id=564} + %broadcast_in_dim.1840 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.2039), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention_layernorm/broadcast_in_dim" stack_frame_id=564} + %mul.3763 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1840), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3764 = f32[1,4096]{1,0} reshape(%mul.3763), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3765 = f32[1,41,4096]{2,1,0} broadcast(%mul.3764), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3766 = f32[1,41,4096]{2,1,0} multiply(%mul.3762, %mul.3765), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention_layernorm/mul" stack_frame_id=564} + %convert_element_type.2040 = bf16[1,41,4096]{2,1,0} convert(%mul.3766), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention_layernorm/convert_element_type" stack_frame_id=568} + %state_1__174_.1 = bf16[4096,8,128]{2,1,0} parameter(177), metadata={op_name="state[1][174]"} + %dot_general.1261 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.2040, %state_1__174_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=647} + %iota.669 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/iota" stack_frame_id=677} + %broadcast_in_dim.1842 = f32[1,41]{1,0} reshape(%iota.669), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/broadcast_in_dim" stack_frame_id=680} + %iota.668 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/iota" stack_frame_id=667} + %mul.3776 = f32[64]{0} multiply(%iota.668, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=667} + %div.1162 = f32[64]{0} divide(%mul.3776, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/div" stack_frame_id=668} + %neg.468 = f32[64]{0} negate(%div.1162), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/neg" stack_frame_id=669} + %pow.462 = f32[64]{0} power(%broadcast.51, %neg.468), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/pow" stack_frame_id=672} + %div.1163 = f32[64]{0} divide(%pow.462, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/div" stack_frame_id=673} + %dot_general.1263 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1842, %div.1163), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/bi,j->bij/dot_general" stack_frame_id=683} + %stack.1824 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1263), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/stack" stack_frame_id=686} + %stack.1825 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1263), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/stack" stack_frame_id=686} + %stack.1826 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1824, %stack.1825), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/stack" stack_frame_id=686} + %reshape.852 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1826), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/reshape"} + %cos.230 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.852), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/cos" stack_frame_id=690} + %convert_element_type.2043 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.230), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/convert_element_type" stack_frame_id=694} + %mul.3777 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2043), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=714} + %mul.3778 = bf16[1,41,128]{2,1,0} reshape(%mul.3777), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=714} + %mul.3779 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3778), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=714} + %mul.3780 = bf16[1,41,8,128]{3,2,1,0} multiply(%dot_general.1261, %mul.3779), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=714} + %split.461 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1261), slice={[0:1], [0:41], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/split" stack_frame_id=706} + %neg.469 = bf16[1,41,8,64]{3,2,1,0} negate(%split.461), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/neg" stack_frame_id=707} + %stack.1827 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%neg.469), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/stack" stack_frame_id=710} + %split.460 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1261), slice={[0:1], [0:41], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/split" stack_frame_id=706} + %stack.1828 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%split.460), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/stack" stack_frame_id=710} + %stack.1829 = bf16[1,41,8,2,64]{4,3,2,1,0} concatenate(%stack.1827, %stack.1828), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/stack" stack_frame_id=710} + %reshape.853 = bf16[1,41,8,128]{3,2,1,0} reshape(%stack.1829), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/reshape" stack_frame_id=713} + %sin.230 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.852), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/sin" stack_frame_id=698} + %convert_element_type.2044 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.230), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/convert_element_type" stack_frame_id=702} + %mul.3781 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2044), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=715} + %mul.3782 = bf16[1,41,128]{2,1,0} reshape(%mul.3781), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=715} + %mul.3783 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3782), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=715} + %mul.3784 = bf16[1,41,8,128]{3,2,1,0} multiply(%reshape.853, %mul.3783), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=715} + %add.1259 = bf16[1,41,8,128]{3,2,1,0} add(%mul.3780, %mul.3784), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/add" stack_frame_id=716} + %stack.1830 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%add.1259), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/stack" stack_frame_id=719} + %state_1__175_.1 = bf16[4096,8,128]{2,1,0} parameter(178), metadata={op_name="state[1][175]"} + %dot_general.1262 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.2040, %state_1__175_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=656} + %stack.1831 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dot_general.1262), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/stack" stack_frame_id=719} + %stack.1832 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1830, %stack.1831), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/stack" stack_frame_id=719} + %stack.2026 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1832), metadata={op_name="jit(compiled_generate_function)/stack" stack_frame_id=804} + %broadcast_in_dim.1844 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dot_general.1262), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/broadcast_in_dim" stack_frame_id=725} + %reshape.855 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1844), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/reshape" stack_frame_id=725} + %iota.662 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/iota" stack_frame_id=525} + %broadcast_in_dim.1833 = f32[41,1]{1,0} reshape(%iota.662), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/broadcast_in_dim" stack_frame_id=528} + %ge.614 = f32[41,1]{1,0} broadcast(%broadcast_in_dim.1833), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/ge" stack_frame_id=532} + %ge.615 = f32[41]{0} reshape(%ge.614), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/ge" stack_frame_id=532} + %ge.616 = f32[41,41]{1,0} broadcast(%ge.615), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/ge" stack_frame_id=532} + %iota.663 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/iota" stack_frame_id=531} + %broadcast_in_dim.1834 = f32[1,41]{1,0} reshape(%iota.663), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/broadcast_in_dim" stack_frame_id=532} + %ge.617 = f32[1,41]{1,0} broadcast(%broadcast_in_dim.1834), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/ge" stack_frame_id=532} + %ge.618 = f32[41]{0} reshape(%ge.617), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/ge" stack_frame_id=532} + %ge.619 = f32[41,41]{1,0} broadcast(%ge.618), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/ge" stack_frame_id=532} + %ge.620 = pred[41,41]{1,0} compare(%ge.616, %ge.619), direction=GE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/ge" stack_frame_id=532} + %broadcast_in_dim.1835 = pred[1,41,41]{2,1,0} reshape(%ge.620), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/broadcast_in_dim" stack_frame_id=535} + %convert_element_type.2037 = s32[1,41,41]{2,1,0} convert(%broadcast_in_dim.1835), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/convert_element_type" stack_frame_id=551} + %iota.664 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/iota" stack_frame_id=538} + %broadcast_in_dim.1836 = s32[41,1]{1,0} reshape(%iota.664), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/broadcast_in_dim" stack_frame_id=536} + %lt.756 = s32[41,1]{1,0} broadcast(%broadcast_in_dim.1836), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/lt" stack_frame_id=543} + %lt.757 = s32[41]{0} reshape(%lt.756), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/lt" stack_frame_id=543} + %lt.758 = s32[41,41]{1,0} broadcast(%lt.757), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/lt" stack_frame_id=543} + %iota.665 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/iota" stack_frame_id=541} + %broadcast_in_dim.1837 = s32[1,41]{1,0} reshape(%iota.665), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/broadcast_in_dim" stack_frame_id=539} + %add.1256 = s32[1,41]{1,0} add(%broadcast_in_dim.1837, %broadcast.57), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/add" stack_frame_id=542} + %lt.759 = s32[1,41]{1,0} broadcast(%add.1256), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/lt" stack_frame_id=543} + %lt.760 = s32[41]{0} reshape(%lt.759), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/lt" stack_frame_id=543} + %lt.761 = s32[41,41]{1,0} broadcast(%lt.760), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/lt" stack_frame_id=543} + %lt.762 = pred[41,41]{1,0} compare(%lt.758, %lt.761), direction=LT, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/lt" stack_frame_id=543} + %convert_element_type.2036 = s32[41,41]{1,0} convert(%lt.762), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/convert_element_type" stack_frame_id=547} + %broadcast_in_dim.1838 = s32[1,41,41]{2,1,0} reshape(%convert_element_type.2036), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/broadcast_in_dim" stack_frame_id=551} + %min.114 = s32[1,41,41]{2,1,0} minimum(%convert_element_type.2037, %broadcast_in_dim.1838), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/min" stack_frame_id=551} + %broadcast_in_dim.1845 = s32[1,1,41,41]{3,2,1,0} reshape(%min.114), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/broadcast_in_dim" stack_frame_id=728} + %convert_element_type.2045 = pred[1,1,41,41]{3,2,1,0} compare(%broadcast_in_dim.1845, %broadcast.49), direction=NE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/convert_element_type" stack_frame_id=732} + %broadcast_in_dim.1846 = pred[1,1,1,41,41]{4,3,2,1,0} reshape(%convert_element_type.2045), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/broadcast_in_dim" stack_frame_id=34} + %and.355 = pred[1,1,1,41,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1846), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/vmap()/and" stack_frame_id=34} + %and.356 = pred[1,1,41,41]{3,2,1,0} reshape(%and.355), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/vmap()/and" stack_frame_id=34} + %and.357 = pred[1,1,32,41,41]{4,3,2,1,0} broadcast(%and.356), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/vmap()/and" stack_frame_id=34} + %state_1__173_.1 = bf16[4096,32,128]{2,1,0} parameter(176), metadata={op_name="state[1][173]"} + %dot_general.1259 = bf16[1,41,32,128]{3,2,1,0} dot(%convert_element_type.2040, %state_1__173_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=577} + %iota.667 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/iota" stack_frame_id=598} + %broadcast_in_dim.1841 = f32[1,41]{1,0} reshape(%iota.667), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/broadcast_in_dim" stack_frame_id=601} + %iota.666 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/iota" stack_frame_id=588} + %mul.3767 = f32[64]{0} multiply(%iota.666, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=588} + %div.1160 = f32[64]{0} divide(%mul.3767, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/div" stack_frame_id=589} + %neg.466 = f32[64]{0} negate(%div.1160), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/neg" stack_frame_id=590} + %pow.461 = f32[64]{0} power(%broadcast.51, %neg.466), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/pow" stack_frame_id=593} + %div.1161 = f32[64]{0} divide(%pow.461, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/div" stack_frame_id=594} + %dot_general.1260 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1841, %div.1161), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/bi,j->bij/dot_general" stack_frame_id=604} + %stack.1818 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1260), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/stack" stack_frame_id=607} + %stack.1819 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1260), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/stack" stack_frame_id=607} + %stack.1820 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1818, %stack.1819), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/stack" stack_frame_id=607} + %reshape.850 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1820), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/reshape"} + %cos.229 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.850), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/cos" stack_frame_id=611} + %convert_element_type.2041 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.229), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/convert_element_type" stack_frame_id=615} + %mul.3768 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2041), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=635} + %mul.3769 = bf16[1,41,128]{2,1,0} reshape(%mul.3768), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=635} + %mul.3770 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3769), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=635} + %mul.3771 = bf16[1,41,32,128]{3,2,1,0} multiply(%dot_general.1259, %mul.3770), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=635} + %split.459 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1259), slice={[0:1], [0:41], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/split" stack_frame_id=627} + %neg.467 = bf16[1,41,32,64]{3,2,1,0} negate(%split.459), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/neg" stack_frame_id=628} + %stack.1821 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%neg.467), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/stack" stack_frame_id=631} + %split.458 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1259), slice={[0:1], [0:41], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/split" stack_frame_id=627} + %stack.1822 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%split.458), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/stack" stack_frame_id=631} + %stack.1823 = bf16[1,41,32,2,64]{4,3,2,1,0} concatenate(%stack.1821, %stack.1822), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/stack" stack_frame_id=631} + %reshape.851 = bf16[1,41,32,128]{3,2,1,0} reshape(%stack.1823), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/reshape" stack_frame_id=634} + %sin.229 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.850), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/sin" stack_frame_id=619} + %convert_element_type.2042 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.229), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/convert_element_type" stack_frame_id=623} + %mul.3772 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2042), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=636} + %mul.3773 = bf16[1,41,128]{2,1,0} reshape(%mul.3772), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=636} + %mul.3774 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3773), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=636} + %mul.3775 = bf16[1,41,32,128]{3,2,1,0} multiply(%reshape.851, %mul.3774), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/mul" stack_frame_id=636} + %add.1258 = bf16[1,41,32,128]{3,2,1,0} add(%mul.3771, %mul.3775), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/rotary_embedding_19/add" stack_frame_id=637} + %reshape.856 = bf16[1,41,32,1,128]{4,3,2,1,0} reshape(%add.1258), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/reshape" stack_frame_id=34} + %broadcast_in_dim.1843 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%add.1259), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/broadcast_in_dim" stack_frame_id=722} + %reshape.854 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1843), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/reshape" stack_frame_id=722} + %dot_general.1264 = f32[1,32,41,1,41]{4,3,2,1,0} dot(%reshape.856, %reshape.854), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=34} + %mul.3785 = f32[1,32,41,1,41]{4,3,2,1,0} multiply(%dot_general.1264, %broadcast.48), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/vmap()/mul" stack_frame_id=34} + %vmap_jit__where__.114 = f32[1,1,32,41,41]{4,3,2,1,0} call(%and.357, %mul.3785, %constant.153), to_apply=%_where_0.5, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/vmap(jit(_where))" stack_frame_id=34} + %reduce_max.500 = f32[1,1,32,41]{3,2,1,0} reduce(%vmap_jit__where__.114, %constant.152), dimensions={4}, to_apply=%region_78.83, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/vmap()/reduce_max" stack_frame_id=34} + %max.124 = f32[1,1,32,41]{3,2,1,0} maximum(%reduce_max.500, %broadcast.47), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/vmap()/max" stack_frame_id=34} + %broadcast_in_dim.1847 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%max.124), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %sub.480 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1847), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/vmap()/sub" stack_frame_id=34} + %sub.481 = f32[1,1,32,41]{3,2,1,0} reshape(%sub.480), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/vmap()/sub" stack_frame_id=34} + %sub.482 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%sub.481), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/vmap()/sub" stack_frame_id=34} + %sub.483 = f32[1,1,32,41,41]{4,3,2,1,0} subtract(%vmap_jit__where__.114, %sub.482), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/vmap()/sub" stack_frame_id=34} + %exp.120 = f32[1,1,32,41,41]{4,3,2,1,0} exponential(%sub.483), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/vmap()/exp" stack_frame_id=34} + %reduce_sum.1507 = f32[1,1,32,41]{3,2,1,0} reduce(%exp.120, %constant.155), dimensions={4}, to_apply=%region_79.84, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/vmap()/reduce_sum" stack_frame_id=34} + %broadcast_in_dim.1848 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%reduce_sum.1507), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %div.1164 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1848), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/vmap()/div" stack_frame_id=34} + %div.1165 = f32[1,1,32,41]{3,2,1,0} reshape(%div.1164), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/vmap()/div" stack_frame_id=34} + %div.1166 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%div.1165), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/vmap()/div" stack_frame_id=34} + %div.1167 = f32[1,1,32,41,41]{4,3,2,1,0} divide(%exp.120, %div.1166), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/vmap()/div" stack_frame_id=34} + %convert_element_type.2046 = bf16[1,1,32,41,41]{4,3,2,1,0} convert(%div.1167), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/vmap()/convert_element_type" stack_frame_id=34} + %dot_general.1265 = bf16[1,32,128,1,41]{4,3,2,1,0} dot(%reshape.855, %convert_element_type.2046), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=34} + %transpose.118 = bf16[1,41,32,1,128]{1,3,4,2,0} transpose(%dot_general.1265), dimensions={0,4,1,3,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/vmap()/transpose;jit(compiled_generate_function)/transformer_layer_19/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %reshape.857 = bf16[1,41,32,128]{3,2,1,0} reshape(%transpose.118), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/reshape" stack_frame_id=34} + %state_1__176_.1 = bf16[32,128,4096]{2,1,0} parameter(179), metadata={op_name="state[1][176]"} + %dot_general.1266 = bf16[1,41,4096]{2,1,0} dot(%reshape.857, %state_1__176_.1), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=741} + %add.1260 = bf16[1,41,4096]{2,1,0} add(%dot_general.1266, %add.1255), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/add" stack_frame_id=742} + %convert_element_type.2047 = f32[1,41,4096]{2,1,0} convert(%add.1260), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/feedforward_layernorm/convert_element_type" stack_frame_id=746} + %pow.463 = f32[1,41,4096]{2,1,0} power(%convert_element_type.2047, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/feedforward_layernorm/pow" stack_frame_id=749} + %reduce_sum.1508 = f32[1,41]{1,0} reduce(%pow.463, %constant.155), dimensions={2}, to_apply=%region_80.85, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/feedforward_layernorm/reduce_sum" stack_frame_id=43} + %broadcast_in_dim.1849 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1508), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/feedforward_layernorm/broadcast_in_dim" stack_frame_id=43} + %div.1168 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1849, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/feedforward_layernorm/div" stack_frame_id=43} + %add.1261 = f32[1,41,1]{2,1,0} add(%div.1168, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/feedforward_layernorm/add" stack_frame_id=750} + %rsqrt.232 = f32[1,41,1]{2,1,0} rsqrt(%add.1261), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/feedforward_layernorm/rsqrt" stack_frame_id=753} + %mul.3786 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.232), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3787 = f32[1,41]{1,0} reshape(%mul.3786), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3788 = f32[1,41,4096]{2,1,0} broadcast(%mul.3787), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3789 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.2047, %mul.3788), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/feedforward_layernorm/mul" stack_frame_id=754} + %state_1__181_.1 = bf16[4096]{0} parameter(184), metadata={op_name="state[1][181]"} + %convert_element_type.2048 = f32[4096]{0} convert(%state_1__181_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/feedforward_layernorm/convert_element_type" stack_frame_id=755} + %broadcast_in_dim.1850 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.2048), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/feedforward_layernorm/broadcast_in_dim" stack_frame_id=755} + %mul.3790 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1850), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3791 = f32[1,4096]{1,0} reshape(%mul.3790), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3792 = f32[1,41,4096]{2,1,0} broadcast(%mul.3791), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3793 = f32[1,41,4096]{2,1,0} multiply(%mul.3789, %mul.3792), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/feedforward_layernorm/mul" stack_frame_id=755} + %convert_element_type.2049 = bf16[1,41,4096]{2,1,0} convert(%mul.3793), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/feedforward_layernorm/convert_element_type" stack_frame_id=759} + %state_1__178_.1 = bf16[4096,14336]{1,0} parameter(181), metadata={op_name="state[1][178]"} + %dot_general.1268 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.2049, %state_1__178_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/feedforward_intermediate_dense/dot_general" stack_frame_id=788} + %state_1__179_.1 = bf16[4096,14336]{1,0} parameter(182), metadata={op_name="state[1][179]"} + %dot_general.1267 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.2049, %state_1__179_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/feedforward_gate_dense/dot_general" stack_frame_id=768} + %convert_element_type.2050 = f32[1,41,14336]{2,1,0} convert(%dot_general.1267), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/convert_element_type" stack_frame_id=772} + %jit_silu_.114 = f32[1,41,14336]{2,1,0} call(%convert_element_type.2050), to_apply=%silu.9, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/jit(silu)" stack_frame_id=775} + %convert_element_type.2051 = bf16[1,41,14336]{2,1,0} convert(%jit_silu_.114), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/convert_element_type" stack_frame_id=779} + %mul.3794 = bf16[1,41,14336]{2,1,0} multiply(%dot_general.1268, %convert_element_type.2051), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/mul" stack_frame_id=791} + %state_1__180_.1 = bf16[14336,4096]{1,0} parameter(183), metadata={op_name="state[1][180]"} + %dot_general.1269 = bf16[1,41,4096]{2,1,0} dot(%mul.3794, %state_1__180_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/feedforward_output_dense/dot_general" stack_frame_id=800} + %add.1262 = bf16[1,41,4096]{2,1,0} add(%dot_general.1269, %add.1260), metadata={op_name="jit(compiled_generate_function)/transformer_layer_19/add" stack_frame_id=801} + %convert_element_type.2054 = f32[1,41,4096]{2,1,0} convert(%add.1262), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention_layernorm/convert_element_type" stack_frame_id=555} + %pow.464 = f32[1,41,4096]{2,1,0} power(%convert_element_type.2054, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention_layernorm/pow" stack_frame_id=558} + %reduce_sum.1509 = f32[1,41]{1,0} reduce(%pow.464, %constant.155), dimensions={2}, to_apply=%region_81.86, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention_layernorm/reduce_sum" stack_frame_id=24} + %broadcast_in_dim.1857 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1509), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention_layernorm/broadcast_in_dim" stack_frame_id=24} + %div.1169 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1857, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention_layernorm/div" stack_frame_id=24} + %add.1264 = f32[1,41,1]{2,1,0} add(%div.1169, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention_layernorm/add" stack_frame_id=559} + %rsqrt.233 = f32[1,41,1]{2,1,0} rsqrt(%add.1264), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention_layernorm/rsqrt" stack_frame_id=562} + %mul.3795 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.233), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3796 = f32[1,41]{1,0} reshape(%mul.3795), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3797 = f32[1,41,4096]{2,1,0} broadcast(%mul.3796), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3798 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.2054, %mul.3797), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention_layernorm/mul" stack_frame_id=563} + %state_1__186_.1 = bf16[4096]{0} parameter(189), metadata={op_name="state[1][186]"} + %convert_element_type.2055 = f32[4096]{0} convert(%state_1__186_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention_layernorm/convert_element_type" stack_frame_id=564} + %broadcast_in_dim.1858 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.2055), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention_layernorm/broadcast_in_dim" stack_frame_id=564} + %mul.3799 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1858), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3800 = f32[1,4096]{1,0} reshape(%mul.3799), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3801 = f32[1,41,4096]{2,1,0} broadcast(%mul.3800), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3802 = f32[1,41,4096]{2,1,0} multiply(%mul.3798, %mul.3801), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention_layernorm/mul" stack_frame_id=564} + %convert_element_type.2056 = bf16[1,41,4096]{2,1,0} convert(%mul.3802), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention_layernorm/convert_element_type" stack_frame_id=568} + %state_1__183_.1 = bf16[4096,8,128]{2,1,0} parameter(186), metadata={op_name="state[1][183]"} + %dot_general.1272 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.2056, %state_1__183_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=647} + %iota.677 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/iota" stack_frame_id=677} + %broadcast_in_dim.1860 = f32[1,41]{1,0} reshape(%iota.677), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/broadcast_in_dim" stack_frame_id=680} + %iota.676 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/iota" stack_frame_id=667} + %mul.3812 = f32[64]{0} multiply(%iota.676, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=667} + %div.1172 = f32[64]{0} divide(%mul.3812, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/div" stack_frame_id=668} + %neg.472 = f32[64]{0} negate(%div.1172), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/neg" stack_frame_id=669} + %pow.466 = f32[64]{0} power(%broadcast.51, %neg.472), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/pow" stack_frame_id=672} + %div.1173 = f32[64]{0} divide(%pow.466, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/div" stack_frame_id=673} + %dot_general.1274 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1860, %div.1173), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/bi,j->bij/dot_general" stack_frame_id=683} + %stack.1839 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1274), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/stack" stack_frame_id=686} + %stack.1840 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1274), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/stack" stack_frame_id=686} + %stack.1841 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1839, %stack.1840), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/stack" stack_frame_id=686} + %reshape.860 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1841), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/reshape"} + %cos.232 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.860), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/cos" stack_frame_id=690} + %convert_element_type.2059 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.232), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/convert_element_type" stack_frame_id=694} + %mul.3813 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2059), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=714} + %mul.3814 = bf16[1,41,128]{2,1,0} reshape(%mul.3813), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=714} + %mul.3815 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3814), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=714} + %mul.3816 = bf16[1,41,8,128]{3,2,1,0} multiply(%dot_general.1272, %mul.3815), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=714} + %split.465 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1272), slice={[0:1], [0:41], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/split" stack_frame_id=706} + %neg.473 = bf16[1,41,8,64]{3,2,1,0} negate(%split.465), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/neg" stack_frame_id=707} + %stack.1842 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%neg.473), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/stack" stack_frame_id=710} + %split.464 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1272), slice={[0:1], [0:41], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/split" stack_frame_id=706} + %stack.1843 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%split.464), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/stack" stack_frame_id=710} + %stack.1844 = bf16[1,41,8,2,64]{4,3,2,1,0} concatenate(%stack.1842, %stack.1843), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/stack" stack_frame_id=710} + %reshape.861 = bf16[1,41,8,128]{3,2,1,0} reshape(%stack.1844), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/reshape" stack_frame_id=713} + %sin.232 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.860), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/sin" stack_frame_id=698} + %convert_element_type.2060 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.232), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/convert_element_type" stack_frame_id=702} + %mul.3817 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2060), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=715} + %mul.3818 = bf16[1,41,128]{2,1,0} reshape(%mul.3817), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=715} + %mul.3819 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3818), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=715} + %mul.3820 = bf16[1,41,8,128]{3,2,1,0} multiply(%reshape.861, %mul.3819), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=715} + %add.1266 = bf16[1,41,8,128]{3,2,1,0} add(%mul.3816, %mul.3820), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/add" stack_frame_id=716} + %stack.1845 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%add.1266), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/stack" stack_frame_id=719} + %state_1__184_.1 = bf16[4096,8,128]{2,1,0} parameter(187), metadata={op_name="state[1][184]"} + %dot_general.1273 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.2056, %state_1__184_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=656} + %stack.1846 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dot_general.1273), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/stack" stack_frame_id=719} + %stack.1847 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1845, %stack.1846), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/stack" stack_frame_id=719} + %stack.2027 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1847), metadata={op_name="jit(compiled_generate_function)/stack" stack_frame_id=804} + %broadcast_in_dim.1862 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dot_general.1273), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/broadcast_in_dim" stack_frame_id=725} + %reshape.863 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1862), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/reshape" stack_frame_id=725} + %iota.670 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/iota" stack_frame_id=525} + %broadcast_in_dim.1851 = f32[41,1]{1,0} reshape(%iota.670), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/broadcast_in_dim" stack_frame_id=528} + %ge.621 = f32[41,1]{1,0} broadcast(%broadcast_in_dim.1851), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/ge" stack_frame_id=532} + %ge.622 = f32[41]{0} reshape(%ge.621), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/ge" stack_frame_id=532} + %ge.623 = f32[41,41]{1,0} broadcast(%ge.622), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/ge" stack_frame_id=532} + %iota.671 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/iota" stack_frame_id=531} + %broadcast_in_dim.1852 = f32[1,41]{1,0} reshape(%iota.671), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/broadcast_in_dim" stack_frame_id=532} + %ge.624 = f32[1,41]{1,0} broadcast(%broadcast_in_dim.1852), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/ge" stack_frame_id=532} + %ge.625 = f32[41]{0} reshape(%ge.624), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/ge" stack_frame_id=532} + %ge.626 = f32[41,41]{1,0} broadcast(%ge.625), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/ge" stack_frame_id=532} + %ge.627 = pred[41,41]{1,0} compare(%ge.623, %ge.626), direction=GE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/ge" stack_frame_id=532} + %broadcast_in_dim.1853 = pred[1,41,41]{2,1,0} reshape(%ge.627), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/broadcast_in_dim" stack_frame_id=535} + %convert_element_type.2053 = s32[1,41,41]{2,1,0} convert(%broadcast_in_dim.1853), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/convert_element_type" stack_frame_id=551} + %iota.672 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/iota" stack_frame_id=538} + %broadcast_in_dim.1854 = s32[41,1]{1,0} reshape(%iota.672), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/broadcast_in_dim" stack_frame_id=536} + %lt.763 = s32[41,1]{1,0} broadcast(%broadcast_in_dim.1854), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/lt" stack_frame_id=543} + %lt.764 = s32[41]{0} reshape(%lt.763), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/lt" stack_frame_id=543} + %lt.765 = s32[41,41]{1,0} broadcast(%lt.764), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/lt" stack_frame_id=543} + %iota.673 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/iota" stack_frame_id=541} + %broadcast_in_dim.1855 = s32[1,41]{1,0} reshape(%iota.673), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/broadcast_in_dim" stack_frame_id=539} + %add.1263 = s32[1,41]{1,0} add(%broadcast_in_dim.1855, %broadcast.57), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/add" stack_frame_id=542} + %lt.766 = s32[1,41]{1,0} broadcast(%add.1263), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/lt" stack_frame_id=543} + %lt.767 = s32[41]{0} reshape(%lt.766), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/lt" stack_frame_id=543} + %lt.768 = s32[41,41]{1,0} broadcast(%lt.767), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/lt" stack_frame_id=543} + %lt.769 = pred[41,41]{1,0} compare(%lt.765, %lt.768), direction=LT, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/lt" stack_frame_id=543} + %convert_element_type.2052 = s32[41,41]{1,0} convert(%lt.769), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/convert_element_type" stack_frame_id=547} + %broadcast_in_dim.1856 = s32[1,41,41]{2,1,0} reshape(%convert_element_type.2052), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/broadcast_in_dim" stack_frame_id=551} + %min.115 = s32[1,41,41]{2,1,0} minimum(%convert_element_type.2053, %broadcast_in_dim.1856), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/min" stack_frame_id=551} + %broadcast_in_dim.1863 = s32[1,1,41,41]{3,2,1,0} reshape(%min.115), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/broadcast_in_dim" stack_frame_id=728} + %convert_element_type.2061 = pred[1,1,41,41]{3,2,1,0} compare(%broadcast_in_dim.1863, %broadcast.49), direction=NE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/convert_element_type" stack_frame_id=732} + %broadcast_in_dim.1864 = pred[1,1,1,41,41]{4,3,2,1,0} reshape(%convert_element_type.2061), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/broadcast_in_dim" stack_frame_id=34} + %and.358 = pred[1,1,1,41,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1864), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/vmap()/and" stack_frame_id=34} + %and.359 = pred[1,1,41,41]{3,2,1,0} reshape(%and.358), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/vmap()/and" stack_frame_id=34} + %and.360 = pred[1,1,32,41,41]{4,3,2,1,0} broadcast(%and.359), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/vmap()/and" stack_frame_id=34} + %state_1__182_.1 = bf16[4096,32,128]{2,1,0} parameter(185), metadata={op_name="state[1][182]"} + %dot_general.1270 = bf16[1,41,32,128]{3,2,1,0} dot(%convert_element_type.2056, %state_1__182_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=577} + %iota.675 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/iota" stack_frame_id=598} + %broadcast_in_dim.1859 = f32[1,41]{1,0} reshape(%iota.675), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/broadcast_in_dim" stack_frame_id=601} + %iota.674 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/iota" stack_frame_id=588} + %mul.3803 = f32[64]{0} multiply(%iota.674, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=588} + %div.1170 = f32[64]{0} divide(%mul.3803, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/div" stack_frame_id=589} + %neg.470 = f32[64]{0} negate(%div.1170), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/neg" stack_frame_id=590} + %pow.465 = f32[64]{0} power(%broadcast.51, %neg.470), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/pow" stack_frame_id=593} + %div.1171 = f32[64]{0} divide(%pow.465, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/div" stack_frame_id=594} + %dot_general.1271 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1859, %div.1171), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/bi,j->bij/dot_general" stack_frame_id=604} + %stack.1833 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1271), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/stack" stack_frame_id=607} + %stack.1834 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1271), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/stack" stack_frame_id=607} + %stack.1835 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1833, %stack.1834), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/stack" stack_frame_id=607} + %reshape.858 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1835), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/reshape"} + %cos.231 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.858), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/cos" stack_frame_id=611} + %convert_element_type.2057 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.231), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/convert_element_type" stack_frame_id=615} + %mul.3804 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2057), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=635} + %mul.3805 = bf16[1,41,128]{2,1,0} reshape(%mul.3804), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=635} + %mul.3806 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3805), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=635} + %mul.3807 = bf16[1,41,32,128]{3,2,1,0} multiply(%dot_general.1270, %mul.3806), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=635} + %split.463 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1270), slice={[0:1], [0:41], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/split" stack_frame_id=627} + %neg.471 = bf16[1,41,32,64]{3,2,1,0} negate(%split.463), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/neg" stack_frame_id=628} + %stack.1836 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%neg.471), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/stack" stack_frame_id=631} + %split.462 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1270), slice={[0:1], [0:41], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/split" stack_frame_id=627} + %stack.1837 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%split.462), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/stack" stack_frame_id=631} + %stack.1838 = bf16[1,41,32,2,64]{4,3,2,1,0} concatenate(%stack.1836, %stack.1837), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/stack" stack_frame_id=631} + %reshape.859 = bf16[1,41,32,128]{3,2,1,0} reshape(%stack.1838), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/reshape" stack_frame_id=634} + %sin.231 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.858), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/sin" stack_frame_id=619} + %convert_element_type.2058 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.231), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/convert_element_type" stack_frame_id=623} + %mul.3808 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2058), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=636} + %mul.3809 = bf16[1,41,128]{2,1,0} reshape(%mul.3808), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=636} + %mul.3810 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3809), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=636} + %mul.3811 = bf16[1,41,32,128]{3,2,1,0} multiply(%reshape.859, %mul.3810), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/mul" stack_frame_id=636} + %add.1265 = bf16[1,41,32,128]{3,2,1,0} add(%mul.3807, %mul.3811), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/rotary_embedding_20/add" stack_frame_id=637} + %reshape.864 = bf16[1,41,32,1,128]{4,3,2,1,0} reshape(%add.1265), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/reshape" stack_frame_id=34} + %broadcast_in_dim.1861 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%add.1266), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/broadcast_in_dim" stack_frame_id=722} + %reshape.862 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1861), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/reshape" stack_frame_id=722} + %dot_general.1275 = f32[1,32,41,1,41]{4,3,2,1,0} dot(%reshape.864, %reshape.862), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=34} + %mul.3821 = f32[1,32,41,1,41]{4,3,2,1,0} multiply(%dot_general.1275, %broadcast.48), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/vmap()/mul" stack_frame_id=34} + %vmap_jit__where__.115 = f32[1,1,32,41,41]{4,3,2,1,0} call(%and.360, %mul.3821, %constant.153), to_apply=%_where_0.5, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/vmap(jit(_where))" stack_frame_id=34} + %reduce_max.501 = f32[1,1,32,41]{3,2,1,0} reduce(%vmap_jit__where__.115, %constant.152), dimensions={4}, to_apply=%region_82.87, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/vmap()/reduce_max" stack_frame_id=34} + %max.125 = f32[1,1,32,41]{3,2,1,0} maximum(%reduce_max.501, %broadcast.47), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/vmap()/max" stack_frame_id=34} + %broadcast_in_dim.1865 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%max.125), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %sub.484 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1865), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/vmap()/sub" stack_frame_id=34} + %sub.485 = f32[1,1,32,41]{3,2,1,0} reshape(%sub.484), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/vmap()/sub" stack_frame_id=34} + %sub.486 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%sub.485), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/vmap()/sub" stack_frame_id=34} + %sub.487 = f32[1,1,32,41,41]{4,3,2,1,0} subtract(%vmap_jit__where__.115, %sub.486), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/vmap()/sub" stack_frame_id=34} + %exp.121 = f32[1,1,32,41,41]{4,3,2,1,0} exponential(%sub.487), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/vmap()/exp" stack_frame_id=34} + %reduce_sum.1510 = f32[1,1,32,41]{3,2,1,0} reduce(%exp.121, %constant.155), dimensions={4}, to_apply=%region_83.88, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/vmap()/reduce_sum" stack_frame_id=34} + %broadcast_in_dim.1866 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%reduce_sum.1510), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %div.1174 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1866), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/vmap()/div" stack_frame_id=34} + %div.1175 = f32[1,1,32,41]{3,2,1,0} reshape(%div.1174), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/vmap()/div" stack_frame_id=34} + %div.1176 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%div.1175), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/vmap()/div" stack_frame_id=34} + %div.1177 = f32[1,1,32,41,41]{4,3,2,1,0} divide(%exp.121, %div.1176), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/vmap()/div" stack_frame_id=34} + %convert_element_type.2062 = bf16[1,1,32,41,41]{4,3,2,1,0} convert(%div.1177), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/vmap()/convert_element_type" stack_frame_id=34} + %dot_general.1276 = bf16[1,32,128,1,41]{4,3,2,1,0} dot(%reshape.863, %convert_element_type.2062), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=34} + %transpose.119 = bf16[1,41,32,1,128]{1,3,4,2,0} transpose(%dot_general.1276), dimensions={0,4,1,3,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/vmap()/transpose;jit(compiled_generate_function)/transformer_layer_20/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %reshape.865 = bf16[1,41,32,128]{3,2,1,0} reshape(%transpose.119), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/reshape" stack_frame_id=34} + %state_1__185_.1 = bf16[32,128,4096]{2,1,0} parameter(188), metadata={op_name="state[1][185]"} + %dot_general.1277 = bf16[1,41,4096]{2,1,0} dot(%reshape.865, %state_1__185_.1), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=741} + %add.1267 = bf16[1,41,4096]{2,1,0} add(%dot_general.1277, %add.1262), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/add" stack_frame_id=742} + %convert_element_type.2063 = f32[1,41,4096]{2,1,0} convert(%add.1267), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/feedforward_layernorm/convert_element_type" stack_frame_id=746} + %pow.467 = f32[1,41,4096]{2,1,0} power(%convert_element_type.2063, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/feedforward_layernorm/pow" stack_frame_id=749} + %reduce_sum.1511 = f32[1,41]{1,0} reduce(%pow.467, %constant.155), dimensions={2}, to_apply=%region_84.89, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/feedforward_layernorm/reduce_sum" stack_frame_id=43} + %broadcast_in_dim.1867 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1511), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/feedforward_layernorm/broadcast_in_dim" stack_frame_id=43} + %div.1178 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1867, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/feedforward_layernorm/div" stack_frame_id=43} + %add.1268 = f32[1,41,1]{2,1,0} add(%div.1178, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/feedforward_layernorm/add" stack_frame_id=750} + %rsqrt.234 = f32[1,41,1]{2,1,0} rsqrt(%add.1268), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/feedforward_layernorm/rsqrt" stack_frame_id=753} + %mul.3822 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.234), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3823 = f32[1,41]{1,0} reshape(%mul.3822), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3824 = f32[1,41,4096]{2,1,0} broadcast(%mul.3823), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3825 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.2063, %mul.3824), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/feedforward_layernorm/mul" stack_frame_id=754} + %state_1__190_.1 = bf16[4096]{0} parameter(193), metadata={op_name="state[1][190]"} + %convert_element_type.2064 = f32[4096]{0} convert(%state_1__190_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/feedforward_layernorm/convert_element_type" stack_frame_id=755} + %broadcast_in_dim.1868 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.2064), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/feedforward_layernorm/broadcast_in_dim" stack_frame_id=755} + %mul.3826 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1868), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3827 = f32[1,4096]{1,0} reshape(%mul.3826), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3828 = f32[1,41,4096]{2,1,0} broadcast(%mul.3827), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3829 = f32[1,41,4096]{2,1,0} multiply(%mul.3825, %mul.3828), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/feedforward_layernorm/mul" stack_frame_id=755} + %convert_element_type.2065 = bf16[1,41,4096]{2,1,0} convert(%mul.3829), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/feedforward_layernorm/convert_element_type" stack_frame_id=759} + %state_1__187_.1 = bf16[4096,14336]{1,0} parameter(190), metadata={op_name="state[1][187]"} + %dot_general.1279 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.2065, %state_1__187_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/feedforward_intermediate_dense/dot_general" stack_frame_id=788} + %state_1__188_.1 = bf16[4096,14336]{1,0} parameter(191), metadata={op_name="state[1][188]"} + %dot_general.1278 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.2065, %state_1__188_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/feedforward_gate_dense/dot_general" stack_frame_id=768} + %convert_element_type.2066 = f32[1,41,14336]{2,1,0} convert(%dot_general.1278), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/convert_element_type" stack_frame_id=772} + %jit_silu_.115 = f32[1,41,14336]{2,1,0} call(%convert_element_type.2066), to_apply=%silu.9, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/jit(silu)" stack_frame_id=775} + %convert_element_type.2067 = bf16[1,41,14336]{2,1,0} convert(%jit_silu_.115), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/convert_element_type" stack_frame_id=779} + %mul.3830 = bf16[1,41,14336]{2,1,0} multiply(%dot_general.1279, %convert_element_type.2067), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/mul" stack_frame_id=791} + %state_1__189_.1 = bf16[14336,4096]{1,0} parameter(192), metadata={op_name="state[1][189]"} + %dot_general.1280 = bf16[1,41,4096]{2,1,0} dot(%mul.3830, %state_1__189_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/feedforward_output_dense/dot_general" stack_frame_id=800} + %add.1269 = bf16[1,41,4096]{2,1,0} add(%dot_general.1280, %add.1267), metadata={op_name="jit(compiled_generate_function)/transformer_layer_20/add" stack_frame_id=801} + %convert_element_type.2070 = f32[1,41,4096]{2,1,0} convert(%add.1269), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention_layernorm/convert_element_type" stack_frame_id=555} + %pow.468 = f32[1,41,4096]{2,1,0} power(%convert_element_type.2070, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention_layernorm/pow" stack_frame_id=558} + %reduce_sum.1512 = f32[1,41]{1,0} reduce(%pow.468, %constant.155), dimensions={2}, to_apply=%region_85.90, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention_layernorm/reduce_sum" stack_frame_id=24} + %broadcast_in_dim.1875 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1512), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention_layernorm/broadcast_in_dim" stack_frame_id=24} + %div.1179 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1875, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention_layernorm/div" stack_frame_id=24} + %add.1271 = f32[1,41,1]{2,1,0} add(%div.1179, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention_layernorm/add" stack_frame_id=559} + %rsqrt.235 = f32[1,41,1]{2,1,0} rsqrt(%add.1271), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention_layernorm/rsqrt" stack_frame_id=562} + %mul.3831 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.235), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3832 = f32[1,41]{1,0} reshape(%mul.3831), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3833 = f32[1,41,4096]{2,1,0} broadcast(%mul.3832), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3834 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.2070, %mul.3833), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention_layernorm/mul" stack_frame_id=563} + %state_1__195_.1 = bf16[4096]{0} parameter(198), metadata={op_name="state[1][195]"} + %convert_element_type.2071 = f32[4096]{0} convert(%state_1__195_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention_layernorm/convert_element_type" stack_frame_id=564} + %broadcast_in_dim.1876 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.2071), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention_layernorm/broadcast_in_dim" stack_frame_id=564} + %mul.3835 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1876), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3836 = f32[1,4096]{1,0} reshape(%mul.3835), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3837 = f32[1,41,4096]{2,1,0} broadcast(%mul.3836), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3838 = f32[1,41,4096]{2,1,0} multiply(%mul.3834, %mul.3837), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention_layernorm/mul" stack_frame_id=564} + %convert_element_type.2072 = bf16[1,41,4096]{2,1,0} convert(%mul.3838), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention_layernorm/convert_element_type" stack_frame_id=568} + %state_1__192_.1 = bf16[4096,8,128]{2,1,0} parameter(195), metadata={op_name="state[1][192]"} + %dot_general.1283 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.2072, %state_1__192_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=647} + %iota.685 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/iota" stack_frame_id=677} + %broadcast_in_dim.1878 = f32[1,41]{1,0} reshape(%iota.685), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/broadcast_in_dim" stack_frame_id=680} + %iota.684 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/iota" stack_frame_id=667} + %mul.3848 = f32[64]{0} multiply(%iota.684, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=667} + %div.1182 = f32[64]{0} divide(%mul.3848, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/div" stack_frame_id=668} + %neg.476 = f32[64]{0} negate(%div.1182), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/neg" stack_frame_id=669} + %pow.470 = f32[64]{0} power(%broadcast.51, %neg.476), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/pow" stack_frame_id=672} + %div.1183 = f32[64]{0} divide(%pow.470, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/div" stack_frame_id=673} + %dot_general.1285 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1878, %div.1183), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/bi,j->bij/dot_general" stack_frame_id=683} + %stack.1854 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1285), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/stack" stack_frame_id=686} + %stack.1855 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1285), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/stack" stack_frame_id=686} + %stack.1856 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1854, %stack.1855), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/stack" stack_frame_id=686} + %reshape.868 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1856), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/reshape"} + %cos.234 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.868), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/cos" stack_frame_id=690} + %convert_element_type.2075 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.234), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/convert_element_type" stack_frame_id=694} + %mul.3849 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2075), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=714} + %mul.3850 = bf16[1,41,128]{2,1,0} reshape(%mul.3849), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=714} + %mul.3851 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3850), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=714} + %mul.3852 = bf16[1,41,8,128]{3,2,1,0} multiply(%dot_general.1283, %mul.3851), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=714} + %split.469 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1283), slice={[0:1], [0:41], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/split" stack_frame_id=706} + %neg.477 = bf16[1,41,8,64]{3,2,1,0} negate(%split.469), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/neg" stack_frame_id=707} + %stack.1857 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%neg.477), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/stack" stack_frame_id=710} + %split.468 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1283), slice={[0:1], [0:41], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/split" stack_frame_id=706} + %stack.1858 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%split.468), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/stack" stack_frame_id=710} + %stack.1859 = bf16[1,41,8,2,64]{4,3,2,1,0} concatenate(%stack.1857, %stack.1858), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/stack" stack_frame_id=710} + %reshape.869 = bf16[1,41,8,128]{3,2,1,0} reshape(%stack.1859), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/reshape" stack_frame_id=713} + %sin.234 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.868), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/sin" stack_frame_id=698} + %convert_element_type.2076 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.234), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/convert_element_type" stack_frame_id=702} + %mul.3853 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2076), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=715} + %mul.3854 = bf16[1,41,128]{2,1,0} reshape(%mul.3853), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=715} + %mul.3855 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3854), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=715} + %mul.3856 = bf16[1,41,8,128]{3,2,1,0} multiply(%reshape.869, %mul.3855), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=715} + %add.1273 = bf16[1,41,8,128]{3,2,1,0} add(%mul.3852, %mul.3856), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/add" stack_frame_id=716} + %stack.1860 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%add.1273), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/stack" stack_frame_id=719} + %state_1__193_.1 = bf16[4096,8,128]{2,1,0} parameter(196), metadata={op_name="state[1][193]"} + %dot_general.1284 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.2072, %state_1__193_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=656} + %stack.1861 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dot_general.1284), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/stack" stack_frame_id=719} + %stack.1862 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1860, %stack.1861), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/stack" stack_frame_id=719} + %stack.2028 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1862), metadata={op_name="jit(compiled_generate_function)/stack" stack_frame_id=804} + %broadcast_in_dim.1880 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dot_general.1284), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/broadcast_in_dim" stack_frame_id=725} + %reshape.871 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1880), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/reshape" stack_frame_id=725} + %iota.678 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/iota" stack_frame_id=525} + %broadcast_in_dim.1869 = f32[41,1]{1,0} reshape(%iota.678), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/broadcast_in_dim" stack_frame_id=528} + %ge.628 = f32[41,1]{1,0} broadcast(%broadcast_in_dim.1869), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/ge" stack_frame_id=532} + %ge.629 = f32[41]{0} reshape(%ge.628), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/ge" stack_frame_id=532} + %ge.630 = f32[41,41]{1,0} broadcast(%ge.629), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/ge" stack_frame_id=532} + %iota.679 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/iota" stack_frame_id=531} + %broadcast_in_dim.1870 = f32[1,41]{1,0} reshape(%iota.679), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/broadcast_in_dim" stack_frame_id=532} + %ge.631 = f32[1,41]{1,0} broadcast(%broadcast_in_dim.1870), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/ge" stack_frame_id=532} + %ge.632 = f32[41]{0} reshape(%ge.631), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/ge" stack_frame_id=532} + %ge.633 = f32[41,41]{1,0} broadcast(%ge.632), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/ge" stack_frame_id=532} + %ge.634 = pred[41,41]{1,0} compare(%ge.630, %ge.633), direction=GE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/ge" stack_frame_id=532} + %broadcast_in_dim.1871 = pred[1,41,41]{2,1,0} reshape(%ge.634), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/broadcast_in_dim" stack_frame_id=535} + %convert_element_type.2069 = s32[1,41,41]{2,1,0} convert(%broadcast_in_dim.1871), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/convert_element_type" stack_frame_id=551} + %iota.680 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/iota" stack_frame_id=538} + %broadcast_in_dim.1872 = s32[41,1]{1,0} reshape(%iota.680), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/broadcast_in_dim" stack_frame_id=536} + %lt.770 = s32[41,1]{1,0} broadcast(%broadcast_in_dim.1872), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/lt" stack_frame_id=543} + %lt.771 = s32[41]{0} reshape(%lt.770), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/lt" stack_frame_id=543} + %lt.772 = s32[41,41]{1,0} broadcast(%lt.771), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/lt" stack_frame_id=543} + %iota.681 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/iota" stack_frame_id=541} + %broadcast_in_dim.1873 = s32[1,41]{1,0} reshape(%iota.681), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/broadcast_in_dim" stack_frame_id=539} + %add.1270 = s32[1,41]{1,0} add(%broadcast_in_dim.1873, %broadcast.57), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/add" stack_frame_id=542} + %lt.773 = s32[1,41]{1,0} broadcast(%add.1270), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/lt" stack_frame_id=543} + %lt.774 = s32[41]{0} reshape(%lt.773), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/lt" stack_frame_id=543} + %lt.775 = s32[41,41]{1,0} broadcast(%lt.774), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/lt" stack_frame_id=543} + %lt.776 = pred[41,41]{1,0} compare(%lt.772, %lt.775), direction=LT, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/lt" stack_frame_id=543} + %convert_element_type.2068 = s32[41,41]{1,0} convert(%lt.776), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/convert_element_type" stack_frame_id=547} + %broadcast_in_dim.1874 = s32[1,41,41]{2,1,0} reshape(%convert_element_type.2068), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/broadcast_in_dim" stack_frame_id=551} + %min.116 = s32[1,41,41]{2,1,0} minimum(%convert_element_type.2069, %broadcast_in_dim.1874), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/min" stack_frame_id=551} + %broadcast_in_dim.1881 = s32[1,1,41,41]{3,2,1,0} reshape(%min.116), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/broadcast_in_dim" stack_frame_id=728} + %convert_element_type.2077 = pred[1,1,41,41]{3,2,1,0} compare(%broadcast_in_dim.1881, %broadcast.49), direction=NE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/convert_element_type" stack_frame_id=732} + %broadcast_in_dim.1882 = pred[1,1,1,41,41]{4,3,2,1,0} reshape(%convert_element_type.2077), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/broadcast_in_dim" stack_frame_id=34} + %and.361 = pred[1,1,1,41,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1882), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/vmap()/and" stack_frame_id=34} + %and.362 = pred[1,1,41,41]{3,2,1,0} reshape(%and.361), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/vmap()/and" stack_frame_id=34} + %and.363 = pred[1,1,32,41,41]{4,3,2,1,0} broadcast(%and.362), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/vmap()/and" stack_frame_id=34} + %state_1__191_.1 = bf16[4096,32,128]{2,1,0} parameter(194), metadata={op_name="state[1][191]"} + %dot_general.1281 = bf16[1,41,32,128]{3,2,1,0} dot(%convert_element_type.2072, %state_1__191_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=577} + %iota.683 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/iota" stack_frame_id=598} + %broadcast_in_dim.1877 = f32[1,41]{1,0} reshape(%iota.683), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/broadcast_in_dim" stack_frame_id=601} + %iota.682 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/iota" stack_frame_id=588} + %mul.3839 = f32[64]{0} multiply(%iota.682, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=588} + %div.1180 = f32[64]{0} divide(%mul.3839, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/div" stack_frame_id=589} + %neg.474 = f32[64]{0} negate(%div.1180), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/neg" stack_frame_id=590} + %pow.469 = f32[64]{0} power(%broadcast.51, %neg.474), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/pow" stack_frame_id=593} + %div.1181 = f32[64]{0} divide(%pow.469, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/div" stack_frame_id=594} + %dot_general.1282 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1877, %div.1181), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/bi,j->bij/dot_general" stack_frame_id=604} + %stack.1848 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1282), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/stack" stack_frame_id=607} + %stack.1849 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1282), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/stack" stack_frame_id=607} + %stack.1850 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1848, %stack.1849), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/stack" stack_frame_id=607} + %reshape.866 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1850), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/reshape"} + %cos.233 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.866), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/cos" stack_frame_id=611} + %convert_element_type.2073 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.233), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/convert_element_type" stack_frame_id=615} + %mul.3840 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2073), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=635} + %mul.3841 = bf16[1,41,128]{2,1,0} reshape(%mul.3840), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=635} + %mul.3842 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3841), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=635} + %mul.3843 = bf16[1,41,32,128]{3,2,1,0} multiply(%dot_general.1281, %mul.3842), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=635} + %split.467 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1281), slice={[0:1], [0:41], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/split" stack_frame_id=627} + %neg.475 = bf16[1,41,32,64]{3,2,1,0} negate(%split.467), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/neg" stack_frame_id=628} + %stack.1851 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%neg.475), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/stack" stack_frame_id=631} + %split.466 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1281), slice={[0:1], [0:41], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/split" stack_frame_id=627} + %stack.1852 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%split.466), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/stack" stack_frame_id=631} + %stack.1853 = bf16[1,41,32,2,64]{4,3,2,1,0} concatenate(%stack.1851, %stack.1852), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/stack" stack_frame_id=631} + %reshape.867 = bf16[1,41,32,128]{3,2,1,0} reshape(%stack.1853), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/reshape" stack_frame_id=634} + %sin.233 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.866), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/sin" stack_frame_id=619} + %convert_element_type.2074 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.233), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/convert_element_type" stack_frame_id=623} + %mul.3844 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2074), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=636} + %mul.3845 = bf16[1,41,128]{2,1,0} reshape(%mul.3844), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=636} + %mul.3846 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3845), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=636} + %mul.3847 = bf16[1,41,32,128]{3,2,1,0} multiply(%reshape.867, %mul.3846), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/mul" stack_frame_id=636} + %add.1272 = bf16[1,41,32,128]{3,2,1,0} add(%mul.3843, %mul.3847), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/rotary_embedding_21/add" stack_frame_id=637} + %reshape.872 = bf16[1,41,32,1,128]{4,3,2,1,0} reshape(%add.1272), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/reshape" stack_frame_id=34} + %broadcast_in_dim.1879 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%add.1273), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/broadcast_in_dim" stack_frame_id=722} + %reshape.870 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1879), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/reshape" stack_frame_id=722} + %dot_general.1286 = f32[1,32,41,1,41]{4,3,2,1,0} dot(%reshape.872, %reshape.870), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=34} + %mul.3857 = f32[1,32,41,1,41]{4,3,2,1,0} multiply(%dot_general.1286, %broadcast.48), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/vmap()/mul" stack_frame_id=34} + %vmap_jit__where__.116 = f32[1,1,32,41,41]{4,3,2,1,0} call(%and.363, %mul.3857, %constant.153), to_apply=%_where_0.5, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/vmap(jit(_where))" stack_frame_id=34} + %reduce_max.502 = f32[1,1,32,41]{3,2,1,0} reduce(%vmap_jit__where__.116, %constant.152), dimensions={4}, to_apply=%region_86.91, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/vmap()/reduce_max" stack_frame_id=34} + %max.126 = f32[1,1,32,41]{3,2,1,0} maximum(%reduce_max.502, %broadcast.47), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/vmap()/max" stack_frame_id=34} + %broadcast_in_dim.1883 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%max.126), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %sub.488 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1883), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/vmap()/sub" stack_frame_id=34} + %sub.489 = f32[1,1,32,41]{3,2,1,0} reshape(%sub.488), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/vmap()/sub" stack_frame_id=34} + %sub.490 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%sub.489), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/vmap()/sub" stack_frame_id=34} + %sub.491 = f32[1,1,32,41,41]{4,3,2,1,0} subtract(%vmap_jit__where__.116, %sub.490), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/vmap()/sub" stack_frame_id=34} + %exp.122 = f32[1,1,32,41,41]{4,3,2,1,0} exponential(%sub.491), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/vmap()/exp" stack_frame_id=34} + %reduce_sum.1513 = f32[1,1,32,41]{3,2,1,0} reduce(%exp.122, %constant.155), dimensions={4}, to_apply=%region_87.92, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/vmap()/reduce_sum" stack_frame_id=34} + %broadcast_in_dim.1884 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%reduce_sum.1513), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %div.1184 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1884), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/vmap()/div" stack_frame_id=34} + %div.1185 = f32[1,1,32,41]{3,2,1,0} reshape(%div.1184), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/vmap()/div" stack_frame_id=34} + %div.1186 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%div.1185), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/vmap()/div" stack_frame_id=34} + %div.1187 = f32[1,1,32,41,41]{4,3,2,1,0} divide(%exp.122, %div.1186), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/vmap()/div" stack_frame_id=34} + %convert_element_type.2078 = bf16[1,1,32,41,41]{4,3,2,1,0} convert(%div.1187), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/vmap()/convert_element_type" stack_frame_id=34} + %dot_general.1287 = bf16[1,32,128,1,41]{4,3,2,1,0} dot(%reshape.871, %convert_element_type.2078), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=34} + %transpose.120 = bf16[1,41,32,1,128]{1,3,4,2,0} transpose(%dot_general.1287), dimensions={0,4,1,3,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/vmap()/transpose;jit(compiled_generate_function)/transformer_layer_21/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %reshape.873 = bf16[1,41,32,128]{3,2,1,0} reshape(%transpose.120), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/reshape" stack_frame_id=34} + %state_1__194_.1 = bf16[32,128,4096]{2,1,0} parameter(197), metadata={op_name="state[1][194]"} + %dot_general.1288 = bf16[1,41,4096]{2,1,0} dot(%reshape.873, %state_1__194_.1), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=741} + %add.1274 = bf16[1,41,4096]{2,1,0} add(%dot_general.1288, %add.1269), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/add" stack_frame_id=742} + %convert_element_type.2079 = f32[1,41,4096]{2,1,0} convert(%add.1274), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/feedforward_layernorm/convert_element_type" stack_frame_id=746} + %pow.471 = f32[1,41,4096]{2,1,0} power(%convert_element_type.2079, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/feedforward_layernorm/pow" stack_frame_id=749} + %reduce_sum.1514 = f32[1,41]{1,0} reduce(%pow.471, %constant.155), dimensions={2}, to_apply=%region_88.93, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/feedforward_layernorm/reduce_sum" stack_frame_id=43} + %broadcast_in_dim.1885 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1514), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/feedforward_layernorm/broadcast_in_dim" stack_frame_id=43} + %div.1188 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1885, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/feedforward_layernorm/div" stack_frame_id=43} + %add.1275 = f32[1,41,1]{2,1,0} add(%div.1188, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/feedforward_layernorm/add" stack_frame_id=750} + %rsqrt.236 = f32[1,41,1]{2,1,0} rsqrt(%add.1275), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/feedforward_layernorm/rsqrt" stack_frame_id=753} + %mul.3858 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.236), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3859 = f32[1,41]{1,0} reshape(%mul.3858), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3860 = f32[1,41,4096]{2,1,0} broadcast(%mul.3859), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3861 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.2079, %mul.3860), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/feedforward_layernorm/mul" stack_frame_id=754} + %state_1__199_.1 = bf16[4096]{0} parameter(202), metadata={op_name="state[1][199]"} + %convert_element_type.2080 = f32[4096]{0} convert(%state_1__199_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/feedforward_layernorm/convert_element_type" stack_frame_id=755} + %broadcast_in_dim.1886 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.2080), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/feedforward_layernorm/broadcast_in_dim" stack_frame_id=755} + %mul.3862 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1886), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3863 = f32[1,4096]{1,0} reshape(%mul.3862), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3864 = f32[1,41,4096]{2,1,0} broadcast(%mul.3863), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3865 = f32[1,41,4096]{2,1,0} multiply(%mul.3861, %mul.3864), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/feedforward_layernorm/mul" stack_frame_id=755} + %convert_element_type.2081 = bf16[1,41,4096]{2,1,0} convert(%mul.3865), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/feedforward_layernorm/convert_element_type" stack_frame_id=759} + %state_1__196_.1 = bf16[4096,14336]{1,0} parameter(199), metadata={op_name="state[1][196]"} + %dot_general.1290 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.2081, %state_1__196_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/feedforward_intermediate_dense/dot_general" stack_frame_id=788} + %state_1__197_.1 = bf16[4096,14336]{1,0} parameter(200), metadata={op_name="state[1][197]"} + %dot_general.1289 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.2081, %state_1__197_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/feedforward_gate_dense/dot_general" stack_frame_id=768} + %convert_element_type.2082 = f32[1,41,14336]{2,1,0} convert(%dot_general.1289), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/convert_element_type" stack_frame_id=772} + %jit_silu_.116 = f32[1,41,14336]{2,1,0} call(%convert_element_type.2082), to_apply=%silu.9, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/jit(silu)" stack_frame_id=775} + %convert_element_type.2083 = bf16[1,41,14336]{2,1,0} convert(%jit_silu_.116), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/convert_element_type" stack_frame_id=779} + %mul.3866 = bf16[1,41,14336]{2,1,0} multiply(%dot_general.1290, %convert_element_type.2083), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/mul" stack_frame_id=791} + %state_1__198_.1 = bf16[14336,4096]{1,0} parameter(201), metadata={op_name="state[1][198]"} + %dot_general.1291 = bf16[1,41,4096]{2,1,0} dot(%mul.3866, %state_1__198_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/feedforward_output_dense/dot_general" stack_frame_id=800} + %add.1276 = bf16[1,41,4096]{2,1,0} add(%dot_general.1291, %add.1274), metadata={op_name="jit(compiled_generate_function)/transformer_layer_21/add" stack_frame_id=801} + %convert_element_type.2086 = f32[1,41,4096]{2,1,0} convert(%add.1276), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention_layernorm/convert_element_type" stack_frame_id=555} + %pow.472 = f32[1,41,4096]{2,1,0} power(%convert_element_type.2086, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention_layernorm/pow" stack_frame_id=558} + %reduce_sum.1515 = f32[1,41]{1,0} reduce(%pow.472, %constant.155), dimensions={2}, to_apply=%region_89.94, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention_layernorm/reduce_sum" stack_frame_id=24} + %broadcast_in_dim.1893 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1515), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention_layernorm/broadcast_in_dim" stack_frame_id=24} + %div.1189 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1893, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention_layernorm/div" stack_frame_id=24} + %add.1278 = f32[1,41,1]{2,1,0} add(%div.1189, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention_layernorm/add" stack_frame_id=559} + %rsqrt.237 = f32[1,41,1]{2,1,0} rsqrt(%add.1278), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention_layernorm/rsqrt" stack_frame_id=562} + %mul.3867 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.237), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3868 = f32[1,41]{1,0} reshape(%mul.3867), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3869 = f32[1,41,4096]{2,1,0} broadcast(%mul.3868), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3870 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.2086, %mul.3869), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention_layernorm/mul" stack_frame_id=563} + %state_1__204_.1 = bf16[4096]{0} parameter(207), metadata={op_name="state[1][204]"} + %convert_element_type.2087 = f32[4096]{0} convert(%state_1__204_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention_layernorm/convert_element_type" stack_frame_id=564} + %broadcast_in_dim.1894 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.2087), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention_layernorm/broadcast_in_dim" stack_frame_id=564} + %mul.3871 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1894), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3872 = f32[1,4096]{1,0} reshape(%mul.3871), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3873 = f32[1,41,4096]{2,1,0} broadcast(%mul.3872), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3874 = f32[1,41,4096]{2,1,0} multiply(%mul.3870, %mul.3873), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention_layernorm/mul" stack_frame_id=564} + %convert_element_type.2088 = bf16[1,41,4096]{2,1,0} convert(%mul.3874), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention_layernorm/convert_element_type" stack_frame_id=568} + %state_1__201_.1 = bf16[4096,8,128]{2,1,0} parameter(204), metadata={op_name="state[1][201]"} + %dot_general.1294 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.2088, %state_1__201_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=647} + %iota.693 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/iota" stack_frame_id=677} + %broadcast_in_dim.1896 = f32[1,41]{1,0} reshape(%iota.693), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/broadcast_in_dim" stack_frame_id=680} + %iota.692 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/iota" stack_frame_id=667} + %mul.3884 = f32[64]{0} multiply(%iota.692, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=667} + %div.1192 = f32[64]{0} divide(%mul.3884, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/div" stack_frame_id=668} + %neg.480 = f32[64]{0} negate(%div.1192), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/neg" stack_frame_id=669} + %pow.474 = f32[64]{0} power(%broadcast.51, %neg.480), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/pow" stack_frame_id=672} + %div.1193 = f32[64]{0} divide(%pow.474, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/div" stack_frame_id=673} + %dot_general.1296 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1896, %div.1193), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/bi,j->bij/dot_general" stack_frame_id=683} + %stack.1869 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1296), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/stack" stack_frame_id=686} + %stack.1870 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1296), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/stack" stack_frame_id=686} + %stack.1871 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1869, %stack.1870), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/stack" stack_frame_id=686} + %reshape.876 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1871), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/reshape"} + %cos.236 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.876), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/cos" stack_frame_id=690} + %convert_element_type.2091 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.236), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/convert_element_type" stack_frame_id=694} + %mul.3885 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2091), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=714} + %mul.3886 = bf16[1,41,128]{2,1,0} reshape(%mul.3885), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=714} + %mul.3887 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3886), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=714} + %mul.3888 = bf16[1,41,8,128]{3,2,1,0} multiply(%dot_general.1294, %mul.3887), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=714} + %split.473 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1294), slice={[0:1], [0:41], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/split" stack_frame_id=706} + %neg.481 = bf16[1,41,8,64]{3,2,1,0} negate(%split.473), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/neg" stack_frame_id=707} + %stack.1872 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%neg.481), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/stack" stack_frame_id=710} + %split.472 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1294), slice={[0:1], [0:41], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/split" stack_frame_id=706} + %stack.1873 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%split.472), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/stack" stack_frame_id=710} + %stack.1874 = bf16[1,41,8,2,64]{4,3,2,1,0} concatenate(%stack.1872, %stack.1873), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/stack" stack_frame_id=710} + %reshape.877 = bf16[1,41,8,128]{3,2,1,0} reshape(%stack.1874), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/reshape" stack_frame_id=713} + %sin.236 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.876), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/sin" stack_frame_id=698} + %convert_element_type.2092 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.236), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/convert_element_type" stack_frame_id=702} + %mul.3889 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2092), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=715} + %mul.3890 = bf16[1,41,128]{2,1,0} reshape(%mul.3889), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=715} + %mul.3891 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3890), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=715} + %mul.3892 = bf16[1,41,8,128]{3,2,1,0} multiply(%reshape.877, %mul.3891), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=715} + %add.1280 = bf16[1,41,8,128]{3,2,1,0} add(%mul.3888, %mul.3892), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/add" stack_frame_id=716} + %stack.1875 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%add.1280), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/stack" stack_frame_id=719} + %state_1__202_.1 = bf16[4096,8,128]{2,1,0} parameter(205), metadata={op_name="state[1][202]"} + %dot_general.1295 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.2088, %state_1__202_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=656} + %stack.1876 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dot_general.1295), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/stack" stack_frame_id=719} + %stack.1877 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1875, %stack.1876), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/stack" stack_frame_id=719} + %stack.2029 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1877), metadata={op_name="jit(compiled_generate_function)/stack" stack_frame_id=804} + %broadcast_in_dim.1898 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dot_general.1295), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/broadcast_in_dim" stack_frame_id=725} + %reshape.879 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1898), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/reshape" stack_frame_id=725} + %iota.686 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/iota" stack_frame_id=525} + %broadcast_in_dim.1887 = f32[41,1]{1,0} reshape(%iota.686), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/broadcast_in_dim" stack_frame_id=528} + %ge.635 = f32[41,1]{1,0} broadcast(%broadcast_in_dim.1887), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/ge" stack_frame_id=532} + %ge.636 = f32[41]{0} reshape(%ge.635), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/ge" stack_frame_id=532} + %ge.637 = f32[41,41]{1,0} broadcast(%ge.636), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/ge" stack_frame_id=532} + %iota.687 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/iota" stack_frame_id=531} + %broadcast_in_dim.1888 = f32[1,41]{1,0} reshape(%iota.687), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/broadcast_in_dim" stack_frame_id=532} + %ge.638 = f32[1,41]{1,0} broadcast(%broadcast_in_dim.1888), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/ge" stack_frame_id=532} + %ge.639 = f32[41]{0} reshape(%ge.638), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/ge" stack_frame_id=532} + %ge.640 = f32[41,41]{1,0} broadcast(%ge.639), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/ge" stack_frame_id=532} + %ge.641 = pred[41,41]{1,0} compare(%ge.637, %ge.640), direction=GE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/ge" stack_frame_id=532} + %broadcast_in_dim.1889 = pred[1,41,41]{2,1,0} reshape(%ge.641), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/broadcast_in_dim" stack_frame_id=535} + %convert_element_type.2085 = s32[1,41,41]{2,1,0} convert(%broadcast_in_dim.1889), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/convert_element_type" stack_frame_id=551} + %iota.688 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/iota" stack_frame_id=538} + %broadcast_in_dim.1890 = s32[41,1]{1,0} reshape(%iota.688), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/broadcast_in_dim" stack_frame_id=536} + %lt.777 = s32[41,1]{1,0} broadcast(%broadcast_in_dim.1890), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/lt" stack_frame_id=543} + %lt.778 = s32[41]{0} reshape(%lt.777), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/lt" stack_frame_id=543} + %lt.779 = s32[41,41]{1,0} broadcast(%lt.778), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/lt" stack_frame_id=543} + %iota.689 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/iota" stack_frame_id=541} + %broadcast_in_dim.1891 = s32[1,41]{1,0} reshape(%iota.689), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/broadcast_in_dim" stack_frame_id=539} + %add.1277 = s32[1,41]{1,0} add(%broadcast_in_dim.1891, %broadcast.57), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/add" stack_frame_id=542} + %lt.780 = s32[1,41]{1,0} broadcast(%add.1277), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/lt" stack_frame_id=543} + %lt.781 = s32[41]{0} reshape(%lt.780), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/lt" stack_frame_id=543} + %lt.782 = s32[41,41]{1,0} broadcast(%lt.781), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/lt" stack_frame_id=543} + %lt.783 = pred[41,41]{1,0} compare(%lt.779, %lt.782), direction=LT, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/lt" stack_frame_id=543} + %convert_element_type.2084 = s32[41,41]{1,0} convert(%lt.783), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/convert_element_type" stack_frame_id=547} + %broadcast_in_dim.1892 = s32[1,41,41]{2,1,0} reshape(%convert_element_type.2084), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/broadcast_in_dim" stack_frame_id=551} + %min.117 = s32[1,41,41]{2,1,0} minimum(%convert_element_type.2085, %broadcast_in_dim.1892), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/min" stack_frame_id=551} + %broadcast_in_dim.1899 = s32[1,1,41,41]{3,2,1,0} reshape(%min.117), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/broadcast_in_dim" stack_frame_id=728} + %convert_element_type.2093 = pred[1,1,41,41]{3,2,1,0} compare(%broadcast_in_dim.1899, %broadcast.49), direction=NE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/convert_element_type" stack_frame_id=732} + %broadcast_in_dim.1900 = pred[1,1,1,41,41]{4,3,2,1,0} reshape(%convert_element_type.2093), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/broadcast_in_dim" stack_frame_id=34} + %and.364 = pred[1,1,1,41,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1900), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/vmap()/and" stack_frame_id=34} + %and.365 = pred[1,1,41,41]{3,2,1,0} reshape(%and.364), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/vmap()/and" stack_frame_id=34} + %and.366 = pred[1,1,32,41,41]{4,3,2,1,0} broadcast(%and.365), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/vmap()/and" stack_frame_id=34} + %state_1__200_.1 = bf16[4096,32,128]{2,1,0} parameter(203), metadata={op_name="state[1][200]"} + %dot_general.1292 = bf16[1,41,32,128]{3,2,1,0} dot(%convert_element_type.2088, %state_1__200_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=577} + %iota.691 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/iota" stack_frame_id=598} + %broadcast_in_dim.1895 = f32[1,41]{1,0} reshape(%iota.691), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/broadcast_in_dim" stack_frame_id=601} + %iota.690 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/iota" stack_frame_id=588} + %mul.3875 = f32[64]{0} multiply(%iota.690, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=588} + %div.1190 = f32[64]{0} divide(%mul.3875, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/div" stack_frame_id=589} + %neg.478 = f32[64]{0} negate(%div.1190), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/neg" stack_frame_id=590} + %pow.473 = f32[64]{0} power(%broadcast.51, %neg.478), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/pow" stack_frame_id=593} + %div.1191 = f32[64]{0} divide(%pow.473, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/div" stack_frame_id=594} + %dot_general.1293 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1895, %div.1191), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/bi,j->bij/dot_general" stack_frame_id=604} + %stack.1863 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1293), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/stack" stack_frame_id=607} + %stack.1864 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1293), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/stack" stack_frame_id=607} + %stack.1865 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1863, %stack.1864), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/stack" stack_frame_id=607} + %reshape.874 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1865), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/reshape"} + %cos.235 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.874), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/cos" stack_frame_id=611} + %convert_element_type.2089 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.235), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/convert_element_type" stack_frame_id=615} + %mul.3876 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2089), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=635} + %mul.3877 = bf16[1,41,128]{2,1,0} reshape(%mul.3876), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=635} + %mul.3878 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3877), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=635} + %mul.3879 = bf16[1,41,32,128]{3,2,1,0} multiply(%dot_general.1292, %mul.3878), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=635} + %split.471 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1292), slice={[0:1], [0:41], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/split" stack_frame_id=627} + %neg.479 = bf16[1,41,32,64]{3,2,1,0} negate(%split.471), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/neg" stack_frame_id=628} + %stack.1866 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%neg.479), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/stack" stack_frame_id=631} + %split.470 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1292), slice={[0:1], [0:41], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/split" stack_frame_id=627} + %stack.1867 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%split.470), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/stack" stack_frame_id=631} + %stack.1868 = bf16[1,41,32,2,64]{4,3,2,1,0} concatenate(%stack.1866, %stack.1867), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/stack" stack_frame_id=631} + %reshape.875 = bf16[1,41,32,128]{3,2,1,0} reshape(%stack.1868), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/reshape" stack_frame_id=634} + %sin.235 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.874), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/sin" stack_frame_id=619} + %convert_element_type.2090 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.235), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/convert_element_type" stack_frame_id=623} + %mul.3880 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2090), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=636} + %mul.3881 = bf16[1,41,128]{2,1,0} reshape(%mul.3880), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=636} + %mul.3882 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3881), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=636} + %mul.3883 = bf16[1,41,32,128]{3,2,1,0} multiply(%reshape.875, %mul.3882), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/mul" stack_frame_id=636} + %add.1279 = bf16[1,41,32,128]{3,2,1,0} add(%mul.3879, %mul.3883), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/rotary_embedding_22/add" stack_frame_id=637} + %reshape.880 = bf16[1,41,32,1,128]{4,3,2,1,0} reshape(%add.1279), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/reshape" stack_frame_id=34} + %broadcast_in_dim.1897 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%add.1280), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/broadcast_in_dim" stack_frame_id=722} + %reshape.878 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1897), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/reshape" stack_frame_id=722} + %dot_general.1297 = f32[1,32,41,1,41]{4,3,2,1,0} dot(%reshape.880, %reshape.878), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=34} + %mul.3893 = f32[1,32,41,1,41]{4,3,2,1,0} multiply(%dot_general.1297, %broadcast.48), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/vmap()/mul" stack_frame_id=34} + %vmap_jit__where__.117 = f32[1,1,32,41,41]{4,3,2,1,0} call(%and.366, %mul.3893, %constant.153), to_apply=%_where_0.5, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/vmap(jit(_where))" stack_frame_id=34} + %reduce_max.503 = f32[1,1,32,41]{3,2,1,0} reduce(%vmap_jit__where__.117, %constant.152), dimensions={4}, to_apply=%region_90.95, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/vmap()/reduce_max" stack_frame_id=34} + %max.127 = f32[1,1,32,41]{3,2,1,0} maximum(%reduce_max.503, %broadcast.47), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/vmap()/max" stack_frame_id=34} + %broadcast_in_dim.1901 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%max.127), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %sub.492 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1901), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/vmap()/sub" stack_frame_id=34} + %sub.493 = f32[1,1,32,41]{3,2,1,0} reshape(%sub.492), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/vmap()/sub" stack_frame_id=34} + %sub.494 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%sub.493), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/vmap()/sub" stack_frame_id=34} + %sub.495 = f32[1,1,32,41,41]{4,3,2,1,0} subtract(%vmap_jit__where__.117, %sub.494), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/vmap()/sub" stack_frame_id=34} + %exp.123 = f32[1,1,32,41,41]{4,3,2,1,0} exponential(%sub.495), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/vmap()/exp" stack_frame_id=34} + %reduce_sum.1516 = f32[1,1,32,41]{3,2,1,0} reduce(%exp.123, %constant.155), dimensions={4}, to_apply=%region_91.96, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/vmap()/reduce_sum" stack_frame_id=34} + %broadcast_in_dim.1902 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%reduce_sum.1516), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %div.1194 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1902), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/vmap()/div" stack_frame_id=34} + %div.1195 = f32[1,1,32,41]{3,2,1,0} reshape(%div.1194), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/vmap()/div" stack_frame_id=34} + %div.1196 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%div.1195), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/vmap()/div" stack_frame_id=34} + %div.1197 = f32[1,1,32,41,41]{4,3,2,1,0} divide(%exp.123, %div.1196), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/vmap()/div" stack_frame_id=34} + %convert_element_type.2094 = bf16[1,1,32,41,41]{4,3,2,1,0} convert(%div.1197), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/vmap()/convert_element_type" stack_frame_id=34} + %dot_general.1298 = bf16[1,32,128,1,41]{4,3,2,1,0} dot(%reshape.879, %convert_element_type.2094), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=34} + %transpose.121 = bf16[1,41,32,1,128]{1,3,4,2,0} transpose(%dot_general.1298), dimensions={0,4,1,3,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/vmap()/transpose;jit(compiled_generate_function)/transformer_layer_22/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %reshape.881 = bf16[1,41,32,128]{3,2,1,0} reshape(%transpose.121), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/reshape" stack_frame_id=34} + %state_1__203_.1 = bf16[32,128,4096]{2,1,0} parameter(206), metadata={op_name="state[1][203]"} + %dot_general.1299 = bf16[1,41,4096]{2,1,0} dot(%reshape.881, %state_1__203_.1), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=741} + %add.1281 = bf16[1,41,4096]{2,1,0} add(%dot_general.1299, %add.1276), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/add" stack_frame_id=742} + %convert_element_type.2095 = f32[1,41,4096]{2,1,0} convert(%add.1281), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/feedforward_layernorm/convert_element_type" stack_frame_id=746} + %pow.475 = f32[1,41,4096]{2,1,0} power(%convert_element_type.2095, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/feedforward_layernorm/pow" stack_frame_id=749} + %reduce_sum.1517 = f32[1,41]{1,0} reduce(%pow.475, %constant.155), dimensions={2}, to_apply=%region_92.97, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/feedforward_layernorm/reduce_sum" stack_frame_id=43} + %broadcast_in_dim.1903 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1517), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/feedforward_layernorm/broadcast_in_dim" stack_frame_id=43} + %div.1198 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1903, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/feedforward_layernorm/div" stack_frame_id=43} + %add.1282 = f32[1,41,1]{2,1,0} add(%div.1198, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/feedforward_layernorm/add" stack_frame_id=750} + %rsqrt.238 = f32[1,41,1]{2,1,0} rsqrt(%add.1282), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/feedforward_layernorm/rsqrt" stack_frame_id=753} + %mul.3894 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.238), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3895 = f32[1,41]{1,0} reshape(%mul.3894), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3896 = f32[1,41,4096]{2,1,0} broadcast(%mul.3895), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3897 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.2095, %mul.3896), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/feedforward_layernorm/mul" stack_frame_id=754} + %state_1__208_.1 = bf16[4096]{0} parameter(211), metadata={op_name="state[1][208]"} + %convert_element_type.2096 = f32[4096]{0} convert(%state_1__208_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/feedforward_layernorm/convert_element_type" stack_frame_id=755} + %broadcast_in_dim.1904 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.2096), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/feedforward_layernorm/broadcast_in_dim" stack_frame_id=755} + %mul.3898 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1904), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3899 = f32[1,4096]{1,0} reshape(%mul.3898), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3900 = f32[1,41,4096]{2,1,0} broadcast(%mul.3899), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3901 = f32[1,41,4096]{2,1,0} multiply(%mul.3897, %mul.3900), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/feedforward_layernorm/mul" stack_frame_id=755} + %convert_element_type.2097 = bf16[1,41,4096]{2,1,0} convert(%mul.3901), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/feedforward_layernorm/convert_element_type" stack_frame_id=759} + %state_1__205_.1 = bf16[4096,14336]{1,0} parameter(208), metadata={op_name="state[1][205]"} + %dot_general.1301 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.2097, %state_1__205_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/feedforward_intermediate_dense/dot_general" stack_frame_id=788} + %state_1__206_.1 = bf16[4096,14336]{1,0} parameter(209), metadata={op_name="state[1][206]"} + %dot_general.1300 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.2097, %state_1__206_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/feedforward_gate_dense/dot_general" stack_frame_id=768} + %convert_element_type.2098 = f32[1,41,14336]{2,1,0} convert(%dot_general.1300), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/convert_element_type" stack_frame_id=772} + %jit_silu_.117 = f32[1,41,14336]{2,1,0} call(%convert_element_type.2098), to_apply=%silu.9, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/jit(silu)" stack_frame_id=775} + %convert_element_type.2099 = bf16[1,41,14336]{2,1,0} convert(%jit_silu_.117), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/convert_element_type" stack_frame_id=779} + %mul.3902 = bf16[1,41,14336]{2,1,0} multiply(%dot_general.1301, %convert_element_type.2099), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/mul" stack_frame_id=791} + %state_1__207_.1 = bf16[14336,4096]{1,0} parameter(210), metadata={op_name="state[1][207]"} + %dot_general.1302 = bf16[1,41,4096]{2,1,0} dot(%mul.3902, %state_1__207_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/feedforward_output_dense/dot_general" stack_frame_id=800} + %add.1283 = bf16[1,41,4096]{2,1,0} add(%dot_general.1302, %add.1281), metadata={op_name="jit(compiled_generate_function)/transformer_layer_22/add" stack_frame_id=801} + %convert_element_type.2102 = f32[1,41,4096]{2,1,0} convert(%add.1283), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention_layernorm/convert_element_type" stack_frame_id=555} + %pow.476 = f32[1,41,4096]{2,1,0} power(%convert_element_type.2102, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention_layernorm/pow" stack_frame_id=558} + %reduce_sum.1518 = f32[1,41]{1,0} reduce(%pow.476, %constant.155), dimensions={2}, to_apply=%region_93.98, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention_layernorm/reduce_sum" stack_frame_id=24} + %broadcast_in_dim.1911 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1518), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention_layernorm/broadcast_in_dim" stack_frame_id=24} + %div.1199 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1911, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention_layernorm/div" stack_frame_id=24} + %add.1285 = f32[1,41,1]{2,1,0} add(%div.1199, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention_layernorm/add" stack_frame_id=559} + %rsqrt.239 = f32[1,41,1]{2,1,0} rsqrt(%add.1285), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention_layernorm/rsqrt" stack_frame_id=562} + %mul.3903 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.239), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3904 = f32[1,41]{1,0} reshape(%mul.3903), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3905 = f32[1,41,4096]{2,1,0} broadcast(%mul.3904), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3906 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.2102, %mul.3905), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention_layernorm/mul" stack_frame_id=563} + %state_1__213_.1 = bf16[4096]{0} parameter(216), metadata={op_name="state[1][213]"} + %convert_element_type.2103 = f32[4096]{0} convert(%state_1__213_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention_layernorm/convert_element_type" stack_frame_id=564} + %broadcast_in_dim.1912 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.2103), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention_layernorm/broadcast_in_dim" stack_frame_id=564} + %mul.3907 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1912), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3908 = f32[1,4096]{1,0} reshape(%mul.3907), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3909 = f32[1,41,4096]{2,1,0} broadcast(%mul.3908), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3910 = f32[1,41,4096]{2,1,0} multiply(%mul.3906, %mul.3909), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention_layernorm/mul" stack_frame_id=564} + %convert_element_type.2104 = bf16[1,41,4096]{2,1,0} convert(%mul.3910), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention_layernorm/convert_element_type" stack_frame_id=568} + %state_1__210_.1 = bf16[4096,8,128]{2,1,0} parameter(213), metadata={op_name="state[1][210]"} + %dot_general.1305 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.2104, %state_1__210_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=647} + %iota.701 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/iota" stack_frame_id=677} + %broadcast_in_dim.1914 = f32[1,41]{1,0} reshape(%iota.701), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/broadcast_in_dim" stack_frame_id=680} + %iota.700 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/iota" stack_frame_id=667} + %mul.3920 = f32[64]{0} multiply(%iota.700, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=667} + %div.1202 = f32[64]{0} divide(%mul.3920, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/div" stack_frame_id=668} + %neg.484 = f32[64]{0} negate(%div.1202), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/neg" stack_frame_id=669} + %pow.478 = f32[64]{0} power(%broadcast.51, %neg.484), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/pow" stack_frame_id=672} + %div.1203 = f32[64]{0} divide(%pow.478, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/div" stack_frame_id=673} + %dot_general.1307 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1914, %div.1203), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/bi,j->bij/dot_general" stack_frame_id=683} + %stack.1884 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1307), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/stack" stack_frame_id=686} + %stack.1885 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1307), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/stack" stack_frame_id=686} + %stack.1886 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1884, %stack.1885), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/stack" stack_frame_id=686} + %reshape.884 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1886), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/reshape"} + %cos.238 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.884), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/cos" stack_frame_id=690} + %convert_element_type.2107 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.238), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/convert_element_type" stack_frame_id=694} + %mul.3921 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2107), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=714} + %mul.3922 = bf16[1,41,128]{2,1,0} reshape(%mul.3921), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=714} + %mul.3923 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3922), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=714} + %mul.3924 = bf16[1,41,8,128]{3,2,1,0} multiply(%dot_general.1305, %mul.3923), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=714} + %split.477 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1305), slice={[0:1], [0:41], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/split" stack_frame_id=706} + %neg.485 = bf16[1,41,8,64]{3,2,1,0} negate(%split.477), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/neg" stack_frame_id=707} + %stack.1887 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%neg.485), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/stack" stack_frame_id=710} + %split.476 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1305), slice={[0:1], [0:41], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/split" stack_frame_id=706} + %stack.1888 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%split.476), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/stack" stack_frame_id=710} + %stack.1889 = bf16[1,41,8,2,64]{4,3,2,1,0} concatenate(%stack.1887, %stack.1888), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/stack" stack_frame_id=710} + %reshape.885 = bf16[1,41,8,128]{3,2,1,0} reshape(%stack.1889), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/reshape" stack_frame_id=713} + %sin.238 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.884), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/sin" stack_frame_id=698} + %convert_element_type.2108 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.238), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/convert_element_type" stack_frame_id=702} + %mul.3925 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2108), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=715} + %mul.3926 = bf16[1,41,128]{2,1,0} reshape(%mul.3925), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=715} + %mul.3927 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3926), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=715} + %mul.3928 = bf16[1,41,8,128]{3,2,1,0} multiply(%reshape.885, %mul.3927), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=715} + %add.1287 = bf16[1,41,8,128]{3,2,1,0} add(%mul.3924, %mul.3928), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/add" stack_frame_id=716} + %stack.1890 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%add.1287), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/stack" stack_frame_id=719} + %state_1__211_.1 = bf16[4096,8,128]{2,1,0} parameter(214), metadata={op_name="state[1][211]"} + %dot_general.1306 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.2104, %state_1__211_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=656} + %stack.1891 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dot_general.1306), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/stack" stack_frame_id=719} + %stack.1892 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1890, %stack.1891), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/stack" stack_frame_id=719} + %stack.2030 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1892), metadata={op_name="jit(compiled_generate_function)/stack" stack_frame_id=804} + %broadcast_in_dim.1916 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dot_general.1306), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/broadcast_in_dim" stack_frame_id=725} + %reshape.887 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1916), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/reshape" stack_frame_id=725} + %iota.694 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/iota" stack_frame_id=525} + %broadcast_in_dim.1905 = f32[41,1]{1,0} reshape(%iota.694), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/broadcast_in_dim" stack_frame_id=528} + %ge.642 = f32[41,1]{1,0} broadcast(%broadcast_in_dim.1905), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/ge" stack_frame_id=532} + %ge.643 = f32[41]{0} reshape(%ge.642), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/ge" stack_frame_id=532} + %ge.644 = f32[41,41]{1,0} broadcast(%ge.643), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/ge" stack_frame_id=532} + %iota.695 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/iota" stack_frame_id=531} + %broadcast_in_dim.1906 = f32[1,41]{1,0} reshape(%iota.695), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/broadcast_in_dim" stack_frame_id=532} + %ge.645 = f32[1,41]{1,0} broadcast(%broadcast_in_dim.1906), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/ge" stack_frame_id=532} + %ge.646 = f32[41]{0} reshape(%ge.645), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/ge" stack_frame_id=532} + %ge.647 = f32[41,41]{1,0} broadcast(%ge.646), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/ge" stack_frame_id=532} + %ge.648 = pred[41,41]{1,0} compare(%ge.644, %ge.647), direction=GE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/ge" stack_frame_id=532} + %broadcast_in_dim.1907 = pred[1,41,41]{2,1,0} reshape(%ge.648), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/broadcast_in_dim" stack_frame_id=535} + %convert_element_type.2101 = s32[1,41,41]{2,1,0} convert(%broadcast_in_dim.1907), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/convert_element_type" stack_frame_id=551} + %iota.696 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/iota" stack_frame_id=538} + %broadcast_in_dim.1908 = s32[41,1]{1,0} reshape(%iota.696), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/broadcast_in_dim" stack_frame_id=536} + %lt.784 = s32[41,1]{1,0} broadcast(%broadcast_in_dim.1908), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/lt" stack_frame_id=543} + %lt.785 = s32[41]{0} reshape(%lt.784), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/lt" stack_frame_id=543} + %lt.786 = s32[41,41]{1,0} broadcast(%lt.785), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/lt" stack_frame_id=543} + %iota.697 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/iota" stack_frame_id=541} + %broadcast_in_dim.1909 = s32[1,41]{1,0} reshape(%iota.697), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/broadcast_in_dim" stack_frame_id=539} + %add.1284 = s32[1,41]{1,0} add(%broadcast_in_dim.1909, %broadcast.57), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/add" stack_frame_id=542} + %lt.787 = s32[1,41]{1,0} broadcast(%add.1284), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/lt" stack_frame_id=543} + %lt.788 = s32[41]{0} reshape(%lt.787), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/lt" stack_frame_id=543} + %lt.789 = s32[41,41]{1,0} broadcast(%lt.788), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/lt" stack_frame_id=543} + %lt.790 = pred[41,41]{1,0} compare(%lt.786, %lt.789), direction=LT, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/lt" stack_frame_id=543} + %convert_element_type.2100 = s32[41,41]{1,0} convert(%lt.790), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/convert_element_type" stack_frame_id=547} + %broadcast_in_dim.1910 = s32[1,41,41]{2,1,0} reshape(%convert_element_type.2100), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/broadcast_in_dim" stack_frame_id=551} + %min.118 = s32[1,41,41]{2,1,0} minimum(%convert_element_type.2101, %broadcast_in_dim.1910), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/min" stack_frame_id=551} + %broadcast_in_dim.1917 = s32[1,1,41,41]{3,2,1,0} reshape(%min.118), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/broadcast_in_dim" stack_frame_id=728} + %convert_element_type.2109 = pred[1,1,41,41]{3,2,1,0} compare(%broadcast_in_dim.1917, %broadcast.49), direction=NE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/convert_element_type" stack_frame_id=732} + %broadcast_in_dim.1918 = pred[1,1,1,41,41]{4,3,2,1,0} reshape(%convert_element_type.2109), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/broadcast_in_dim" stack_frame_id=34} + %and.367 = pred[1,1,1,41,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1918), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/vmap()/and" stack_frame_id=34} + %and.368 = pred[1,1,41,41]{3,2,1,0} reshape(%and.367), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/vmap()/and" stack_frame_id=34} + %and.369 = pred[1,1,32,41,41]{4,3,2,1,0} broadcast(%and.368), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/vmap()/and" stack_frame_id=34} + %state_1__209_.1 = bf16[4096,32,128]{2,1,0} parameter(212), metadata={op_name="state[1][209]"} + %dot_general.1303 = bf16[1,41,32,128]{3,2,1,0} dot(%convert_element_type.2104, %state_1__209_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=577} + %iota.699 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/iota" stack_frame_id=598} + %broadcast_in_dim.1913 = f32[1,41]{1,0} reshape(%iota.699), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/broadcast_in_dim" stack_frame_id=601} + %iota.698 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/iota" stack_frame_id=588} + %mul.3911 = f32[64]{0} multiply(%iota.698, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=588} + %div.1200 = f32[64]{0} divide(%mul.3911, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/div" stack_frame_id=589} + %neg.482 = f32[64]{0} negate(%div.1200), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/neg" stack_frame_id=590} + %pow.477 = f32[64]{0} power(%broadcast.51, %neg.482), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/pow" stack_frame_id=593} + %div.1201 = f32[64]{0} divide(%pow.477, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/div" stack_frame_id=594} + %dot_general.1304 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1913, %div.1201), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/bi,j->bij/dot_general" stack_frame_id=604} + %stack.1878 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1304), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/stack" stack_frame_id=607} + %stack.1879 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1304), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/stack" stack_frame_id=607} + %stack.1880 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1878, %stack.1879), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/stack" stack_frame_id=607} + %reshape.882 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1880), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/reshape"} + %cos.237 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.882), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/cos" stack_frame_id=611} + %convert_element_type.2105 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.237), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/convert_element_type" stack_frame_id=615} + %mul.3912 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2105), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=635} + %mul.3913 = bf16[1,41,128]{2,1,0} reshape(%mul.3912), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=635} + %mul.3914 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3913), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=635} + %mul.3915 = bf16[1,41,32,128]{3,2,1,0} multiply(%dot_general.1303, %mul.3914), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=635} + %split.475 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1303), slice={[0:1], [0:41], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/split" stack_frame_id=627} + %neg.483 = bf16[1,41,32,64]{3,2,1,0} negate(%split.475), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/neg" stack_frame_id=628} + %stack.1881 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%neg.483), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/stack" stack_frame_id=631} + %split.474 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1303), slice={[0:1], [0:41], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/split" stack_frame_id=627} + %stack.1882 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%split.474), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/stack" stack_frame_id=631} + %stack.1883 = bf16[1,41,32,2,64]{4,3,2,1,0} concatenate(%stack.1881, %stack.1882), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/stack" stack_frame_id=631} + %reshape.883 = bf16[1,41,32,128]{3,2,1,0} reshape(%stack.1883), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/reshape" stack_frame_id=634} + %sin.237 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.882), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/sin" stack_frame_id=619} + %convert_element_type.2106 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.237), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/convert_element_type" stack_frame_id=623} + %mul.3916 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2106), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=636} + %mul.3917 = bf16[1,41,128]{2,1,0} reshape(%mul.3916), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=636} + %mul.3918 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3917), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=636} + %mul.3919 = bf16[1,41,32,128]{3,2,1,0} multiply(%reshape.883, %mul.3918), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/mul" stack_frame_id=636} + %add.1286 = bf16[1,41,32,128]{3,2,1,0} add(%mul.3915, %mul.3919), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/rotary_embedding_23/add" stack_frame_id=637} + %reshape.888 = bf16[1,41,32,1,128]{4,3,2,1,0} reshape(%add.1286), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/reshape" stack_frame_id=34} + %broadcast_in_dim.1915 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%add.1287), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/broadcast_in_dim" stack_frame_id=722} + %reshape.886 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1915), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/reshape" stack_frame_id=722} + %dot_general.1308 = f32[1,32,41,1,41]{4,3,2,1,0} dot(%reshape.888, %reshape.886), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=34} + %mul.3929 = f32[1,32,41,1,41]{4,3,2,1,0} multiply(%dot_general.1308, %broadcast.48), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/vmap()/mul" stack_frame_id=34} + %vmap_jit__where__.118 = f32[1,1,32,41,41]{4,3,2,1,0} call(%and.369, %mul.3929, %constant.153), to_apply=%_where_0.5, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/vmap(jit(_where))" stack_frame_id=34} + %reduce_max.504 = f32[1,1,32,41]{3,2,1,0} reduce(%vmap_jit__where__.118, %constant.152), dimensions={4}, to_apply=%region_94.99, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/vmap()/reduce_max" stack_frame_id=34} + %max.128 = f32[1,1,32,41]{3,2,1,0} maximum(%reduce_max.504, %broadcast.47), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/vmap()/max" stack_frame_id=34} + %broadcast_in_dim.1919 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%max.128), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %sub.496 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1919), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/vmap()/sub" stack_frame_id=34} + %sub.497 = f32[1,1,32,41]{3,2,1,0} reshape(%sub.496), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/vmap()/sub" stack_frame_id=34} + %sub.498 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%sub.497), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/vmap()/sub" stack_frame_id=34} + %sub.499 = f32[1,1,32,41,41]{4,3,2,1,0} subtract(%vmap_jit__where__.118, %sub.498), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/vmap()/sub" stack_frame_id=34} + %exp.124 = f32[1,1,32,41,41]{4,3,2,1,0} exponential(%sub.499), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/vmap()/exp" stack_frame_id=34} + %reduce_sum.1519 = f32[1,1,32,41]{3,2,1,0} reduce(%exp.124, %constant.155), dimensions={4}, to_apply=%region_95.100, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/vmap()/reduce_sum" stack_frame_id=34} + %broadcast_in_dim.1920 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%reduce_sum.1519), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %div.1204 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1920), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/vmap()/div" stack_frame_id=34} + %div.1205 = f32[1,1,32,41]{3,2,1,0} reshape(%div.1204), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/vmap()/div" stack_frame_id=34} + %div.1206 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%div.1205), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/vmap()/div" stack_frame_id=34} + %div.1207 = f32[1,1,32,41,41]{4,3,2,1,0} divide(%exp.124, %div.1206), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/vmap()/div" stack_frame_id=34} + %convert_element_type.2110 = bf16[1,1,32,41,41]{4,3,2,1,0} convert(%div.1207), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/vmap()/convert_element_type" stack_frame_id=34} + %dot_general.1309 = bf16[1,32,128,1,41]{4,3,2,1,0} dot(%reshape.887, %convert_element_type.2110), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=34} + %transpose.122 = bf16[1,41,32,1,128]{1,3,4,2,0} transpose(%dot_general.1309), dimensions={0,4,1,3,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/vmap()/transpose;jit(compiled_generate_function)/transformer_layer_23/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %reshape.889 = bf16[1,41,32,128]{3,2,1,0} reshape(%transpose.122), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/reshape" stack_frame_id=34} + %state_1__212_.1 = bf16[32,128,4096]{2,1,0} parameter(215), metadata={op_name="state[1][212]"} + %dot_general.1310 = bf16[1,41,4096]{2,1,0} dot(%reshape.889, %state_1__212_.1), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=741} + %add.1288 = bf16[1,41,4096]{2,1,0} add(%dot_general.1310, %add.1283), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/add" stack_frame_id=742} + %convert_element_type.2111 = f32[1,41,4096]{2,1,0} convert(%add.1288), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/feedforward_layernorm/convert_element_type" stack_frame_id=746} + %pow.479 = f32[1,41,4096]{2,1,0} power(%convert_element_type.2111, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/feedforward_layernorm/pow" stack_frame_id=749} + %reduce_sum.1520 = f32[1,41]{1,0} reduce(%pow.479, %constant.155), dimensions={2}, to_apply=%region_96.101, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/feedforward_layernorm/reduce_sum" stack_frame_id=43} + %broadcast_in_dim.1921 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1520), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/feedforward_layernorm/broadcast_in_dim" stack_frame_id=43} + %div.1208 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1921, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/feedforward_layernorm/div" stack_frame_id=43} + %add.1289 = f32[1,41,1]{2,1,0} add(%div.1208, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/feedforward_layernorm/add" stack_frame_id=750} + %rsqrt.240 = f32[1,41,1]{2,1,0} rsqrt(%add.1289), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/feedforward_layernorm/rsqrt" stack_frame_id=753} + %mul.3930 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.240), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3931 = f32[1,41]{1,0} reshape(%mul.3930), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3932 = f32[1,41,4096]{2,1,0} broadcast(%mul.3931), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3933 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.2111, %mul.3932), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/feedforward_layernorm/mul" stack_frame_id=754} + %state_1__217_.1 = bf16[4096]{0} parameter(220), metadata={op_name="state[1][217]"} + %convert_element_type.2112 = f32[4096]{0} convert(%state_1__217_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/feedforward_layernorm/convert_element_type" stack_frame_id=755} + %broadcast_in_dim.1922 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.2112), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/feedforward_layernorm/broadcast_in_dim" stack_frame_id=755} + %mul.3934 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1922), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3935 = f32[1,4096]{1,0} reshape(%mul.3934), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3936 = f32[1,41,4096]{2,1,0} broadcast(%mul.3935), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3937 = f32[1,41,4096]{2,1,0} multiply(%mul.3933, %mul.3936), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/feedforward_layernorm/mul" stack_frame_id=755} + %convert_element_type.2113 = bf16[1,41,4096]{2,1,0} convert(%mul.3937), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/feedforward_layernorm/convert_element_type" stack_frame_id=759} + %state_1__214_.1 = bf16[4096,14336]{1,0} parameter(217), metadata={op_name="state[1][214]"} + %dot_general.1312 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.2113, %state_1__214_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/feedforward_intermediate_dense/dot_general" stack_frame_id=788} + %state_1__215_.1 = bf16[4096,14336]{1,0} parameter(218), metadata={op_name="state[1][215]"} + %dot_general.1311 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.2113, %state_1__215_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/feedforward_gate_dense/dot_general" stack_frame_id=768} + %convert_element_type.2114 = f32[1,41,14336]{2,1,0} convert(%dot_general.1311), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/convert_element_type" stack_frame_id=772} + %jit_silu_.118 = f32[1,41,14336]{2,1,0} call(%convert_element_type.2114), to_apply=%silu.9, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/jit(silu)" stack_frame_id=775} + %convert_element_type.2115 = bf16[1,41,14336]{2,1,0} convert(%jit_silu_.118), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/convert_element_type" stack_frame_id=779} + %mul.3938 = bf16[1,41,14336]{2,1,0} multiply(%dot_general.1312, %convert_element_type.2115), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/mul" stack_frame_id=791} + %state_1__216_.1 = bf16[14336,4096]{1,0} parameter(219), metadata={op_name="state[1][216]"} + %dot_general.1313 = bf16[1,41,4096]{2,1,0} dot(%mul.3938, %state_1__216_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/feedforward_output_dense/dot_general" stack_frame_id=800} + %add.1290 = bf16[1,41,4096]{2,1,0} add(%dot_general.1313, %add.1288), metadata={op_name="jit(compiled_generate_function)/transformer_layer_23/add" stack_frame_id=801} + %convert_element_type.2118 = f32[1,41,4096]{2,1,0} convert(%add.1290), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention_layernorm/convert_element_type" stack_frame_id=555} + %pow.480 = f32[1,41,4096]{2,1,0} power(%convert_element_type.2118, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention_layernorm/pow" stack_frame_id=558} + %reduce_sum.1521 = f32[1,41]{1,0} reduce(%pow.480, %constant.155), dimensions={2}, to_apply=%region_97.102, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention_layernorm/reduce_sum" stack_frame_id=24} + %broadcast_in_dim.1929 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1521), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention_layernorm/broadcast_in_dim" stack_frame_id=24} + %div.1209 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1929, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention_layernorm/div" stack_frame_id=24} + %add.1292 = f32[1,41,1]{2,1,0} add(%div.1209, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention_layernorm/add" stack_frame_id=559} + %rsqrt.241 = f32[1,41,1]{2,1,0} rsqrt(%add.1292), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention_layernorm/rsqrt" stack_frame_id=562} + %mul.3939 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.241), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3940 = f32[1,41]{1,0} reshape(%mul.3939), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3941 = f32[1,41,4096]{2,1,0} broadcast(%mul.3940), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3942 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.2118, %mul.3941), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention_layernorm/mul" stack_frame_id=563} + %state_1__222_.1 = bf16[4096]{0} parameter(225), metadata={op_name="state[1][222]"} + %convert_element_type.2119 = f32[4096]{0} convert(%state_1__222_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention_layernorm/convert_element_type" stack_frame_id=564} + %broadcast_in_dim.1930 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.2119), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention_layernorm/broadcast_in_dim" stack_frame_id=564} + %mul.3943 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1930), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3944 = f32[1,4096]{1,0} reshape(%mul.3943), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3945 = f32[1,41,4096]{2,1,0} broadcast(%mul.3944), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3946 = f32[1,41,4096]{2,1,0} multiply(%mul.3942, %mul.3945), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention_layernorm/mul" stack_frame_id=564} + %convert_element_type.2120 = bf16[1,41,4096]{2,1,0} convert(%mul.3946), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention_layernorm/convert_element_type" stack_frame_id=568} + %state_1__219_.1 = bf16[4096,8,128]{2,1,0} parameter(222), metadata={op_name="state[1][219]"} + %dot_general.1316 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.2120, %state_1__219_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=647} + %iota.709 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/iota" stack_frame_id=677} + %broadcast_in_dim.1932 = f32[1,41]{1,0} reshape(%iota.709), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/broadcast_in_dim" stack_frame_id=680} + %iota.708 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/iota" stack_frame_id=667} + %mul.3956 = f32[64]{0} multiply(%iota.708, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=667} + %div.1212 = f32[64]{0} divide(%mul.3956, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/div" stack_frame_id=668} + %neg.488 = f32[64]{0} negate(%div.1212), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/neg" stack_frame_id=669} + %pow.482 = f32[64]{0} power(%broadcast.51, %neg.488), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/pow" stack_frame_id=672} + %div.1213 = f32[64]{0} divide(%pow.482, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/div" stack_frame_id=673} + %dot_general.1318 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1932, %div.1213), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/bi,j->bij/dot_general" stack_frame_id=683} + %stack.1899 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1318), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/stack" stack_frame_id=686} + %stack.1900 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1318), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/stack" stack_frame_id=686} + %stack.1901 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1899, %stack.1900), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/stack" stack_frame_id=686} + %reshape.892 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1901), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/reshape"} + %cos.240 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.892), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/cos" stack_frame_id=690} + %convert_element_type.2123 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.240), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/convert_element_type" stack_frame_id=694} + %mul.3957 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2123), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=714} + %mul.3958 = bf16[1,41,128]{2,1,0} reshape(%mul.3957), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=714} + %mul.3959 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3958), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=714} + %mul.3960 = bf16[1,41,8,128]{3,2,1,0} multiply(%dot_general.1316, %mul.3959), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=714} + %split.481 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1316), slice={[0:1], [0:41], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/split" stack_frame_id=706} + %neg.489 = bf16[1,41,8,64]{3,2,1,0} negate(%split.481), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/neg" stack_frame_id=707} + %stack.1902 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%neg.489), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/stack" stack_frame_id=710} + %split.480 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1316), slice={[0:1], [0:41], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/split" stack_frame_id=706} + %stack.1903 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%split.480), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/stack" stack_frame_id=710} + %stack.1904 = bf16[1,41,8,2,64]{4,3,2,1,0} concatenate(%stack.1902, %stack.1903), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/stack" stack_frame_id=710} + %reshape.893 = bf16[1,41,8,128]{3,2,1,0} reshape(%stack.1904), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/reshape" stack_frame_id=713} + %sin.240 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.892), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/sin" stack_frame_id=698} + %convert_element_type.2124 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.240), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/convert_element_type" stack_frame_id=702} + %mul.3961 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2124), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=715} + %mul.3962 = bf16[1,41,128]{2,1,0} reshape(%mul.3961), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=715} + %mul.3963 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3962), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=715} + %mul.3964 = bf16[1,41,8,128]{3,2,1,0} multiply(%reshape.893, %mul.3963), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=715} + %add.1294 = bf16[1,41,8,128]{3,2,1,0} add(%mul.3960, %mul.3964), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/add" stack_frame_id=716} + %stack.1905 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%add.1294), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/stack" stack_frame_id=719} + %state_1__220_.1 = bf16[4096,8,128]{2,1,0} parameter(223), metadata={op_name="state[1][220]"} + %dot_general.1317 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.2120, %state_1__220_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=656} + %stack.1906 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dot_general.1317), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/stack" stack_frame_id=719} + %stack.1907 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1905, %stack.1906), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/stack" stack_frame_id=719} + %stack.2031 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1907), metadata={op_name="jit(compiled_generate_function)/stack" stack_frame_id=804} + %broadcast_in_dim.1934 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dot_general.1317), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/broadcast_in_dim" stack_frame_id=725} + %reshape.895 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1934), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/reshape" stack_frame_id=725} + %iota.702 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/iota" stack_frame_id=525} + %broadcast_in_dim.1923 = f32[41,1]{1,0} reshape(%iota.702), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/broadcast_in_dim" stack_frame_id=528} + %ge.649 = f32[41,1]{1,0} broadcast(%broadcast_in_dim.1923), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/ge" stack_frame_id=532} + %ge.650 = f32[41]{0} reshape(%ge.649), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/ge" stack_frame_id=532} + %ge.651 = f32[41,41]{1,0} broadcast(%ge.650), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/ge" stack_frame_id=532} + %iota.703 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/iota" stack_frame_id=531} + %broadcast_in_dim.1924 = f32[1,41]{1,0} reshape(%iota.703), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/broadcast_in_dim" stack_frame_id=532} + %ge.652 = f32[1,41]{1,0} broadcast(%broadcast_in_dim.1924), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/ge" stack_frame_id=532} + %ge.653 = f32[41]{0} reshape(%ge.652), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/ge" stack_frame_id=532} + %ge.654 = f32[41,41]{1,0} broadcast(%ge.653), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/ge" stack_frame_id=532} + %ge.655 = pred[41,41]{1,0} compare(%ge.651, %ge.654), direction=GE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/ge" stack_frame_id=532} + %broadcast_in_dim.1925 = pred[1,41,41]{2,1,0} reshape(%ge.655), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/broadcast_in_dim" stack_frame_id=535} + %convert_element_type.2117 = s32[1,41,41]{2,1,0} convert(%broadcast_in_dim.1925), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/convert_element_type" stack_frame_id=551} + %iota.704 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/iota" stack_frame_id=538} + %broadcast_in_dim.1926 = s32[41,1]{1,0} reshape(%iota.704), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/broadcast_in_dim" stack_frame_id=536} + %lt.791 = s32[41,1]{1,0} broadcast(%broadcast_in_dim.1926), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/lt" stack_frame_id=543} + %lt.792 = s32[41]{0} reshape(%lt.791), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/lt" stack_frame_id=543} + %lt.793 = s32[41,41]{1,0} broadcast(%lt.792), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/lt" stack_frame_id=543} + %iota.705 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/iota" stack_frame_id=541} + %broadcast_in_dim.1927 = s32[1,41]{1,0} reshape(%iota.705), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/broadcast_in_dim" stack_frame_id=539} + %add.1291 = s32[1,41]{1,0} add(%broadcast_in_dim.1927, %broadcast.57), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/add" stack_frame_id=542} + %lt.794 = s32[1,41]{1,0} broadcast(%add.1291), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/lt" stack_frame_id=543} + %lt.795 = s32[41]{0} reshape(%lt.794), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/lt" stack_frame_id=543} + %lt.796 = s32[41,41]{1,0} broadcast(%lt.795), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/lt" stack_frame_id=543} + %lt.797 = pred[41,41]{1,0} compare(%lt.793, %lt.796), direction=LT, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/lt" stack_frame_id=543} + %convert_element_type.2116 = s32[41,41]{1,0} convert(%lt.797), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/convert_element_type" stack_frame_id=547} + %broadcast_in_dim.1928 = s32[1,41,41]{2,1,0} reshape(%convert_element_type.2116), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/broadcast_in_dim" stack_frame_id=551} + %min.119 = s32[1,41,41]{2,1,0} minimum(%convert_element_type.2117, %broadcast_in_dim.1928), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/min" stack_frame_id=551} + %broadcast_in_dim.1935 = s32[1,1,41,41]{3,2,1,0} reshape(%min.119), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/broadcast_in_dim" stack_frame_id=728} + %convert_element_type.2125 = pred[1,1,41,41]{3,2,1,0} compare(%broadcast_in_dim.1935, %broadcast.49), direction=NE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/convert_element_type" stack_frame_id=732} + %broadcast_in_dim.1936 = pred[1,1,1,41,41]{4,3,2,1,0} reshape(%convert_element_type.2125), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/broadcast_in_dim" stack_frame_id=34} + %and.370 = pred[1,1,1,41,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1936), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/vmap()/and" stack_frame_id=34} + %and.371 = pred[1,1,41,41]{3,2,1,0} reshape(%and.370), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/vmap()/and" stack_frame_id=34} + %and.372 = pred[1,1,32,41,41]{4,3,2,1,0} broadcast(%and.371), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/vmap()/and" stack_frame_id=34} + %state_1__218_.1 = bf16[4096,32,128]{2,1,0} parameter(221), metadata={op_name="state[1][218]"} + %dot_general.1314 = bf16[1,41,32,128]{3,2,1,0} dot(%convert_element_type.2120, %state_1__218_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=577} + %iota.707 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/iota" stack_frame_id=598} + %broadcast_in_dim.1931 = f32[1,41]{1,0} reshape(%iota.707), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/broadcast_in_dim" stack_frame_id=601} + %iota.706 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/iota" stack_frame_id=588} + %mul.3947 = f32[64]{0} multiply(%iota.706, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=588} + %div.1210 = f32[64]{0} divide(%mul.3947, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/div" stack_frame_id=589} + %neg.486 = f32[64]{0} negate(%div.1210), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/neg" stack_frame_id=590} + %pow.481 = f32[64]{0} power(%broadcast.51, %neg.486), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/pow" stack_frame_id=593} + %div.1211 = f32[64]{0} divide(%pow.481, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/div" stack_frame_id=594} + %dot_general.1315 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1931, %div.1211), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/bi,j->bij/dot_general" stack_frame_id=604} + %stack.1893 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1315), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/stack" stack_frame_id=607} + %stack.1894 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1315), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/stack" stack_frame_id=607} + %stack.1895 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1893, %stack.1894), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/stack" stack_frame_id=607} + %reshape.890 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1895), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/reshape"} + %cos.239 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.890), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/cos" stack_frame_id=611} + %convert_element_type.2121 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.239), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/convert_element_type" stack_frame_id=615} + %mul.3948 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2121), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=635} + %mul.3949 = bf16[1,41,128]{2,1,0} reshape(%mul.3948), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=635} + %mul.3950 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3949), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=635} + %mul.3951 = bf16[1,41,32,128]{3,2,1,0} multiply(%dot_general.1314, %mul.3950), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=635} + %split.479 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1314), slice={[0:1], [0:41], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/split" stack_frame_id=627} + %neg.487 = bf16[1,41,32,64]{3,2,1,0} negate(%split.479), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/neg" stack_frame_id=628} + %stack.1896 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%neg.487), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/stack" stack_frame_id=631} + %split.478 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1314), slice={[0:1], [0:41], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/split" stack_frame_id=627} + %stack.1897 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%split.478), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/stack" stack_frame_id=631} + %stack.1898 = bf16[1,41,32,2,64]{4,3,2,1,0} concatenate(%stack.1896, %stack.1897), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/stack" stack_frame_id=631} + %reshape.891 = bf16[1,41,32,128]{3,2,1,0} reshape(%stack.1898), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/reshape" stack_frame_id=634} + %sin.239 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.890), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/sin" stack_frame_id=619} + %convert_element_type.2122 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.239), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/convert_element_type" stack_frame_id=623} + %mul.3952 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2122), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=636} + %mul.3953 = bf16[1,41,128]{2,1,0} reshape(%mul.3952), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=636} + %mul.3954 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3953), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=636} + %mul.3955 = bf16[1,41,32,128]{3,2,1,0} multiply(%reshape.891, %mul.3954), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/mul" stack_frame_id=636} + %add.1293 = bf16[1,41,32,128]{3,2,1,0} add(%mul.3951, %mul.3955), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/rotary_embedding_24/add" stack_frame_id=637} + %reshape.896 = bf16[1,41,32,1,128]{4,3,2,1,0} reshape(%add.1293), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/reshape" stack_frame_id=34} + %broadcast_in_dim.1933 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%add.1294), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/broadcast_in_dim" stack_frame_id=722} + %reshape.894 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1933), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/reshape" stack_frame_id=722} + %dot_general.1319 = f32[1,32,41,1,41]{4,3,2,1,0} dot(%reshape.896, %reshape.894), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=34} + %mul.3965 = f32[1,32,41,1,41]{4,3,2,1,0} multiply(%dot_general.1319, %broadcast.48), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/vmap()/mul" stack_frame_id=34} + %vmap_jit__where__.119 = f32[1,1,32,41,41]{4,3,2,1,0} call(%and.372, %mul.3965, %constant.153), to_apply=%_where_0.5, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/vmap(jit(_where))" stack_frame_id=34} + %reduce_max.505 = f32[1,1,32,41]{3,2,1,0} reduce(%vmap_jit__where__.119, %constant.152), dimensions={4}, to_apply=%region_98.103, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/vmap()/reduce_max" stack_frame_id=34} + %max.129 = f32[1,1,32,41]{3,2,1,0} maximum(%reduce_max.505, %broadcast.47), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/vmap()/max" stack_frame_id=34} + %broadcast_in_dim.1937 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%max.129), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %sub.500 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1937), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/vmap()/sub" stack_frame_id=34} + %sub.501 = f32[1,1,32,41]{3,2,1,0} reshape(%sub.500), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/vmap()/sub" stack_frame_id=34} + %sub.502 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%sub.501), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/vmap()/sub" stack_frame_id=34} + %sub.503 = f32[1,1,32,41,41]{4,3,2,1,0} subtract(%vmap_jit__where__.119, %sub.502), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/vmap()/sub" stack_frame_id=34} + %exp.125 = f32[1,1,32,41,41]{4,3,2,1,0} exponential(%sub.503), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/vmap()/exp" stack_frame_id=34} + %reduce_sum.1522 = f32[1,1,32,41]{3,2,1,0} reduce(%exp.125, %constant.155), dimensions={4}, to_apply=%region_99.104, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/vmap()/reduce_sum" stack_frame_id=34} + %broadcast_in_dim.1938 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%reduce_sum.1522), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %div.1214 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1938), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/vmap()/div" stack_frame_id=34} + %div.1215 = f32[1,1,32,41]{3,2,1,0} reshape(%div.1214), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/vmap()/div" stack_frame_id=34} + %div.1216 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%div.1215), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/vmap()/div" stack_frame_id=34} + %div.1217 = f32[1,1,32,41,41]{4,3,2,1,0} divide(%exp.125, %div.1216), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/vmap()/div" stack_frame_id=34} + %convert_element_type.2126 = bf16[1,1,32,41,41]{4,3,2,1,0} convert(%div.1217), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/vmap()/convert_element_type" stack_frame_id=34} + %dot_general.1320 = bf16[1,32,128,1,41]{4,3,2,1,0} dot(%reshape.895, %convert_element_type.2126), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=34} + %transpose.123 = bf16[1,41,32,1,128]{1,3,4,2,0} transpose(%dot_general.1320), dimensions={0,4,1,3,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/vmap()/transpose;jit(compiled_generate_function)/transformer_layer_24/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %reshape.897 = bf16[1,41,32,128]{3,2,1,0} reshape(%transpose.123), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/reshape" stack_frame_id=34} + %state_1__221_.1 = bf16[32,128,4096]{2,1,0} parameter(224), metadata={op_name="state[1][221]"} + %dot_general.1321 = bf16[1,41,4096]{2,1,0} dot(%reshape.897, %state_1__221_.1), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=741} + %add.1295 = bf16[1,41,4096]{2,1,0} add(%dot_general.1321, %add.1290), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/add" stack_frame_id=742} + %convert_element_type.2127 = f32[1,41,4096]{2,1,0} convert(%add.1295), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/feedforward_layernorm/convert_element_type" stack_frame_id=746} + %pow.483 = f32[1,41,4096]{2,1,0} power(%convert_element_type.2127, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/feedforward_layernorm/pow" stack_frame_id=749} + %reduce_sum.1523 = f32[1,41]{1,0} reduce(%pow.483, %constant.155), dimensions={2}, to_apply=%region_100.105, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/feedforward_layernorm/reduce_sum" stack_frame_id=43} + %broadcast_in_dim.1939 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1523), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/feedforward_layernorm/broadcast_in_dim" stack_frame_id=43} + %div.1218 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1939, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/feedforward_layernorm/div" stack_frame_id=43} + %add.1296 = f32[1,41,1]{2,1,0} add(%div.1218, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/feedforward_layernorm/add" stack_frame_id=750} + %rsqrt.242 = f32[1,41,1]{2,1,0} rsqrt(%add.1296), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/feedforward_layernorm/rsqrt" stack_frame_id=753} + %mul.3966 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.242), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3967 = f32[1,41]{1,0} reshape(%mul.3966), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3968 = f32[1,41,4096]{2,1,0} broadcast(%mul.3967), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/feedforward_layernorm/mul" stack_frame_id=754} + %mul.3969 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.2127, %mul.3968), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/feedforward_layernorm/mul" stack_frame_id=754} + %state_1__226_.1 = bf16[4096]{0} parameter(229), metadata={op_name="state[1][226]"} + %convert_element_type.2128 = f32[4096]{0} convert(%state_1__226_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/feedforward_layernorm/convert_element_type" stack_frame_id=755} + %broadcast_in_dim.1940 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.2128), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/feedforward_layernorm/broadcast_in_dim" stack_frame_id=755} + %mul.3970 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1940), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3971 = f32[1,4096]{1,0} reshape(%mul.3970), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3972 = f32[1,41,4096]{2,1,0} broadcast(%mul.3971), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/feedforward_layernorm/mul" stack_frame_id=755} + %mul.3973 = f32[1,41,4096]{2,1,0} multiply(%mul.3969, %mul.3972), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/feedforward_layernorm/mul" stack_frame_id=755} + %convert_element_type.2129 = bf16[1,41,4096]{2,1,0} convert(%mul.3973), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/feedforward_layernorm/convert_element_type" stack_frame_id=759} + %state_1__223_.1 = bf16[4096,14336]{1,0} parameter(226), metadata={op_name="state[1][223]"} + %dot_general.1323 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.2129, %state_1__223_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/feedforward_intermediate_dense/dot_general" stack_frame_id=788} + %state_1__224_.1 = bf16[4096,14336]{1,0} parameter(227), metadata={op_name="state[1][224]"} + %dot_general.1322 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.2129, %state_1__224_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/feedforward_gate_dense/dot_general" stack_frame_id=768} + %convert_element_type.2130 = f32[1,41,14336]{2,1,0} convert(%dot_general.1322), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/convert_element_type" stack_frame_id=772} + %jit_silu_.119 = f32[1,41,14336]{2,1,0} call(%convert_element_type.2130), to_apply=%silu.9, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/jit(silu)" stack_frame_id=775} + %convert_element_type.2131 = bf16[1,41,14336]{2,1,0} convert(%jit_silu_.119), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/convert_element_type" stack_frame_id=779} + %mul.3974 = bf16[1,41,14336]{2,1,0} multiply(%dot_general.1323, %convert_element_type.2131), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/mul" stack_frame_id=791} + %state_1__225_.1 = bf16[14336,4096]{1,0} parameter(228), metadata={op_name="state[1][225]"} + %dot_general.1324 = bf16[1,41,4096]{2,1,0} dot(%mul.3974, %state_1__225_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/feedforward_output_dense/dot_general" stack_frame_id=800} + %add.1297 = bf16[1,41,4096]{2,1,0} add(%dot_general.1324, %add.1295), metadata={op_name="jit(compiled_generate_function)/transformer_layer_24/add" stack_frame_id=801} + %convert_element_type.2134 = f32[1,41,4096]{2,1,0} convert(%add.1297), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention_layernorm/convert_element_type" stack_frame_id=555} + %pow.484 = f32[1,41,4096]{2,1,0} power(%convert_element_type.2134, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention_layernorm/pow" stack_frame_id=558} + %reduce_sum.1524 = f32[1,41]{1,0} reduce(%pow.484, %constant.155), dimensions={2}, to_apply=%region_101.106, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention_layernorm/reduce_sum" stack_frame_id=24} + %broadcast_in_dim.1947 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1524), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention_layernorm/broadcast_in_dim" stack_frame_id=24} + %div.1219 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1947, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention_layernorm/div" stack_frame_id=24} + %add.1299 = f32[1,41,1]{2,1,0} add(%div.1219, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention_layernorm/add" stack_frame_id=559} + %rsqrt.243 = f32[1,41,1]{2,1,0} rsqrt(%add.1299), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention_layernorm/rsqrt" stack_frame_id=562} + %mul.3975 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.243), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3976 = f32[1,41]{1,0} reshape(%mul.3975), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3977 = f32[1,41,4096]{2,1,0} broadcast(%mul.3976), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention_layernorm/mul" stack_frame_id=563} + %mul.3978 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.2134, %mul.3977), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention_layernorm/mul" stack_frame_id=563} + %state_1__231_.1 = bf16[4096]{0} parameter(234), metadata={op_name="state[1][231]"} + %convert_element_type.2135 = f32[4096]{0} convert(%state_1__231_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention_layernorm/convert_element_type" stack_frame_id=564} + %broadcast_in_dim.1948 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.2135), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention_layernorm/broadcast_in_dim" stack_frame_id=564} + %mul.3979 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1948), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3980 = f32[1,4096]{1,0} reshape(%mul.3979), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3981 = f32[1,41,4096]{2,1,0} broadcast(%mul.3980), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention_layernorm/mul" stack_frame_id=564} + %mul.3982 = f32[1,41,4096]{2,1,0} multiply(%mul.3978, %mul.3981), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention_layernorm/mul" stack_frame_id=564} + %convert_element_type.2136 = bf16[1,41,4096]{2,1,0} convert(%mul.3982), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention_layernorm/convert_element_type" stack_frame_id=568} + %state_1__228_.1 = bf16[4096,8,128]{2,1,0} parameter(231), metadata={op_name="state[1][228]"} + %dot_general.1327 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.2136, %state_1__228_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=647} + %iota.717 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/iota" stack_frame_id=677} + %broadcast_in_dim.1950 = f32[1,41]{1,0} reshape(%iota.717), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/broadcast_in_dim" stack_frame_id=680} + %iota.716 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/iota" stack_frame_id=667} + %mul.3992 = f32[64]{0} multiply(%iota.716, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=667} + %div.1222 = f32[64]{0} divide(%mul.3992, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/div" stack_frame_id=668} + %neg.492 = f32[64]{0} negate(%div.1222), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/neg" stack_frame_id=669} + %pow.486 = f32[64]{0} power(%broadcast.51, %neg.492), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/pow" stack_frame_id=672} + %div.1223 = f32[64]{0} divide(%pow.486, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/div" stack_frame_id=673} + %dot_general.1329 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1950, %div.1223), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/bi,j->bij/dot_general" stack_frame_id=683} + %stack.1914 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1329), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/stack" stack_frame_id=686} + %stack.1915 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1329), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/stack" stack_frame_id=686} + %stack.1916 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1914, %stack.1915), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/stack" stack_frame_id=686} + %reshape.900 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1916), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/reshape"} + %cos.242 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.900), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/cos" stack_frame_id=690} + %convert_element_type.2139 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.242), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/convert_element_type" stack_frame_id=694} + %mul.3993 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2139), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=714} + %mul.3994 = bf16[1,41,128]{2,1,0} reshape(%mul.3993), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=714} + %mul.3995 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3994), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=714} + %mul.3996 = bf16[1,41,8,128]{3,2,1,0} multiply(%dot_general.1327, %mul.3995), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=714} + %split.485 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1327), slice={[0:1], [0:41], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/split" stack_frame_id=706} + %neg.493 = bf16[1,41,8,64]{3,2,1,0} negate(%split.485), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/neg" stack_frame_id=707} + %stack.1917 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%neg.493), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/stack" stack_frame_id=710} + %split.484 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1327), slice={[0:1], [0:41], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/split" stack_frame_id=706} + %stack.1918 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%split.484), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/stack" stack_frame_id=710} + %stack.1919 = bf16[1,41,8,2,64]{4,3,2,1,0} concatenate(%stack.1917, %stack.1918), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/stack" stack_frame_id=710} + %reshape.901 = bf16[1,41,8,128]{3,2,1,0} reshape(%stack.1919), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/reshape" stack_frame_id=713} + %sin.242 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.900), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/sin" stack_frame_id=698} + %convert_element_type.2140 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.242), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/convert_element_type" stack_frame_id=702} + %mul.3997 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2140), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=715} + %mul.3998 = bf16[1,41,128]{2,1,0} reshape(%mul.3997), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=715} + %mul.3999 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.3998), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=715} + %mul.4000 = bf16[1,41,8,128]{3,2,1,0} multiply(%reshape.901, %mul.3999), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=715} + %add.1301 = bf16[1,41,8,128]{3,2,1,0} add(%mul.3996, %mul.4000), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/add" stack_frame_id=716} + %stack.1920 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%add.1301), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/stack" stack_frame_id=719} + %state_1__229_.1 = bf16[4096,8,128]{2,1,0} parameter(232), metadata={op_name="state[1][229]"} + %dot_general.1328 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.2136, %state_1__229_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=656} + %stack.1921 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dot_general.1328), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/stack" stack_frame_id=719} + %stack.1922 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1920, %stack.1921), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/stack" stack_frame_id=719} + %stack.2032 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1922), metadata={op_name="jit(compiled_generate_function)/stack" stack_frame_id=804} + %broadcast_in_dim.1952 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dot_general.1328), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/broadcast_in_dim" stack_frame_id=725} + %reshape.903 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1952), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/reshape" stack_frame_id=725} + %iota.710 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/iota" stack_frame_id=525} + %broadcast_in_dim.1941 = f32[41,1]{1,0} reshape(%iota.710), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/broadcast_in_dim" stack_frame_id=528} + %ge.656 = f32[41,1]{1,0} broadcast(%broadcast_in_dim.1941), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/ge" stack_frame_id=532} + %ge.657 = f32[41]{0} reshape(%ge.656), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/ge" stack_frame_id=532} + %ge.658 = f32[41,41]{1,0} broadcast(%ge.657), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/ge" stack_frame_id=532} + %iota.711 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/iota" stack_frame_id=531} + %broadcast_in_dim.1942 = f32[1,41]{1,0} reshape(%iota.711), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/broadcast_in_dim" stack_frame_id=532} + %ge.659 = f32[1,41]{1,0} broadcast(%broadcast_in_dim.1942), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/ge" stack_frame_id=532} + %ge.660 = f32[41]{0} reshape(%ge.659), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/ge" stack_frame_id=532} + %ge.661 = f32[41,41]{1,0} broadcast(%ge.660), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/ge" stack_frame_id=532} + %ge.662 = pred[41,41]{1,0} compare(%ge.658, %ge.661), direction=GE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/ge" stack_frame_id=532} + %broadcast_in_dim.1943 = pred[1,41,41]{2,1,0} reshape(%ge.662), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/broadcast_in_dim" stack_frame_id=535} + %convert_element_type.2133 = s32[1,41,41]{2,1,0} convert(%broadcast_in_dim.1943), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/convert_element_type" stack_frame_id=551} + %iota.712 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/iota" stack_frame_id=538} + %broadcast_in_dim.1944 = s32[41,1]{1,0} reshape(%iota.712), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/broadcast_in_dim" stack_frame_id=536} + %lt.798 = s32[41,1]{1,0} broadcast(%broadcast_in_dim.1944), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/lt" stack_frame_id=543} + %lt.799 = s32[41]{0} reshape(%lt.798), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/lt" stack_frame_id=543} + %lt.800 = s32[41,41]{1,0} broadcast(%lt.799), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/lt" stack_frame_id=543} + %iota.713 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/iota" stack_frame_id=541} + %broadcast_in_dim.1945 = s32[1,41]{1,0} reshape(%iota.713), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/broadcast_in_dim" stack_frame_id=539} + %add.1298 = s32[1,41]{1,0} add(%broadcast_in_dim.1945, %broadcast.57), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/add" stack_frame_id=542} + %lt.801 = s32[1,41]{1,0} broadcast(%add.1298), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/lt" stack_frame_id=543} + %lt.802 = s32[41]{0} reshape(%lt.801), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/lt" stack_frame_id=543} + %lt.803 = s32[41,41]{1,0} broadcast(%lt.802), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/lt" stack_frame_id=543} + %lt.804 = pred[41,41]{1,0} compare(%lt.800, %lt.803), direction=LT, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/lt" stack_frame_id=543} + %convert_element_type.2132 = s32[41,41]{1,0} convert(%lt.804), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/convert_element_type" stack_frame_id=547} + %broadcast_in_dim.1946 = s32[1,41,41]{2,1,0} reshape(%convert_element_type.2132), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/broadcast_in_dim" stack_frame_id=551} + %min.120 = s32[1,41,41]{2,1,0} minimum(%convert_element_type.2133, %broadcast_in_dim.1946), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/min" stack_frame_id=551} + %broadcast_in_dim.1953 = s32[1,1,41,41]{3,2,1,0} reshape(%min.120), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/broadcast_in_dim" stack_frame_id=728} + %convert_element_type.2141 = pred[1,1,41,41]{3,2,1,0} compare(%broadcast_in_dim.1953, %broadcast.49), direction=NE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/convert_element_type" stack_frame_id=732} + %broadcast_in_dim.1954 = pred[1,1,1,41,41]{4,3,2,1,0} reshape(%convert_element_type.2141), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/broadcast_in_dim" stack_frame_id=34} + %and.373 = pred[1,1,1,41,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1954), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/vmap()/and" stack_frame_id=34} + %and.374 = pred[1,1,41,41]{3,2,1,0} reshape(%and.373), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/vmap()/and" stack_frame_id=34} + %and.375 = pred[1,1,32,41,41]{4,3,2,1,0} broadcast(%and.374), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/vmap()/and" stack_frame_id=34} + %state_1__227_.1 = bf16[4096,32,128]{2,1,0} parameter(230), metadata={op_name="state[1][227]"} + %dot_general.1325 = bf16[1,41,32,128]{3,2,1,0} dot(%convert_element_type.2136, %state_1__227_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=577} + %iota.715 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/iota" stack_frame_id=598} + %broadcast_in_dim.1949 = f32[1,41]{1,0} reshape(%iota.715), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/broadcast_in_dim" stack_frame_id=601} + %iota.714 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/iota" stack_frame_id=588} + %mul.3983 = f32[64]{0} multiply(%iota.714, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=588} + %div.1220 = f32[64]{0} divide(%mul.3983, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/div" stack_frame_id=589} + %neg.490 = f32[64]{0} negate(%div.1220), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/neg" stack_frame_id=590} + %pow.485 = f32[64]{0} power(%broadcast.51, %neg.490), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/pow" stack_frame_id=593} + %div.1221 = f32[64]{0} divide(%pow.485, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/div" stack_frame_id=594} + %dot_general.1326 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1949, %div.1221), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/bi,j->bij/dot_general" stack_frame_id=604} + %stack.1908 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1326), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/stack" stack_frame_id=607} + %stack.1909 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1326), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/stack" stack_frame_id=607} + %stack.1910 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1908, %stack.1909), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/stack" stack_frame_id=607} + %reshape.898 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1910), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/reshape"} + %cos.241 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.898), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/cos" stack_frame_id=611} + %convert_element_type.2137 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.241), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/convert_element_type" stack_frame_id=615} + %mul.3984 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2137), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=635} + %mul.3985 = bf16[1,41,128]{2,1,0} reshape(%mul.3984), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=635} + %mul.3986 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3985), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=635} + %mul.3987 = bf16[1,41,32,128]{3,2,1,0} multiply(%dot_general.1325, %mul.3986), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=635} + %split.483 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1325), slice={[0:1], [0:41], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/split" stack_frame_id=627} + %neg.491 = bf16[1,41,32,64]{3,2,1,0} negate(%split.483), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/neg" stack_frame_id=628} + %stack.1911 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%neg.491), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/stack" stack_frame_id=631} + %split.482 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1325), slice={[0:1], [0:41], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/split" stack_frame_id=627} + %stack.1912 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%split.482), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/stack" stack_frame_id=631} + %stack.1913 = bf16[1,41,32,2,64]{4,3,2,1,0} concatenate(%stack.1911, %stack.1912), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/stack" stack_frame_id=631} + %reshape.899 = bf16[1,41,32,128]{3,2,1,0} reshape(%stack.1913), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/reshape" stack_frame_id=634} + %sin.241 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.898), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/sin" stack_frame_id=619} + %convert_element_type.2138 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.241), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/convert_element_type" stack_frame_id=623} + %mul.3988 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2138), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=636} + %mul.3989 = bf16[1,41,128]{2,1,0} reshape(%mul.3988), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=636} + %mul.3990 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.3989), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=636} + %mul.3991 = bf16[1,41,32,128]{3,2,1,0} multiply(%reshape.899, %mul.3990), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/mul" stack_frame_id=636} + %add.1300 = bf16[1,41,32,128]{3,2,1,0} add(%mul.3987, %mul.3991), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/rotary_embedding_25/add" stack_frame_id=637} + %reshape.904 = bf16[1,41,32,1,128]{4,3,2,1,0} reshape(%add.1300), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/reshape" stack_frame_id=34} + %broadcast_in_dim.1951 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%add.1301), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/broadcast_in_dim" stack_frame_id=722} + %reshape.902 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1951), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/reshape" stack_frame_id=722} + %dot_general.1330 = f32[1,32,41,1,41]{4,3,2,1,0} dot(%reshape.904, %reshape.902), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=34} + %mul.4001 = f32[1,32,41,1,41]{4,3,2,1,0} multiply(%dot_general.1330, %broadcast.48), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/vmap()/mul" stack_frame_id=34} + %vmap_jit__where__.120 = f32[1,1,32,41,41]{4,3,2,1,0} call(%and.375, %mul.4001, %constant.153), to_apply=%_where_0.5, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/vmap(jit(_where))" stack_frame_id=34} + %reduce_max.506 = f32[1,1,32,41]{3,2,1,0} reduce(%vmap_jit__where__.120, %constant.152), dimensions={4}, to_apply=%region_102.107, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/vmap()/reduce_max" stack_frame_id=34} + %max.130 = f32[1,1,32,41]{3,2,1,0} maximum(%reduce_max.506, %broadcast.47), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/vmap()/max" stack_frame_id=34} + %broadcast_in_dim.1955 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%max.130), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %sub.504 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1955), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/vmap()/sub" stack_frame_id=34} + %sub.505 = f32[1,1,32,41]{3,2,1,0} reshape(%sub.504), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/vmap()/sub" stack_frame_id=34} + %sub.506 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%sub.505), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/vmap()/sub" stack_frame_id=34} + %sub.507 = f32[1,1,32,41,41]{4,3,2,1,0} subtract(%vmap_jit__where__.120, %sub.506), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/vmap()/sub" stack_frame_id=34} + %exp.126 = f32[1,1,32,41,41]{4,3,2,1,0} exponential(%sub.507), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/vmap()/exp" stack_frame_id=34} + %reduce_sum.1525 = f32[1,1,32,41]{3,2,1,0} reduce(%exp.126, %constant.155), dimensions={4}, to_apply=%region_103.108, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/vmap()/reduce_sum" stack_frame_id=34} + %broadcast_in_dim.1956 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%reduce_sum.1525), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %div.1224 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1956), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/vmap()/div" stack_frame_id=34} + %div.1225 = f32[1,1,32,41]{3,2,1,0} reshape(%div.1224), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/vmap()/div" stack_frame_id=34} + %div.1226 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%div.1225), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/vmap()/div" stack_frame_id=34} + %div.1227 = f32[1,1,32,41,41]{4,3,2,1,0} divide(%exp.126, %div.1226), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/vmap()/div" stack_frame_id=34} + %convert_element_type.2142 = bf16[1,1,32,41,41]{4,3,2,1,0} convert(%div.1227), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/vmap()/convert_element_type" stack_frame_id=34} + %dot_general.1331 = bf16[1,32,128,1,41]{4,3,2,1,0} dot(%reshape.903, %convert_element_type.2142), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=34} + %transpose.124 = bf16[1,41,32,1,128]{1,3,4,2,0} transpose(%dot_general.1331), dimensions={0,4,1,3,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/vmap()/transpose;jit(compiled_generate_function)/transformer_layer_25/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %reshape.905 = bf16[1,41,32,128]{3,2,1,0} reshape(%transpose.124), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/reshape" stack_frame_id=34} + %state_1__230_.1 = bf16[32,128,4096]{2,1,0} parameter(233), metadata={op_name="state[1][230]"} + %dot_general.1332 = bf16[1,41,4096]{2,1,0} dot(%reshape.905, %state_1__230_.1), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=741} + %add.1302 = bf16[1,41,4096]{2,1,0} add(%dot_general.1332, %add.1297), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/add" stack_frame_id=742} + %convert_element_type.2143 = f32[1,41,4096]{2,1,0} convert(%add.1302), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/feedforward_layernorm/convert_element_type" stack_frame_id=746} + %pow.487 = f32[1,41,4096]{2,1,0} power(%convert_element_type.2143, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/feedforward_layernorm/pow" stack_frame_id=749} + %reduce_sum.1526 = f32[1,41]{1,0} reduce(%pow.487, %constant.155), dimensions={2}, to_apply=%region_104.109, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/feedforward_layernorm/reduce_sum" stack_frame_id=43} + %broadcast_in_dim.1957 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1526), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/feedforward_layernorm/broadcast_in_dim" stack_frame_id=43} + %div.1228 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1957, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/feedforward_layernorm/div" stack_frame_id=43} + %add.1303 = f32[1,41,1]{2,1,0} add(%div.1228, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/feedforward_layernorm/add" stack_frame_id=750} + %rsqrt.244 = f32[1,41,1]{2,1,0} rsqrt(%add.1303), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/feedforward_layernorm/rsqrt" stack_frame_id=753} + %mul.4002 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.244), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/feedforward_layernorm/mul" stack_frame_id=754} + %mul.4003 = f32[1,41]{1,0} reshape(%mul.4002), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/feedforward_layernorm/mul" stack_frame_id=754} + %mul.4004 = f32[1,41,4096]{2,1,0} broadcast(%mul.4003), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/feedforward_layernorm/mul" stack_frame_id=754} + %mul.4005 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.2143, %mul.4004), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/feedforward_layernorm/mul" stack_frame_id=754} + %state_1__235_.1 = bf16[4096]{0} parameter(238), metadata={op_name="state[1][235]"} + %convert_element_type.2144 = f32[4096]{0} convert(%state_1__235_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/feedforward_layernorm/convert_element_type" stack_frame_id=755} + %broadcast_in_dim.1958 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.2144), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/feedforward_layernorm/broadcast_in_dim" stack_frame_id=755} + %mul.4006 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1958), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/feedforward_layernorm/mul" stack_frame_id=755} + %mul.4007 = f32[1,4096]{1,0} reshape(%mul.4006), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/feedforward_layernorm/mul" stack_frame_id=755} + %mul.4008 = f32[1,41,4096]{2,1,0} broadcast(%mul.4007), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/feedforward_layernorm/mul" stack_frame_id=755} + %mul.4009 = f32[1,41,4096]{2,1,0} multiply(%mul.4005, %mul.4008), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/feedforward_layernorm/mul" stack_frame_id=755} + %convert_element_type.2145 = bf16[1,41,4096]{2,1,0} convert(%mul.4009), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/feedforward_layernorm/convert_element_type" stack_frame_id=759} + %state_1__232_.1 = bf16[4096,14336]{1,0} parameter(235), metadata={op_name="state[1][232]"} + %dot_general.1334 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.2145, %state_1__232_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/feedforward_intermediate_dense/dot_general" stack_frame_id=788} + %state_1__233_.1 = bf16[4096,14336]{1,0} parameter(236), metadata={op_name="state[1][233]"} + %dot_general.1333 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.2145, %state_1__233_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/feedforward_gate_dense/dot_general" stack_frame_id=768} + %convert_element_type.2146 = f32[1,41,14336]{2,1,0} convert(%dot_general.1333), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/convert_element_type" stack_frame_id=772} + %jit_silu_.120 = f32[1,41,14336]{2,1,0} call(%convert_element_type.2146), to_apply=%silu.9, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/jit(silu)" stack_frame_id=775} + %convert_element_type.2147 = bf16[1,41,14336]{2,1,0} convert(%jit_silu_.120), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/convert_element_type" stack_frame_id=779} + %mul.4010 = bf16[1,41,14336]{2,1,0} multiply(%dot_general.1334, %convert_element_type.2147), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/mul" stack_frame_id=791} + %state_1__234_.1 = bf16[14336,4096]{1,0} parameter(237), metadata={op_name="state[1][234]"} + %dot_general.1335 = bf16[1,41,4096]{2,1,0} dot(%mul.4010, %state_1__234_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/feedforward_output_dense/dot_general" stack_frame_id=800} + %add.1304 = bf16[1,41,4096]{2,1,0} add(%dot_general.1335, %add.1302), metadata={op_name="jit(compiled_generate_function)/transformer_layer_25/add" stack_frame_id=801} + %convert_element_type.2150 = f32[1,41,4096]{2,1,0} convert(%add.1304), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention_layernorm/convert_element_type" stack_frame_id=555} + %pow.488 = f32[1,41,4096]{2,1,0} power(%convert_element_type.2150, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention_layernorm/pow" stack_frame_id=558} + %reduce_sum.1527 = f32[1,41]{1,0} reduce(%pow.488, %constant.155), dimensions={2}, to_apply=%region_105.110, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention_layernorm/reduce_sum" stack_frame_id=24} + %broadcast_in_dim.1965 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1527), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention_layernorm/broadcast_in_dim" stack_frame_id=24} + %div.1229 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1965, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention_layernorm/div" stack_frame_id=24} + %add.1306 = f32[1,41,1]{2,1,0} add(%div.1229, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention_layernorm/add" stack_frame_id=559} + %rsqrt.245 = f32[1,41,1]{2,1,0} rsqrt(%add.1306), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention_layernorm/rsqrt" stack_frame_id=562} + %mul.4011 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.245), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention_layernorm/mul" stack_frame_id=563} + %mul.4012 = f32[1,41]{1,0} reshape(%mul.4011), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention_layernorm/mul" stack_frame_id=563} + %mul.4013 = f32[1,41,4096]{2,1,0} broadcast(%mul.4012), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention_layernorm/mul" stack_frame_id=563} + %mul.4014 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.2150, %mul.4013), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention_layernorm/mul" stack_frame_id=563} + %state_1__240_.1 = bf16[4096]{0} parameter(243), metadata={op_name="state[1][240]"} + %convert_element_type.2151 = f32[4096]{0} convert(%state_1__240_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention_layernorm/convert_element_type" stack_frame_id=564} + %broadcast_in_dim.1966 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.2151), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention_layernorm/broadcast_in_dim" stack_frame_id=564} + %mul.4015 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1966), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention_layernorm/mul" stack_frame_id=564} + %mul.4016 = f32[1,4096]{1,0} reshape(%mul.4015), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention_layernorm/mul" stack_frame_id=564} + %mul.4017 = f32[1,41,4096]{2,1,0} broadcast(%mul.4016), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention_layernorm/mul" stack_frame_id=564} + %mul.4018 = f32[1,41,4096]{2,1,0} multiply(%mul.4014, %mul.4017), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention_layernorm/mul" stack_frame_id=564} + %convert_element_type.2152 = bf16[1,41,4096]{2,1,0} convert(%mul.4018), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention_layernorm/convert_element_type" stack_frame_id=568} + %state_1__237_.1 = bf16[4096,8,128]{2,1,0} parameter(240), metadata={op_name="state[1][237]"} + %dot_general.1338 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.2152, %state_1__237_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=647} + %iota.725 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/iota" stack_frame_id=677} + %broadcast_in_dim.1968 = f32[1,41]{1,0} reshape(%iota.725), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/broadcast_in_dim" stack_frame_id=680} + %iota.724 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/iota" stack_frame_id=667} + %mul.4028 = f32[64]{0} multiply(%iota.724, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=667} + %div.1232 = f32[64]{0} divide(%mul.4028, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/div" stack_frame_id=668} + %neg.496 = f32[64]{0} negate(%div.1232), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/neg" stack_frame_id=669} + %pow.490 = f32[64]{0} power(%broadcast.51, %neg.496), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/pow" stack_frame_id=672} + %div.1233 = f32[64]{0} divide(%pow.490, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/div" stack_frame_id=673} + %dot_general.1340 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1968, %div.1233), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/bi,j->bij/dot_general" stack_frame_id=683} + %stack.1929 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1340), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/stack" stack_frame_id=686} + %stack.1930 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1340), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/stack" stack_frame_id=686} + %stack.1931 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1929, %stack.1930), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/stack" stack_frame_id=686} + %reshape.908 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1931), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/reshape"} + %cos.244 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.908), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/cos" stack_frame_id=690} + %convert_element_type.2155 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.244), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/convert_element_type" stack_frame_id=694} + %mul.4029 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2155), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=714} + %mul.4030 = bf16[1,41,128]{2,1,0} reshape(%mul.4029), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=714} + %mul.4031 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.4030), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=714} + %mul.4032 = bf16[1,41,8,128]{3,2,1,0} multiply(%dot_general.1338, %mul.4031), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=714} + %split.489 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1338), slice={[0:1], [0:41], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/split" stack_frame_id=706} + %neg.497 = bf16[1,41,8,64]{3,2,1,0} negate(%split.489), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/neg" stack_frame_id=707} + %stack.1932 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%neg.497), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/stack" stack_frame_id=710} + %split.488 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1338), slice={[0:1], [0:41], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/split" stack_frame_id=706} + %stack.1933 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%split.488), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/stack" stack_frame_id=710} + %stack.1934 = bf16[1,41,8,2,64]{4,3,2,1,0} concatenate(%stack.1932, %stack.1933), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/stack" stack_frame_id=710} + %reshape.909 = bf16[1,41,8,128]{3,2,1,0} reshape(%stack.1934), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/reshape" stack_frame_id=713} + %sin.244 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.908), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/sin" stack_frame_id=698} + %convert_element_type.2156 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.244), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/convert_element_type" stack_frame_id=702} + %mul.4033 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2156), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=715} + %mul.4034 = bf16[1,41,128]{2,1,0} reshape(%mul.4033), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=715} + %mul.4035 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.4034), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=715} + %mul.4036 = bf16[1,41,8,128]{3,2,1,0} multiply(%reshape.909, %mul.4035), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=715} + %add.1308 = bf16[1,41,8,128]{3,2,1,0} add(%mul.4032, %mul.4036), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/add" stack_frame_id=716} + %stack.1935 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%add.1308), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/stack" stack_frame_id=719} + %state_1__238_.1 = bf16[4096,8,128]{2,1,0} parameter(241), metadata={op_name="state[1][238]"} + %dot_general.1339 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.2152, %state_1__238_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=656} + %stack.1936 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dot_general.1339), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/stack" stack_frame_id=719} + %stack.1937 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1935, %stack.1936), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/stack" stack_frame_id=719} + %stack.2033 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1937), metadata={op_name="jit(compiled_generate_function)/stack" stack_frame_id=804} + %broadcast_in_dim.1970 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dot_general.1339), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/broadcast_in_dim" stack_frame_id=725} + %reshape.911 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1970), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/reshape" stack_frame_id=725} + %iota.718 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/iota" stack_frame_id=525} + %broadcast_in_dim.1959 = f32[41,1]{1,0} reshape(%iota.718), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/broadcast_in_dim" stack_frame_id=528} + %ge.663 = f32[41,1]{1,0} broadcast(%broadcast_in_dim.1959), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/ge" stack_frame_id=532} + %ge.664 = f32[41]{0} reshape(%ge.663), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/ge" stack_frame_id=532} + %ge.665 = f32[41,41]{1,0} broadcast(%ge.664), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/ge" stack_frame_id=532} + %iota.719 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/iota" stack_frame_id=531} + %broadcast_in_dim.1960 = f32[1,41]{1,0} reshape(%iota.719), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/broadcast_in_dim" stack_frame_id=532} + %ge.666 = f32[1,41]{1,0} broadcast(%broadcast_in_dim.1960), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/ge" stack_frame_id=532} + %ge.667 = f32[41]{0} reshape(%ge.666), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/ge" stack_frame_id=532} + %ge.668 = f32[41,41]{1,0} broadcast(%ge.667), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/ge" stack_frame_id=532} + %ge.669 = pred[41,41]{1,0} compare(%ge.665, %ge.668), direction=GE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/ge" stack_frame_id=532} + %broadcast_in_dim.1961 = pred[1,41,41]{2,1,0} reshape(%ge.669), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/broadcast_in_dim" stack_frame_id=535} + %convert_element_type.2149 = s32[1,41,41]{2,1,0} convert(%broadcast_in_dim.1961), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/convert_element_type" stack_frame_id=551} + %iota.720 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/iota" stack_frame_id=538} + %broadcast_in_dim.1962 = s32[41,1]{1,0} reshape(%iota.720), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/broadcast_in_dim" stack_frame_id=536} + %lt.805 = s32[41,1]{1,0} broadcast(%broadcast_in_dim.1962), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/lt" stack_frame_id=543} + %lt.806 = s32[41]{0} reshape(%lt.805), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/lt" stack_frame_id=543} + %lt.807 = s32[41,41]{1,0} broadcast(%lt.806), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/lt" stack_frame_id=543} + %iota.721 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/iota" stack_frame_id=541} + %broadcast_in_dim.1963 = s32[1,41]{1,0} reshape(%iota.721), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/broadcast_in_dim" stack_frame_id=539} + %add.1305 = s32[1,41]{1,0} add(%broadcast_in_dim.1963, %broadcast.57), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/add" stack_frame_id=542} + %lt.808 = s32[1,41]{1,0} broadcast(%add.1305), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/lt" stack_frame_id=543} + %lt.809 = s32[41]{0} reshape(%lt.808), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/lt" stack_frame_id=543} + %lt.810 = s32[41,41]{1,0} broadcast(%lt.809), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/lt" stack_frame_id=543} + %lt.811 = pred[41,41]{1,0} compare(%lt.807, %lt.810), direction=LT, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/lt" stack_frame_id=543} + %convert_element_type.2148 = s32[41,41]{1,0} convert(%lt.811), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/convert_element_type" stack_frame_id=547} + %broadcast_in_dim.1964 = s32[1,41,41]{2,1,0} reshape(%convert_element_type.2148), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/broadcast_in_dim" stack_frame_id=551} + %min.121 = s32[1,41,41]{2,1,0} minimum(%convert_element_type.2149, %broadcast_in_dim.1964), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/min" stack_frame_id=551} + %broadcast_in_dim.1971 = s32[1,1,41,41]{3,2,1,0} reshape(%min.121), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/broadcast_in_dim" stack_frame_id=728} + %convert_element_type.2157 = pred[1,1,41,41]{3,2,1,0} compare(%broadcast_in_dim.1971, %broadcast.49), direction=NE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/convert_element_type" stack_frame_id=732} + %broadcast_in_dim.1972 = pred[1,1,1,41,41]{4,3,2,1,0} reshape(%convert_element_type.2157), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/broadcast_in_dim" stack_frame_id=34} + %and.376 = pred[1,1,1,41,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1972), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/vmap()/and" stack_frame_id=34} + %and.377 = pred[1,1,41,41]{3,2,1,0} reshape(%and.376), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/vmap()/and" stack_frame_id=34} + %and.378 = pred[1,1,32,41,41]{4,3,2,1,0} broadcast(%and.377), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/vmap()/and" stack_frame_id=34} + %state_1__236_.1 = bf16[4096,32,128]{2,1,0} parameter(239), metadata={op_name="state[1][236]"} + %dot_general.1336 = bf16[1,41,32,128]{3,2,1,0} dot(%convert_element_type.2152, %state_1__236_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=577} + %iota.723 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/iota" stack_frame_id=598} + %broadcast_in_dim.1967 = f32[1,41]{1,0} reshape(%iota.723), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/broadcast_in_dim" stack_frame_id=601} + %iota.722 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/iota" stack_frame_id=588} + %mul.4019 = f32[64]{0} multiply(%iota.722, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=588} + %div.1230 = f32[64]{0} divide(%mul.4019, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/div" stack_frame_id=589} + %neg.494 = f32[64]{0} negate(%div.1230), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/neg" stack_frame_id=590} + %pow.489 = f32[64]{0} power(%broadcast.51, %neg.494), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/pow" stack_frame_id=593} + %div.1231 = f32[64]{0} divide(%pow.489, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/div" stack_frame_id=594} + %dot_general.1337 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1967, %div.1231), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/bi,j->bij/dot_general" stack_frame_id=604} + %stack.1923 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1337), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/stack" stack_frame_id=607} + %stack.1924 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1337), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/stack" stack_frame_id=607} + %stack.1925 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1923, %stack.1924), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/stack" stack_frame_id=607} + %reshape.906 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1925), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/reshape"} + %cos.243 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.906), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/cos" stack_frame_id=611} + %convert_element_type.2153 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.243), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/convert_element_type" stack_frame_id=615} + %mul.4020 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2153), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=635} + %mul.4021 = bf16[1,41,128]{2,1,0} reshape(%mul.4020), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=635} + %mul.4022 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.4021), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=635} + %mul.4023 = bf16[1,41,32,128]{3,2,1,0} multiply(%dot_general.1336, %mul.4022), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=635} + %split.487 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1336), slice={[0:1], [0:41], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/split" stack_frame_id=627} + %neg.495 = bf16[1,41,32,64]{3,2,1,0} negate(%split.487), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/neg" stack_frame_id=628} + %stack.1926 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%neg.495), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/stack" stack_frame_id=631} + %split.486 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1336), slice={[0:1], [0:41], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/split" stack_frame_id=627} + %stack.1927 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%split.486), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/stack" stack_frame_id=631} + %stack.1928 = bf16[1,41,32,2,64]{4,3,2,1,0} concatenate(%stack.1926, %stack.1927), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/stack" stack_frame_id=631} + %reshape.907 = bf16[1,41,32,128]{3,2,1,0} reshape(%stack.1928), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/reshape" stack_frame_id=634} + %sin.243 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.906), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/sin" stack_frame_id=619} + %convert_element_type.2154 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.243), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/convert_element_type" stack_frame_id=623} + %mul.4024 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2154), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=636} + %mul.4025 = bf16[1,41,128]{2,1,0} reshape(%mul.4024), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=636} + %mul.4026 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.4025), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=636} + %mul.4027 = bf16[1,41,32,128]{3,2,1,0} multiply(%reshape.907, %mul.4026), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/mul" stack_frame_id=636} + %add.1307 = bf16[1,41,32,128]{3,2,1,0} add(%mul.4023, %mul.4027), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/rotary_embedding_26/add" stack_frame_id=637} + %reshape.912 = bf16[1,41,32,1,128]{4,3,2,1,0} reshape(%add.1307), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/reshape" stack_frame_id=34} + %broadcast_in_dim.1969 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%add.1308), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/broadcast_in_dim" stack_frame_id=722} + %reshape.910 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1969), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/reshape" stack_frame_id=722} + %dot_general.1341 = f32[1,32,41,1,41]{4,3,2,1,0} dot(%reshape.912, %reshape.910), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=34} + %mul.4037 = f32[1,32,41,1,41]{4,3,2,1,0} multiply(%dot_general.1341, %broadcast.48), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/vmap()/mul" stack_frame_id=34} + %vmap_jit__where__.121 = f32[1,1,32,41,41]{4,3,2,1,0} call(%and.378, %mul.4037, %constant.153), to_apply=%_where_0.5, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/vmap(jit(_where))" stack_frame_id=34} + %reduce_max.507 = f32[1,1,32,41]{3,2,1,0} reduce(%vmap_jit__where__.121, %constant.152), dimensions={4}, to_apply=%region_106.111, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/vmap()/reduce_max" stack_frame_id=34} + %max.131 = f32[1,1,32,41]{3,2,1,0} maximum(%reduce_max.507, %broadcast.47), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/vmap()/max" stack_frame_id=34} + %broadcast_in_dim.1973 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%max.131), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %sub.508 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1973), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/vmap()/sub" stack_frame_id=34} + %sub.509 = f32[1,1,32,41]{3,2,1,0} reshape(%sub.508), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/vmap()/sub" stack_frame_id=34} + %sub.510 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%sub.509), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/vmap()/sub" stack_frame_id=34} + %sub.511 = f32[1,1,32,41,41]{4,3,2,1,0} subtract(%vmap_jit__where__.121, %sub.510), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/vmap()/sub" stack_frame_id=34} + %exp.127 = f32[1,1,32,41,41]{4,3,2,1,0} exponential(%sub.511), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/vmap()/exp" stack_frame_id=34} + %reduce_sum.1528 = f32[1,1,32,41]{3,2,1,0} reduce(%exp.127, %constant.155), dimensions={4}, to_apply=%region_107.112, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/vmap()/reduce_sum" stack_frame_id=34} + %broadcast_in_dim.1974 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%reduce_sum.1528), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %div.1234 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1974), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/vmap()/div" stack_frame_id=34} + %div.1235 = f32[1,1,32,41]{3,2,1,0} reshape(%div.1234), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/vmap()/div" stack_frame_id=34} + %div.1236 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%div.1235), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/vmap()/div" stack_frame_id=34} + %div.1237 = f32[1,1,32,41,41]{4,3,2,1,0} divide(%exp.127, %div.1236), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/vmap()/div" stack_frame_id=34} + %convert_element_type.2158 = bf16[1,1,32,41,41]{4,3,2,1,0} convert(%div.1237), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/vmap()/convert_element_type" stack_frame_id=34} + %dot_general.1342 = bf16[1,32,128,1,41]{4,3,2,1,0} dot(%reshape.911, %convert_element_type.2158), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=34} + %transpose.125 = bf16[1,41,32,1,128]{1,3,4,2,0} transpose(%dot_general.1342), dimensions={0,4,1,3,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/vmap()/transpose;jit(compiled_generate_function)/transformer_layer_26/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %reshape.913 = bf16[1,41,32,128]{3,2,1,0} reshape(%transpose.125), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/reshape" stack_frame_id=34} + %state_1__239_.1 = bf16[32,128,4096]{2,1,0} parameter(242), metadata={op_name="state[1][239]"} + %dot_general.1343 = bf16[1,41,4096]{2,1,0} dot(%reshape.913, %state_1__239_.1), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=741} + %add.1309 = bf16[1,41,4096]{2,1,0} add(%dot_general.1343, %add.1304), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/add" stack_frame_id=742} + %convert_element_type.2159 = f32[1,41,4096]{2,1,0} convert(%add.1309), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/feedforward_layernorm/convert_element_type" stack_frame_id=746} + %pow.491 = f32[1,41,4096]{2,1,0} power(%convert_element_type.2159, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/feedforward_layernorm/pow" stack_frame_id=749} + %reduce_sum.1529 = f32[1,41]{1,0} reduce(%pow.491, %constant.155), dimensions={2}, to_apply=%region_108.113, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/feedforward_layernorm/reduce_sum" stack_frame_id=43} + %broadcast_in_dim.1975 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1529), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/feedforward_layernorm/broadcast_in_dim" stack_frame_id=43} + %div.1238 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1975, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/feedforward_layernorm/div" stack_frame_id=43} + %add.1310 = f32[1,41,1]{2,1,0} add(%div.1238, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/feedforward_layernorm/add" stack_frame_id=750} + %rsqrt.246 = f32[1,41,1]{2,1,0} rsqrt(%add.1310), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/feedforward_layernorm/rsqrt" stack_frame_id=753} + %mul.4038 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.246), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/feedforward_layernorm/mul" stack_frame_id=754} + %mul.4039 = f32[1,41]{1,0} reshape(%mul.4038), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/feedforward_layernorm/mul" stack_frame_id=754} + %mul.4040 = f32[1,41,4096]{2,1,0} broadcast(%mul.4039), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/feedforward_layernorm/mul" stack_frame_id=754} + %mul.4041 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.2159, %mul.4040), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/feedforward_layernorm/mul" stack_frame_id=754} + %state_1__244_.1 = bf16[4096]{0} parameter(247), metadata={op_name="state[1][244]"} + %convert_element_type.2160 = f32[4096]{0} convert(%state_1__244_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/feedforward_layernorm/convert_element_type" stack_frame_id=755} + %broadcast_in_dim.1976 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.2160), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/feedforward_layernorm/broadcast_in_dim" stack_frame_id=755} + %mul.4042 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1976), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/feedforward_layernorm/mul" stack_frame_id=755} + %mul.4043 = f32[1,4096]{1,0} reshape(%mul.4042), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/feedforward_layernorm/mul" stack_frame_id=755} + %mul.4044 = f32[1,41,4096]{2,1,0} broadcast(%mul.4043), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/feedforward_layernorm/mul" stack_frame_id=755} + %mul.4045 = f32[1,41,4096]{2,1,0} multiply(%mul.4041, %mul.4044), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/feedforward_layernorm/mul" stack_frame_id=755} + %convert_element_type.2161 = bf16[1,41,4096]{2,1,0} convert(%mul.4045), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/feedforward_layernorm/convert_element_type" stack_frame_id=759} + %state_1__241_.1 = bf16[4096,14336]{1,0} parameter(244), metadata={op_name="state[1][241]"} + %dot_general.1345 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.2161, %state_1__241_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/feedforward_intermediate_dense/dot_general" stack_frame_id=788} + %state_1__242_.1 = bf16[4096,14336]{1,0} parameter(245), metadata={op_name="state[1][242]"} + %dot_general.1344 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.2161, %state_1__242_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/feedforward_gate_dense/dot_general" stack_frame_id=768} + %convert_element_type.2162 = f32[1,41,14336]{2,1,0} convert(%dot_general.1344), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/convert_element_type" stack_frame_id=772} + %jit_silu_.121 = f32[1,41,14336]{2,1,0} call(%convert_element_type.2162), to_apply=%silu.9, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/jit(silu)" stack_frame_id=775} + %convert_element_type.2163 = bf16[1,41,14336]{2,1,0} convert(%jit_silu_.121), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/convert_element_type" stack_frame_id=779} + %mul.4046 = bf16[1,41,14336]{2,1,0} multiply(%dot_general.1345, %convert_element_type.2163), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/mul" stack_frame_id=791} + %state_1__243_.1 = bf16[14336,4096]{1,0} parameter(246), metadata={op_name="state[1][243]"} + %dot_general.1346 = bf16[1,41,4096]{2,1,0} dot(%mul.4046, %state_1__243_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/feedforward_output_dense/dot_general" stack_frame_id=800} + %add.1311 = bf16[1,41,4096]{2,1,0} add(%dot_general.1346, %add.1309), metadata={op_name="jit(compiled_generate_function)/transformer_layer_26/add" stack_frame_id=801} + %convert_element_type.2166 = f32[1,41,4096]{2,1,0} convert(%add.1311), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention_layernorm/convert_element_type" stack_frame_id=555} + %pow.492 = f32[1,41,4096]{2,1,0} power(%convert_element_type.2166, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention_layernorm/pow" stack_frame_id=558} + %reduce_sum.1530 = f32[1,41]{1,0} reduce(%pow.492, %constant.155), dimensions={2}, to_apply=%region_109.114, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention_layernorm/reduce_sum" stack_frame_id=24} + %broadcast_in_dim.1983 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1530), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention_layernorm/broadcast_in_dim" stack_frame_id=24} + %div.1239 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1983, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention_layernorm/div" stack_frame_id=24} + %add.1313 = f32[1,41,1]{2,1,0} add(%div.1239, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention_layernorm/add" stack_frame_id=559} + %rsqrt.247 = f32[1,41,1]{2,1,0} rsqrt(%add.1313), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention_layernorm/rsqrt" stack_frame_id=562} + %mul.4047 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.247), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention_layernorm/mul" stack_frame_id=563} + %mul.4048 = f32[1,41]{1,0} reshape(%mul.4047), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention_layernorm/mul" stack_frame_id=563} + %mul.4049 = f32[1,41,4096]{2,1,0} broadcast(%mul.4048), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention_layernorm/mul" stack_frame_id=563} + %mul.4050 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.2166, %mul.4049), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention_layernorm/mul" stack_frame_id=563} + %state_1__249_.1 = bf16[4096]{0} parameter(252), metadata={op_name="state[1][249]"} + %convert_element_type.2167 = f32[4096]{0} convert(%state_1__249_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention_layernorm/convert_element_type" stack_frame_id=564} + %broadcast_in_dim.1984 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.2167), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention_layernorm/broadcast_in_dim" stack_frame_id=564} + %mul.4051 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1984), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention_layernorm/mul" stack_frame_id=564} + %mul.4052 = f32[1,4096]{1,0} reshape(%mul.4051), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention_layernorm/mul" stack_frame_id=564} + %mul.4053 = f32[1,41,4096]{2,1,0} broadcast(%mul.4052), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention_layernorm/mul" stack_frame_id=564} + %mul.4054 = f32[1,41,4096]{2,1,0} multiply(%mul.4050, %mul.4053), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention_layernorm/mul" stack_frame_id=564} + %convert_element_type.2168 = bf16[1,41,4096]{2,1,0} convert(%mul.4054), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention_layernorm/convert_element_type" stack_frame_id=568} + %state_1__246_.1 = bf16[4096,8,128]{2,1,0} parameter(249), metadata={op_name="state[1][246]"} + %dot_general.1349 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.2168, %state_1__246_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=647} + %iota.733 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/iota" stack_frame_id=677} + %broadcast_in_dim.1986 = f32[1,41]{1,0} reshape(%iota.733), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/broadcast_in_dim" stack_frame_id=680} + %iota.732 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/iota" stack_frame_id=667} + %mul.4064 = f32[64]{0} multiply(%iota.732, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=667} + %div.1242 = f32[64]{0} divide(%mul.4064, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/div" stack_frame_id=668} + %neg.500 = f32[64]{0} negate(%div.1242), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/neg" stack_frame_id=669} + %pow.494 = f32[64]{0} power(%broadcast.51, %neg.500), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/pow" stack_frame_id=672} + %div.1243 = f32[64]{0} divide(%pow.494, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/div" stack_frame_id=673} + %dot_general.1351 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1986, %div.1243), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/bi,j->bij/dot_general" stack_frame_id=683} + %stack.1944 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1351), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/stack" stack_frame_id=686} + %stack.1945 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1351), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/stack" stack_frame_id=686} + %stack.1946 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1944, %stack.1945), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/stack" stack_frame_id=686} + %reshape.916 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1946), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/reshape"} + %cos.246 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.916), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/cos" stack_frame_id=690} + %convert_element_type.2171 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.246), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/convert_element_type" stack_frame_id=694} + %mul.4065 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2171), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=714} + %mul.4066 = bf16[1,41,128]{2,1,0} reshape(%mul.4065), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=714} + %mul.4067 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.4066), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=714} + %mul.4068 = bf16[1,41,8,128]{3,2,1,0} multiply(%dot_general.1349, %mul.4067), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=714} + %split.493 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1349), slice={[0:1], [0:41], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/split" stack_frame_id=706} + %neg.501 = bf16[1,41,8,64]{3,2,1,0} negate(%split.493), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/neg" stack_frame_id=707} + %stack.1947 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%neg.501), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/stack" stack_frame_id=710} + %split.492 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1349), slice={[0:1], [0:41], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/split" stack_frame_id=706} + %stack.1948 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%split.492), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/stack" stack_frame_id=710} + %stack.1949 = bf16[1,41,8,2,64]{4,3,2,1,0} concatenate(%stack.1947, %stack.1948), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/stack" stack_frame_id=710} + %reshape.917 = bf16[1,41,8,128]{3,2,1,0} reshape(%stack.1949), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/reshape" stack_frame_id=713} + %sin.246 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.916), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/sin" stack_frame_id=698} + %convert_element_type.2172 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.246), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/convert_element_type" stack_frame_id=702} + %mul.4069 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2172), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=715} + %mul.4070 = bf16[1,41,128]{2,1,0} reshape(%mul.4069), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=715} + %mul.4071 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.4070), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=715} + %mul.4072 = bf16[1,41,8,128]{3,2,1,0} multiply(%reshape.917, %mul.4071), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=715} + %add.1315 = bf16[1,41,8,128]{3,2,1,0} add(%mul.4068, %mul.4072), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/add" stack_frame_id=716} + %stack.1950 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%add.1315), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/stack" stack_frame_id=719} + %state_1__247_.1 = bf16[4096,8,128]{2,1,0} parameter(250), metadata={op_name="state[1][247]"} + %dot_general.1350 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.2168, %state_1__247_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=656} + %stack.1951 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dot_general.1350), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/stack" stack_frame_id=719} + %stack.1952 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1950, %stack.1951), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/stack" stack_frame_id=719} + %stack.2034 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1952), metadata={op_name="jit(compiled_generate_function)/stack" stack_frame_id=804} + %broadcast_in_dim.1988 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dot_general.1350), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/broadcast_in_dim" stack_frame_id=725} + %reshape.919 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1988), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/reshape" stack_frame_id=725} + %iota.726 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/iota" stack_frame_id=525} + %broadcast_in_dim.1977 = f32[41,1]{1,0} reshape(%iota.726), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/broadcast_in_dim" stack_frame_id=528} + %ge.670 = f32[41,1]{1,0} broadcast(%broadcast_in_dim.1977), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/ge" stack_frame_id=532} + %ge.671 = f32[41]{0} reshape(%ge.670), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/ge" stack_frame_id=532} + %ge.672 = f32[41,41]{1,0} broadcast(%ge.671), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/ge" stack_frame_id=532} + %iota.727 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/iota" stack_frame_id=531} + %broadcast_in_dim.1978 = f32[1,41]{1,0} reshape(%iota.727), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/broadcast_in_dim" stack_frame_id=532} + %ge.673 = f32[1,41]{1,0} broadcast(%broadcast_in_dim.1978), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/ge" stack_frame_id=532} + %ge.674 = f32[41]{0} reshape(%ge.673), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/ge" stack_frame_id=532} + %ge.675 = f32[41,41]{1,0} broadcast(%ge.674), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/ge" stack_frame_id=532} + %ge.676 = pred[41,41]{1,0} compare(%ge.672, %ge.675), direction=GE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/ge" stack_frame_id=532} + %broadcast_in_dim.1979 = pred[1,41,41]{2,1,0} reshape(%ge.676), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/broadcast_in_dim" stack_frame_id=535} + %convert_element_type.2165 = s32[1,41,41]{2,1,0} convert(%broadcast_in_dim.1979), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/convert_element_type" stack_frame_id=551} + %iota.728 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/iota" stack_frame_id=538} + %broadcast_in_dim.1980 = s32[41,1]{1,0} reshape(%iota.728), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/broadcast_in_dim" stack_frame_id=536} + %lt.812 = s32[41,1]{1,0} broadcast(%broadcast_in_dim.1980), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/lt" stack_frame_id=543} + %lt.813 = s32[41]{0} reshape(%lt.812), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/lt" stack_frame_id=543} + %lt.814 = s32[41,41]{1,0} broadcast(%lt.813), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/lt" stack_frame_id=543} + %iota.729 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/iota" stack_frame_id=541} + %broadcast_in_dim.1981 = s32[1,41]{1,0} reshape(%iota.729), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/broadcast_in_dim" stack_frame_id=539} + %add.1312 = s32[1,41]{1,0} add(%broadcast_in_dim.1981, %broadcast.57), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/add" stack_frame_id=542} + %lt.815 = s32[1,41]{1,0} broadcast(%add.1312), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/lt" stack_frame_id=543} + %lt.816 = s32[41]{0} reshape(%lt.815), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/lt" stack_frame_id=543} + %lt.817 = s32[41,41]{1,0} broadcast(%lt.816), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/lt" stack_frame_id=543} + %lt.818 = pred[41,41]{1,0} compare(%lt.814, %lt.817), direction=LT, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/lt" stack_frame_id=543} + %convert_element_type.2164 = s32[41,41]{1,0} convert(%lt.818), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/convert_element_type" stack_frame_id=547} + %broadcast_in_dim.1982 = s32[1,41,41]{2,1,0} reshape(%convert_element_type.2164), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/broadcast_in_dim" stack_frame_id=551} + %min.122 = s32[1,41,41]{2,1,0} minimum(%convert_element_type.2165, %broadcast_in_dim.1982), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/min" stack_frame_id=551} + %broadcast_in_dim.1989 = s32[1,1,41,41]{3,2,1,0} reshape(%min.122), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/broadcast_in_dim" stack_frame_id=728} + %convert_element_type.2173 = pred[1,1,41,41]{3,2,1,0} compare(%broadcast_in_dim.1989, %broadcast.49), direction=NE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/convert_element_type" stack_frame_id=732} + %broadcast_in_dim.1990 = pred[1,1,1,41,41]{4,3,2,1,0} reshape(%convert_element_type.2173), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/broadcast_in_dim" stack_frame_id=34} + %and.379 = pred[1,1,1,41,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.1990), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/vmap()/and" stack_frame_id=34} + %and.380 = pred[1,1,41,41]{3,2,1,0} reshape(%and.379), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/vmap()/and" stack_frame_id=34} + %and.381 = pred[1,1,32,41,41]{4,3,2,1,0} broadcast(%and.380), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/vmap()/and" stack_frame_id=34} + %state_1__245_.1 = bf16[4096,32,128]{2,1,0} parameter(248), metadata={op_name="state[1][245]"} + %dot_general.1347 = bf16[1,41,32,128]{3,2,1,0} dot(%convert_element_type.2168, %state_1__245_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=577} + %iota.731 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/iota" stack_frame_id=598} + %broadcast_in_dim.1985 = f32[1,41]{1,0} reshape(%iota.731), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/broadcast_in_dim" stack_frame_id=601} + %iota.730 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/iota" stack_frame_id=588} + %mul.4055 = f32[64]{0} multiply(%iota.730, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=588} + %div.1240 = f32[64]{0} divide(%mul.4055, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/div" stack_frame_id=589} + %neg.498 = f32[64]{0} negate(%div.1240), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/neg" stack_frame_id=590} + %pow.493 = f32[64]{0} power(%broadcast.51, %neg.498), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/pow" stack_frame_id=593} + %div.1241 = f32[64]{0} divide(%pow.493, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/div" stack_frame_id=594} + %dot_general.1348 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.1985, %div.1241), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/bi,j->bij/dot_general" stack_frame_id=604} + %stack.1938 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1348), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/stack" stack_frame_id=607} + %stack.1939 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1348), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/stack" stack_frame_id=607} + %stack.1940 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1938, %stack.1939), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/stack" stack_frame_id=607} + %reshape.914 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1940), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/reshape"} + %cos.245 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.914), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/cos" stack_frame_id=611} + %convert_element_type.2169 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.245), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/convert_element_type" stack_frame_id=615} + %mul.4056 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2169), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=635} + %mul.4057 = bf16[1,41,128]{2,1,0} reshape(%mul.4056), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=635} + %mul.4058 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.4057), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=635} + %mul.4059 = bf16[1,41,32,128]{3,2,1,0} multiply(%dot_general.1347, %mul.4058), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=635} + %split.491 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1347), slice={[0:1], [0:41], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/split" stack_frame_id=627} + %neg.499 = bf16[1,41,32,64]{3,2,1,0} negate(%split.491), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/neg" stack_frame_id=628} + %stack.1941 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%neg.499), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/stack" stack_frame_id=631} + %split.490 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1347), slice={[0:1], [0:41], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/split" stack_frame_id=627} + %stack.1942 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%split.490), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/stack" stack_frame_id=631} + %stack.1943 = bf16[1,41,32,2,64]{4,3,2,1,0} concatenate(%stack.1941, %stack.1942), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/stack" stack_frame_id=631} + %reshape.915 = bf16[1,41,32,128]{3,2,1,0} reshape(%stack.1943), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/reshape" stack_frame_id=634} + %sin.245 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.914), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/sin" stack_frame_id=619} + %convert_element_type.2170 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.245), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/convert_element_type" stack_frame_id=623} + %mul.4060 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2170), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=636} + %mul.4061 = bf16[1,41,128]{2,1,0} reshape(%mul.4060), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=636} + %mul.4062 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.4061), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=636} + %mul.4063 = bf16[1,41,32,128]{3,2,1,0} multiply(%reshape.915, %mul.4062), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/mul" stack_frame_id=636} + %add.1314 = bf16[1,41,32,128]{3,2,1,0} add(%mul.4059, %mul.4063), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/rotary_embedding_27/add" stack_frame_id=637} + %reshape.920 = bf16[1,41,32,1,128]{4,3,2,1,0} reshape(%add.1314), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/reshape" stack_frame_id=34} + %broadcast_in_dim.1987 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%add.1315), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/broadcast_in_dim" stack_frame_id=722} + %reshape.918 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.1987), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/reshape" stack_frame_id=722} + %dot_general.1352 = f32[1,32,41,1,41]{4,3,2,1,0} dot(%reshape.920, %reshape.918), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=34} + %mul.4073 = f32[1,32,41,1,41]{4,3,2,1,0} multiply(%dot_general.1352, %broadcast.48), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/vmap()/mul" stack_frame_id=34} + %vmap_jit__where__.122 = f32[1,1,32,41,41]{4,3,2,1,0} call(%and.381, %mul.4073, %constant.153), to_apply=%_where_0.5, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/vmap(jit(_where))" stack_frame_id=34} + %reduce_max.508 = f32[1,1,32,41]{3,2,1,0} reduce(%vmap_jit__where__.122, %constant.152), dimensions={4}, to_apply=%region_110.115, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/vmap()/reduce_max" stack_frame_id=34} + %max.132 = f32[1,1,32,41]{3,2,1,0} maximum(%reduce_max.508, %broadcast.47), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/vmap()/max" stack_frame_id=34} + %broadcast_in_dim.1991 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%max.132), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %sub.512 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1991), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/vmap()/sub" stack_frame_id=34} + %sub.513 = f32[1,1,32,41]{3,2,1,0} reshape(%sub.512), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/vmap()/sub" stack_frame_id=34} + %sub.514 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%sub.513), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/vmap()/sub" stack_frame_id=34} + %sub.515 = f32[1,1,32,41,41]{4,3,2,1,0} subtract(%vmap_jit__where__.122, %sub.514), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/vmap()/sub" stack_frame_id=34} + %exp.128 = f32[1,1,32,41,41]{4,3,2,1,0} exponential(%sub.515), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/vmap()/exp" stack_frame_id=34} + %reduce_sum.1531 = f32[1,1,32,41]{3,2,1,0} reduce(%exp.128, %constant.155), dimensions={4}, to_apply=%region_111.116, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/vmap()/reduce_sum" stack_frame_id=34} + %broadcast_in_dim.1992 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%reduce_sum.1531), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %div.1244 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.1992), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/vmap()/div" stack_frame_id=34} + %div.1245 = f32[1,1,32,41]{3,2,1,0} reshape(%div.1244), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/vmap()/div" stack_frame_id=34} + %div.1246 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%div.1245), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/vmap()/div" stack_frame_id=34} + %div.1247 = f32[1,1,32,41,41]{4,3,2,1,0} divide(%exp.128, %div.1246), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/vmap()/div" stack_frame_id=34} + %convert_element_type.2174 = bf16[1,1,32,41,41]{4,3,2,1,0} convert(%div.1247), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/vmap()/convert_element_type" stack_frame_id=34} + %dot_general.1353 = bf16[1,32,128,1,41]{4,3,2,1,0} dot(%reshape.919, %convert_element_type.2174), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=34} + %transpose.126 = bf16[1,41,32,1,128]{1,3,4,2,0} transpose(%dot_general.1353), dimensions={0,4,1,3,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/vmap()/transpose;jit(compiled_generate_function)/transformer_layer_27/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %reshape.921 = bf16[1,41,32,128]{3,2,1,0} reshape(%transpose.126), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/reshape" stack_frame_id=34} + %state_1__248_.1 = bf16[32,128,4096]{2,1,0} parameter(251), metadata={op_name="state[1][248]"} + %dot_general.1354 = bf16[1,41,4096]{2,1,0} dot(%reshape.921, %state_1__248_.1), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=741} + %add.1316 = bf16[1,41,4096]{2,1,0} add(%dot_general.1354, %add.1311), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/add" stack_frame_id=742} + %convert_element_type.2175 = f32[1,41,4096]{2,1,0} convert(%add.1316), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/feedforward_layernorm/convert_element_type" stack_frame_id=746} + %pow.495 = f32[1,41,4096]{2,1,0} power(%convert_element_type.2175, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/feedforward_layernorm/pow" stack_frame_id=749} + %reduce_sum.1532 = f32[1,41]{1,0} reduce(%pow.495, %constant.155), dimensions={2}, to_apply=%region_112.117, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/feedforward_layernorm/reduce_sum" stack_frame_id=43} + %broadcast_in_dim.1993 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1532), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/feedforward_layernorm/broadcast_in_dim" stack_frame_id=43} + %div.1248 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.1993, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/feedforward_layernorm/div" stack_frame_id=43} + %add.1317 = f32[1,41,1]{2,1,0} add(%div.1248, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/feedforward_layernorm/add" stack_frame_id=750} + %rsqrt.248 = f32[1,41,1]{2,1,0} rsqrt(%add.1317), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/feedforward_layernorm/rsqrt" stack_frame_id=753} + %mul.4074 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.248), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/feedforward_layernorm/mul" stack_frame_id=754} + %mul.4075 = f32[1,41]{1,0} reshape(%mul.4074), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/feedforward_layernorm/mul" stack_frame_id=754} + %mul.4076 = f32[1,41,4096]{2,1,0} broadcast(%mul.4075), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/feedforward_layernorm/mul" stack_frame_id=754} + %mul.4077 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.2175, %mul.4076), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/feedforward_layernorm/mul" stack_frame_id=754} + %state_1__253_.1 = bf16[4096]{0} parameter(256), metadata={op_name="state[1][253]"} + %convert_element_type.2176 = f32[4096]{0} convert(%state_1__253_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/feedforward_layernorm/convert_element_type" stack_frame_id=755} + %broadcast_in_dim.1994 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.2176), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/feedforward_layernorm/broadcast_in_dim" stack_frame_id=755} + %mul.4078 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.1994), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/feedforward_layernorm/mul" stack_frame_id=755} + %mul.4079 = f32[1,4096]{1,0} reshape(%mul.4078), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/feedforward_layernorm/mul" stack_frame_id=755} + %mul.4080 = f32[1,41,4096]{2,1,0} broadcast(%mul.4079), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/feedforward_layernorm/mul" stack_frame_id=755} + %mul.4081 = f32[1,41,4096]{2,1,0} multiply(%mul.4077, %mul.4080), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/feedforward_layernorm/mul" stack_frame_id=755} + %convert_element_type.2177 = bf16[1,41,4096]{2,1,0} convert(%mul.4081), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/feedforward_layernorm/convert_element_type" stack_frame_id=759} + %state_1__250_.1 = bf16[4096,14336]{1,0} parameter(253), metadata={op_name="state[1][250]"} + %dot_general.1356 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.2177, %state_1__250_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/feedforward_intermediate_dense/dot_general" stack_frame_id=788} + %state_1__251_.1 = bf16[4096,14336]{1,0} parameter(254), metadata={op_name="state[1][251]"} + %dot_general.1355 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.2177, %state_1__251_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/feedforward_gate_dense/dot_general" stack_frame_id=768} + %convert_element_type.2178 = f32[1,41,14336]{2,1,0} convert(%dot_general.1355), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/convert_element_type" stack_frame_id=772} + %jit_silu_.122 = f32[1,41,14336]{2,1,0} call(%convert_element_type.2178), to_apply=%silu.9, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/jit(silu)" stack_frame_id=775} + %convert_element_type.2179 = bf16[1,41,14336]{2,1,0} convert(%jit_silu_.122), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/convert_element_type" stack_frame_id=779} + %mul.4082 = bf16[1,41,14336]{2,1,0} multiply(%dot_general.1356, %convert_element_type.2179), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/mul" stack_frame_id=791} + %state_1__252_.1 = bf16[14336,4096]{1,0} parameter(255), metadata={op_name="state[1][252]"} + %dot_general.1357 = bf16[1,41,4096]{2,1,0} dot(%mul.4082, %state_1__252_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/feedforward_output_dense/dot_general" stack_frame_id=800} + %add.1318 = bf16[1,41,4096]{2,1,0} add(%dot_general.1357, %add.1316), metadata={op_name="jit(compiled_generate_function)/transformer_layer_27/add" stack_frame_id=801} + %convert_element_type.2182 = f32[1,41,4096]{2,1,0} convert(%add.1318), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention_layernorm/convert_element_type" stack_frame_id=555} + %pow.496 = f32[1,41,4096]{2,1,0} power(%convert_element_type.2182, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention_layernorm/pow" stack_frame_id=558} + %reduce_sum.1533 = f32[1,41]{1,0} reduce(%pow.496, %constant.155), dimensions={2}, to_apply=%region_113.118, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention_layernorm/reduce_sum" stack_frame_id=24} + %broadcast_in_dim.2001 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1533), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention_layernorm/broadcast_in_dim" stack_frame_id=24} + %div.1249 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.2001, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention_layernorm/div" stack_frame_id=24} + %add.1320 = f32[1,41,1]{2,1,0} add(%div.1249, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention_layernorm/add" stack_frame_id=559} + %rsqrt.249 = f32[1,41,1]{2,1,0} rsqrt(%add.1320), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention_layernorm/rsqrt" stack_frame_id=562} + %mul.4083 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.249), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention_layernorm/mul" stack_frame_id=563} + %mul.4084 = f32[1,41]{1,0} reshape(%mul.4083), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention_layernorm/mul" stack_frame_id=563} + %mul.4085 = f32[1,41,4096]{2,1,0} broadcast(%mul.4084), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention_layernorm/mul" stack_frame_id=563} + %mul.4086 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.2182, %mul.4085), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention_layernorm/mul" stack_frame_id=563} + %state_1__258_.1 = bf16[4096]{0} parameter(261), metadata={op_name="state[1][258]"} + %convert_element_type.2183 = f32[4096]{0} convert(%state_1__258_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention_layernorm/convert_element_type" stack_frame_id=564} + %broadcast_in_dim.2002 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.2183), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention_layernorm/broadcast_in_dim" stack_frame_id=564} + %mul.4087 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.2002), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention_layernorm/mul" stack_frame_id=564} + %mul.4088 = f32[1,4096]{1,0} reshape(%mul.4087), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention_layernorm/mul" stack_frame_id=564} + %mul.4089 = f32[1,41,4096]{2,1,0} broadcast(%mul.4088), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention_layernorm/mul" stack_frame_id=564} + %mul.4090 = f32[1,41,4096]{2,1,0} multiply(%mul.4086, %mul.4089), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention_layernorm/mul" stack_frame_id=564} + %convert_element_type.2184 = bf16[1,41,4096]{2,1,0} convert(%mul.4090), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention_layernorm/convert_element_type" stack_frame_id=568} + %state_1__255_.1 = bf16[4096,8,128]{2,1,0} parameter(258), metadata={op_name="state[1][255]"} + %dot_general.1360 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.2184, %state_1__255_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=647} + %iota.741 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/iota" stack_frame_id=677} + %broadcast_in_dim.2004 = f32[1,41]{1,0} reshape(%iota.741), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/broadcast_in_dim" stack_frame_id=680} + %iota.740 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/iota" stack_frame_id=667} + %mul.4100 = f32[64]{0} multiply(%iota.740, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=667} + %div.1252 = f32[64]{0} divide(%mul.4100, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/div" stack_frame_id=668} + %neg.504 = f32[64]{0} negate(%div.1252), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/neg" stack_frame_id=669} + %pow.498 = f32[64]{0} power(%broadcast.51, %neg.504), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/pow" stack_frame_id=672} + %div.1253 = f32[64]{0} divide(%pow.498, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/div" stack_frame_id=673} + %dot_general.1362 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.2004, %div.1253), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/bi,j->bij/dot_general" stack_frame_id=683} + %stack.1959 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1362), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/stack" stack_frame_id=686} + %stack.1960 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1362), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/stack" stack_frame_id=686} + %stack.1961 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1959, %stack.1960), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/stack" stack_frame_id=686} + %reshape.924 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1961), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/reshape"} + %cos.248 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.924), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/cos" stack_frame_id=690} + %convert_element_type.2187 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.248), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/convert_element_type" stack_frame_id=694} + %mul.4101 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2187), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=714} + %mul.4102 = bf16[1,41,128]{2,1,0} reshape(%mul.4101), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=714} + %mul.4103 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.4102), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=714} + %mul.4104 = bf16[1,41,8,128]{3,2,1,0} multiply(%dot_general.1360, %mul.4103), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=714} + %split.497 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1360), slice={[0:1], [0:41], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/split" stack_frame_id=706} + %neg.505 = bf16[1,41,8,64]{3,2,1,0} negate(%split.497), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/neg" stack_frame_id=707} + %stack.1962 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%neg.505), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/stack" stack_frame_id=710} + %split.496 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1360), slice={[0:1], [0:41], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/split" stack_frame_id=706} + %stack.1963 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%split.496), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/stack" stack_frame_id=710} + %stack.1964 = bf16[1,41,8,2,64]{4,3,2,1,0} concatenate(%stack.1962, %stack.1963), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/stack" stack_frame_id=710} + %reshape.925 = bf16[1,41,8,128]{3,2,1,0} reshape(%stack.1964), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/reshape" stack_frame_id=713} + %sin.248 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.924), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/sin" stack_frame_id=698} + %convert_element_type.2188 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.248), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/convert_element_type" stack_frame_id=702} + %mul.4105 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2188), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=715} + %mul.4106 = bf16[1,41,128]{2,1,0} reshape(%mul.4105), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=715} + %mul.4107 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.4106), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=715} + %mul.4108 = bf16[1,41,8,128]{3,2,1,0} multiply(%reshape.925, %mul.4107), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=715} + %add.1322 = bf16[1,41,8,128]{3,2,1,0} add(%mul.4104, %mul.4108), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/add" stack_frame_id=716} + %stack.1965 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%add.1322), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/stack" stack_frame_id=719} + %state_1__256_.1 = bf16[4096,8,128]{2,1,0} parameter(259), metadata={op_name="state[1][256]"} + %dot_general.1361 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.2184, %state_1__256_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=656} + %stack.1966 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dot_general.1361), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/stack" stack_frame_id=719} + %stack.1967 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1965, %stack.1966), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/stack" stack_frame_id=719} + %stack.2035 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1967), metadata={op_name="jit(compiled_generate_function)/stack" stack_frame_id=804} + %broadcast_in_dim.2006 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dot_general.1361), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/broadcast_in_dim" stack_frame_id=725} + %reshape.927 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.2006), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/reshape" stack_frame_id=725} + %iota.734 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/iota" stack_frame_id=525} + %broadcast_in_dim.1995 = f32[41,1]{1,0} reshape(%iota.734), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/broadcast_in_dim" stack_frame_id=528} + %ge.677 = f32[41,1]{1,0} broadcast(%broadcast_in_dim.1995), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/ge" stack_frame_id=532} + %ge.678 = f32[41]{0} reshape(%ge.677), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/ge" stack_frame_id=532} + %ge.679 = f32[41,41]{1,0} broadcast(%ge.678), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/ge" stack_frame_id=532} + %iota.735 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/iota" stack_frame_id=531} + %broadcast_in_dim.1996 = f32[1,41]{1,0} reshape(%iota.735), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/broadcast_in_dim" stack_frame_id=532} + %ge.680 = f32[1,41]{1,0} broadcast(%broadcast_in_dim.1996), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/ge" stack_frame_id=532} + %ge.681 = f32[41]{0} reshape(%ge.680), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/ge" stack_frame_id=532} + %ge.682 = f32[41,41]{1,0} broadcast(%ge.681), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/ge" stack_frame_id=532} + %ge.683 = pred[41,41]{1,0} compare(%ge.679, %ge.682), direction=GE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/ge" stack_frame_id=532} + %broadcast_in_dim.1997 = pred[1,41,41]{2,1,0} reshape(%ge.683), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/broadcast_in_dim" stack_frame_id=535} + %convert_element_type.2181 = s32[1,41,41]{2,1,0} convert(%broadcast_in_dim.1997), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/convert_element_type" stack_frame_id=551} + %iota.736 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/iota" stack_frame_id=538} + %broadcast_in_dim.1998 = s32[41,1]{1,0} reshape(%iota.736), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/broadcast_in_dim" stack_frame_id=536} + %lt.819 = s32[41,1]{1,0} broadcast(%broadcast_in_dim.1998), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/lt" stack_frame_id=543} + %lt.820 = s32[41]{0} reshape(%lt.819), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/lt" stack_frame_id=543} + %lt.821 = s32[41,41]{1,0} broadcast(%lt.820), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/lt" stack_frame_id=543} + %iota.737 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/iota" stack_frame_id=541} + %broadcast_in_dim.1999 = s32[1,41]{1,0} reshape(%iota.737), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/broadcast_in_dim" stack_frame_id=539} + %add.1319 = s32[1,41]{1,0} add(%broadcast_in_dim.1999, %broadcast.57), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/add" stack_frame_id=542} + %lt.822 = s32[1,41]{1,0} broadcast(%add.1319), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/lt" stack_frame_id=543} + %lt.823 = s32[41]{0} reshape(%lt.822), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/lt" stack_frame_id=543} + %lt.824 = s32[41,41]{1,0} broadcast(%lt.823), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/lt" stack_frame_id=543} + %lt.825 = pred[41,41]{1,0} compare(%lt.821, %lt.824), direction=LT, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/lt" stack_frame_id=543} + %convert_element_type.2180 = s32[41,41]{1,0} convert(%lt.825), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/convert_element_type" stack_frame_id=547} + %broadcast_in_dim.2000 = s32[1,41,41]{2,1,0} reshape(%convert_element_type.2180), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/broadcast_in_dim" stack_frame_id=551} + %min.123 = s32[1,41,41]{2,1,0} minimum(%convert_element_type.2181, %broadcast_in_dim.2000), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/min" stack_frame_id=551} + %broadcast_in_dim.2007 = s32[1,1,41,41]{3,2,1,0} reshape(%min.123), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/broadcast_in_dim" stack_frame_id=728} + %convert_element_type.2189 = pred[1,1,41,41]{3,2,1,0} compare(%broadcast_in_dim.2007, %broadcast.49), direction=NE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/convert_element_type" stack_frame_id=732} + %broadcast_in_dim.2008 = pred[1,1,1,41,41]{4,3,2,1,0} reshape(%convert_element_type.2189), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/broadcast_in_dim" stack_frame_id=34} + %and.382 = pred[1,1,1,41,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.2008), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/vmap()/and" stack_frame_id=34} + %and.383 = pred[1,1,41,41]{3,2,1,0} reshape(%and.382), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/vmap()/and" stack_frame_id=34} + %and.384 = pred[1,1,32,41,41]{4,3,2,1,0} broadcast(%and.383), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/vmap()/and" stack_frame_id=34} + %state_1__254_.1 = bf16[4096,32,128]{2,1,0} parameter(257), metadata={op_name="state[1][254]"} + %dot_general.1358 = bf16[1,41,32,128]{3,2,1,0} dot(%convert_element_type.2184, %state_1__254_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=577} + %iota.739 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/iota" stack_frame_id=598} + %broadcast_in_dim.2003 = f32[1,41]{1,0} reshape(%iota.739), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/broadcast_in_dim" stack_frame_id=601} + %iota.738 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/iota" stack_frame_id=588} + %mul.4091 = f32[64]{0} multiply(%iota.738, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=588} + %div.1250 = f32[64]{0} divide(%mul.4091, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/div" stack_frame_id=589} + %neg.502 = f32[64]{0} negate(%div.1250), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/neg" stack_frame_id=590} + %pow.497 = f32[64]{0} power(%broadcast.51, %neg.502), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/pow" stack_frame_id=593} + %div.1251 = f32[64]{0} divide(%pow.497, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/div" stack_frame_id=594} + %dot_general.1359 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.2003, %div.1251), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/bi,j->bij/dot_general" stack_frame_id=604} + %stack.1953 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1359), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/stack" stack_frame_id=607} + %stack.1954 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1359), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/stack" stack_frame_id=607} + %stack.1955 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1953, %stack.1954), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/stack" stack_frame_id=607} + %reshape.922 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1955), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/reshape"} + %cos.247 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.922), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/cos" stack_frame_id=611} + %convert_element_type.2185 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.247), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/convert_element_type" stack_frame_id=615} + %mul.4092 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2185), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=635} + %mul.4093 = bf16[1,41,128]{2,1,0} reshape(%mul.4092), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=635} + %mul.4094 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.4093), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=635} + %mul.4095 = bf16[1,41,32,128]{3,2,1,0} multiply(%dot_general.1358, %mul.4094), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=635} + %split.495 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1358), slice={[0:1], [0:41], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/split" stack_frame_id=627} + %neg.503 = bf16[1,41,32,64]{3,2,1,0} negate(%split.495), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/neg" stack_frame_id=628} + %stack.1956 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%neg.503), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/stack" stack_frame_id=631} + %split.494 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1358), slice={[0:1], [0:41], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/split" stack_frame_id=627} + %stack.1957 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%split.494), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/stack" stack_frame_id=631} + %stack.1958 = bf16[1,41,32,2,64]{4,3,2,1,0} concatenate(%stack.1956, %stack.1957), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/stack" stack_frame_id=631} + %reshape.923 = bf16[1,41,32,128]{3,2,1,0} reshape(%stack.1958), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/reshape" stack_frame_id=634} + %sin.247 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.922), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/sin" stack_frame_id=619} + %convert_element_type.2186 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.247), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/convert_element_type" stack_frame_id=623} + %mul.4096 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2186), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=636} + %mul.4097 = bf16[1,41,128]{2,1,0} reshape(%mul.4096), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=636} + %mul.4098 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.4097), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=636} + %mul.4099 = bf16[1,41,32,128]{3,2,1,0} multiply(%reshape.923, %mul.4098), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/mul" stack_frame_id=636} + %add.1321 = bf16[1,41,32,128]{3,2,1,0} add(%mul.4095, %mul.4099), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/rotary_embedding_28/add" stack_frame_id=637} + %reshape.928 = bf16[1,41,32,1,128]{4,3,2,1,0} reshape(%add.1321), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/reshape" stack_frame_id=34} + %broadcast_in_dim.2005 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%add.1322), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/broadcast_in_dim" stack_frame_id=722} + %reshape.926 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.2005), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/reshape" stack_frame_id=722} + %dot_general.1363 = f32[1,32,41,1,41]{4,3,2,1,0} dot(%reshape.928, %reshape.926), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=34} + %mul.4109 = f32[1,32,41,1,41]{4,3,2,1,0} multiply(%dot_general.1363, %broadcast.48), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/vmap()/mul" stack_frame_id=34} + %vmap_jit__where__.123 = f32[1,1,32,41,41]{4,3,2,1,0} call(%and.384, %mul.4109, %constant.153), to_apply=%_where_0.5, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/vmap(jit(_where))" stack_frame_id=34} + %reduce_max.509 = f32[1,1,32,41]{3,2,1,0} reduce(%vmap_jit__where__.123, %constant.152), dimensions={4}, to_apply=%region_114.119, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/vmap()/reduce_max" stack_frame_id=34} + %max.133 = f32[1,1,32,41]{3,2,1,0} maximum(%reduce_max.509, %broadcast.47), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/vmap()/max" stack_frame_id=34} + %broadcast_in_dim.2009 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%max.133), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %sub.516 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.2009), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/vmap()/sub" stack_frame_id=34} + %sub.517 = f32[1,1,32,41]{3,2,1,0} reshape(%sub.516), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/vmap()/sub" stack_frame_id=34} + %sub.518 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%sub.517), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/vmap()/sub" stack_frame_id=34} + %sub.519 = f32[1,1,32,41,41]{4,3,2,1,0} subtract(%vmap_jit__where__.123, %sub.518), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/vmap()/sub" stack_frame_id=34} + %exp.129 = f32[1,1,32,41,41]{4,3,2,1,0} exponential(%sub.519), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/vmap()/exp" stack_frame_id=34} + %reduce_sum.1534 = f32[1,1,32,41]{3,2,1,0} reduce(%exp.129, %constant.155), dimensions={4}, to_apply=%region_115.120, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/vmap()/reduce_sum" stack_frame_id=34} + %broadcast_in_dim.2010 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%reduce_sum.1534), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %div.1254 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.2010), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/vmap()/div" stack_frame_id=34} + %div.1255 = f32[1,1,32,41]{3,2,1,0} reshape(%div.1254), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/vmap()/div" stack_frame_id=34} + %div.1256 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%div.1255), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/vmap()/div" stack_frame_id=34} + %div.1257 = f32[1,1,32,41,41]{4,3,2,1,0} divide(%exp.129, %div.1256), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/vmap()/div" stack_frame_id=34} + %convert_element_type.2190 = bf16[1,1,32,41,41]{4,3,2,1,0} convert(%div.1257), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/vmap()/convert_element_type" stack_frame_id=34} + %dot_general.1364 = bf16[1,32,128,1,41]{4,3,2,1,0} dot(%reshape.927, %convert_element_type.2190), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=34} + %transpose.127 = bf16[1,41,32,1,128]{1,3,4,2,0} transpose(%dot_general.1364), dimensions={0,4,1,3,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/vmap()/transpose;jit(compiled_generate_function)/transformer_layer_28/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %reshape.929 = bf16[1,41,32,128]{3,2,1,0} reshape(%transpose.127), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/reshape" stack_frame_id=34} + %state_1__257_.1 = bf16[32,128,4096]{2,1,0} parameter(260), metadata={op_name="state[1][257]"} + %dot_general.1365 = bf16[1,41,4096]{2,1,0} dot(%reshape.929, %state_1__257_.1), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=741} + %add.1323 = bf16[1,41,4096]{2,1,0} add(%dot_general.1365, %add.1318), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/add" stack_frame_id=742} + %convert_element_type.2191 = f32[1,41,4096]{2,1,0} convert(%add.1323), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/feedforward_layernorm/convert_element_type" stack_frame_id=746} + %pow.499 = f32[1,41,4096]{2,1,0} power(%convert_element_type.2191, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/feedforward_layernorm/pow" stack_frame_id=749} + %reduce_sum.1535 = f32[1,41]{1,0} reduce(%pow.499, %constant.155), dimensions={2}, to_apply=%region_116.121, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/feedforward_layernorm/reduce_sum" stack_frame_id=43} + %broadcast_in_dim.2011 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1535), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/feedforward_layernorm/broadcast_in_dim" stack_frame_id=43} + %div.1258 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.2011, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/feedforward_layernorm/div" stack_frame_id=43} + %add.1324 = f32[1,41,1]{2,1,0} add(%div.1258, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/feedforward_layernorm/add" stack_frame_id=750} + %rsqrt.250 = f32[1,41,1]{2,1,0} rsqrt(%add.1324), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/feedforward_layernorm/rsqrt" stack_frame_id=753} + %mul.4110 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.250), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/feedforward_layernorm/mul" stack_frame_id=754} + %mul.4111 = f32[1,41]{1,0} reshape(%mul.4110), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/feedforward_layernorm/mul" stack_frame_id=754} + %mul.4112 = f32[1,41,4096]{2,1,0} broadcast(%mul.4111), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/feedforward_layernorm/mul" stack_frame_id=754} + %mul.4113 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.2191, %mul.4112), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/feedforward_layernorm/mul" stack_frame_id=754} + %state_1__262_.1 = bf16[4096]{0} parameter(265), metadata={op_name="state[1][262]"} + %convert_element_type.2192 = f32[4096]{0} convert(%state_1__262_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/feedforward_layernorm/convert_element_type" stack_frame_id=755} + %broadcast_in_dim.2012 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.2192), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/feedforward_layernorm/broadcast_in_dim" stack_frame_id=755} + %mul.4114 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.2012), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/feedforward_layernorm/mul" stack_frame_id=755} + %mul.4115 = f32[1,4096]{1,0} reshape(%mul.4114), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/feedforward_layernorm/mul" stack_frame_id=755} + %mul.4116 = f32[1,41,4096]{2,1,0} broadcast(%mul.4115), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/feedforward_layernorm/mul" stack_frame_id=755} + %mul.4117 = f32[1,41,4096]{2,1,0} multiply(%mul.4113, %mul.4116), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/feedforward_layernorm/mul" stack_frame_id=755} + %convert_element_type.2193 = bf16[1,41,4096]{2,1,0} convert(%mul.4117), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/feedforward_layernorm/convert_element_type" stack_frame_id=759} + %state_1__259_.1 = bf16[4096,14336]{1,0} parameter(262), metadata={op_name="state[1][259]"} + %dot_general.1367 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.2193, %state_1__259_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/feedforward_intermediate_dense/dot_general" stack_frame_id=788} + %state_1__260_.1 = bf16[4096,14336]{1,0} parameter(263), metadata={op_name="state[1][260]"} + %dot_general.1366 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.2193, %state_1__260_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/feedforward_gate_dense/dot_general" stack_frame_id=768} + %convert_element_type.2194 = f32[1,41,14336]{2,1,0} convert(%dot_general.1366), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/convert_element_type" stack_frame_id=772} + %jit_silu_.123 = f32[1,41,14336]{2,1,0} call(%convert_element_type.2194), to_apply=%silu.9, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/jit(silu)" stack_frame_id=775} + %convert_element_type.2195 = bf16[1,41,14336]{2,1,0} convert(%jit_silu_.123), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/convert_element_type" stack_frame_id=779} + %mul.4118 = bf16[1,41,14336]{2,1,0} multiply(%dot_general.1367, %convert_element_type.2195), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/mul" stack_frame_id=791} + %state_1__261_.1 = bf16[14336,4096]{1,0} parameter(264), metadata={op_name="state[1][261]"} + %dot_general.1368 = bf16[1,41,4096]{2,1,0} dot(%mul.4118, %state_1__261_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/feedforward_output_dense/dot_general" stack_frame_id=800} + %add.1325 = bf16[1,41,4096]{2,1,0} add(%dot_general.1368, %add.1323), metadata={op_name="jit(compiled_generate_function)/transformer_layer_28/add" stack_frame_id=801} + %convert_element_type.2198 = f32[1,41,4096]{2,1,0} convert(%add.1325), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention_layernorm/convert_element_type" stack_frame_id=555} + %pow.500 = f32[1,41,4096]{2,1,0} power(%convert_element_type.2198, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention_layernorm/pow" stack_frame_id=558} + %reduce_sum.1536 = f32[1,41]{1,0} reduce(%pow.500, %constant.155), dimensions={2}, to_apply=%region_117.122, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention_layernorm/reduce_sum" stack_frame_id=24} + %broadcast_in_dim.2019 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1536), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention_layernorm/broadcast_in_dim" stack_frame_id=24} + %div.1259 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.2019, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention_layernorm/div" stack_frame_id=24} + %add.1327 = f32[1,41,1]{2,1,0} add(%div.1259, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention_layernorm/add" stack_frame_id=559} + %rsqrt.251 = f32[1,41,1]{2,1,0} rsqrt(%add.1327), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention_layernorm/rsqrt" stack_frame_id=562} + %mul.4119 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.251), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention_layernorm/mul" stack_frame_id=563} + %mul.4120 = f32[1,41]{1,0} reshape(%mul.4119), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention_layernorm/mul" stack_frame_id=563} + %mul.4121 = f32[1,41,4096]{2,1,0} broadcast(%mul.4120), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention_layernorm/mul" stack_frame_id=563} + %mul.4122 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.2198, %mul.4121), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention_layernorm/mul" stack_frame_id=563} + %state_1__267_.1 = bf16[4096]{0} parameter(270), metadata={op_name="state[1][267]"} + %convert_element_type.2199 = f32[4096]{0} convert(%state_1__267_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention_layernorm/convert_element_type" stack_frame_id=564} + %broadcast_in_dim.2020 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.2199), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention_layernorm/broadcast_in_dim" stack_frame_id=564} + %mul.4123 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.2020), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention_layernorm/mul" stack_frame_id=564} + %mul.4124 = f32[1,4096]{1,0} reshape(%mul.4123), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention_layernorm/mul" stack_frame_id=564} + %mul.4125 = f32[1,41,4096]{2,1,0} broadcast(%mul.4124), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention_layernorm/mul" stack_frame_id=564} + %mul.4126 = f32[1,41,4096]{2,1,0} multiply(%mul.4122, %mul.4125), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention_layernorm/mul" stack_frame_id=564} + %convert_element_type.2200 = bf16[1,41,4096]{2,1,0} convert(%mul.4126), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention_layernorm/convert_element_type" stack_frame_id=568} + %state_1__264_.1 = bf16[4096,8,128]{2,1,0} parameter(267), metadata={op_name="state[1][264]"} + %dot_general.1371 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.2200, %state_1__264_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=647} + %iota.749 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/iota" stack_frame_id=677} + %broadcast_in_dim.2022 = f32[1,41]{1,0} reshape(%iota.749), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/broadcast_in_dim" stack_frame_id=680} + %iota.748 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/iota" stack_frame_id=667} + %mul.4136 = f32[64]{0} multiply(%iota.748, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=667} + %div.1262 = f32[64]{0} divide(%mul.4136, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/div" stack_frame_id=668} + %neg.508 = f32[64]{0} negate(%div.1262), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/neg" stack_frame_id=669} + %pow.502 = f32[64]{0} power(%broadcast.51, %neg.508), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/pow" stack_frame_id=672} + %div.1263 = f32[64]{0} divide(%pow.502, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/div" stack_frame_id=673} + %dot_general.1373 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.2022, %div.1263), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/bi,j->bij/dot_general" stack_frame_id=683} + %stack.1974 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1373), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/stack" stack_frame_id=686} + %stack.1975 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1373), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/stack" stack_frame_id=686} + %stack.1976 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1974, %stack.1975), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/stack" stack_frame_id=686} + %reshape.932 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1976), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/reshape"} + %cos.250 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.932), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/cos" stack_frame_id=690} + %convert_element_type.2203 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.250), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/convert_element_type" stack_frame_id=694} + %mul.4137 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2203), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=714} + %mul.4138 = bf16[1,41,128]{2,1,0} reshape(%mul.4137), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=714} + %mul.4139 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.4138), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=714} + %mul.4140 = bf16[1,41,8,128]{3,2,1,0} multiply(%dot_general.1371, %mul.4139), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=714} + %split.501 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1371), slice={[0:1], [0:41], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/split" stack_frame_id=706} + %neg.509 = bf16[1,41,8,64]{3,2,1,0} negate(%split.501), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/neg" stack_frame_id=707} + %stack.1977 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%neg.509), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/stack" stack_frame_id=710} + %split.500 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1371), slice={[0:1], [0:41], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/split" stack_frame_id=706} + %stack.1978 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%split.500), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/stack" stack_frame_id=710} + %stack.1979 = bf16[1,41,8,2,64]{4,3,2,1,0} concatenate(%stack.1977, %stack.1978), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/stack" stack_frame_id=710} + %reshape.933 = bf16[1,41,8,128]{3,2,1,0} reshape(%stack.1979), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/reshape" stack_frame_id=713} + %sin.250 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.932), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/sin" stack_frame_id=698} + %convert_element_type.2204 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.250), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/convert_element_type" stack_frame_id=702} + %mul.4141 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2204), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=715} + %mul.4142 = bf16[1,41,128]{2,1,0} reshape(%mul.4141), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=715} + %mul.4143 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.4142), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=715} + %mul.4144 = bf16[1,41,8,128]{3,2,1,0} multiply(%reshape.933, %mul.4143), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=715} + %add.1329 = bf16[1,41,8,128]{3,2,1,0} add(%mul.4140, %mul.4144), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/add" stack_frame_id=716} + %stack.1980 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%add.1329), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/stack" stack_frame_id=719} + %state_1__265_.1 = bf16[4096,8,128]{2,1,0} parameter(268), metadata={op_name="state[1][265]"} + %dot_general.1372 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.2200, %state_1__265_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=656} + %stack.1981 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dot_general.1372), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/stack" stack_frame_id=719} + %stack.1982 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1980, %stack.1981), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/stack" stack_frame_id=719} + %stack.2036 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1982), metadata={op_name="jit(compiled_generate_function)/stack" stack_frame_id=804} + %broadcast_in_dim.2024 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dot_general.1372), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/broadcast_in_dim" stack_frame_id=725} + %reshape.935 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.2024), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/reshape" stack_frame_id=725} + %iota.742 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/iota" stack_frame_id=525} + %broadcast_in_dim.2013 = f32[41,1]{1,0} reshape(%iota.742), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/broadcast_in_dim" stack_frame_id=528} + %ge.684 = f32[41,1]{1,0} broadcast(%broadcast_in_dim.2013), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/ge" stack_frame_id=532} + %ge.685 = f32[41]{0} reshape(%ge.684), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/ge" stack_frame_id=532} + %ge.686 = f32[41,41]{1,0} broadcast(%ge.685), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/ge" stack_frame_id=532} + %iota.743 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/iota" stack_frame_id=531} + %broadcast_in_dim.2014 = f32[1,41]{1,0} reshape(%iota.743), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/broadcast_in_dim" stack_frame_id=532} + %ge.687 = f32[1,41]{1,0} broadcast(%broadcast_in_dim.2014), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/ge" stack_frame_id=532} + %ge.688 = f32[41]{0} reshape(%ge.687), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/ge" stack_frame_id=532} + %ge.689 = f32[41,41]{1,0} broadcast(%ge.688), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/ge" stack_frame_id=532} + %ge.690 = pred[41,41]{1,0} compare(%ge.686, %ge.689), direction=GE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/ge" stack_frame_id=532} + %broadcast_in_dim.2015 = pred[1,41,41]{2,1,0} reshape(%ge.690), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/broadcast_in_dim" stack_frame_id=535} + %convert_element_type.2197 = s32[1,41,41]{2,1,0} convert(%broadcast_in_dim.2015), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/convert_element_type" stack_frame_id=551} + %iota.744 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/iota" stack_frame_id=538} + %broadcast_in_dim.2016 = s32[41,1]{1,0} reshape(%iota.744), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/broadcast_in_dim" stack_frame_id=536} + %lt.826 = s32[41,1]{1,0} broadcast(%broadcast_in_dim.2016), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/lt" stack_frame_id=543} + %lt.827 = s32[41]{0} reshape(%lt.826), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/lt" stack_frame_id=543} + %lt.828 = s32[41,41]{1,0} broadcast(%lt.827), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/lt" stack_frame_id=543} + %iota.745 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/iota" stack_frame_id=541} + %broadcast_in_dim.2017 = s32[1,41]{1,0} reshape(%iota.745), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/broadcast_in_dim" stack_frame_id=539} + %add.1326 = s32[1,41]{1,0} add(%broadcast_in_dim.2017, %broadcast.57), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/add" stack_frame_id=542} + %lt.829 = s32[1,41]{1,0} broadcast(%add.1326), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/lt" stack_frame_id=543} + %lt.830 = s32[41]{0} reshape(%lt.829), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/lt" stack_frame_id=543} + %lt.831 = s32[41,41]{1,0} broadcast(%lt.830), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/lt" stack_frame_id=543} + %lt.832 = pred[41,41]{1,0} compare(%lt.828, %lt.831), direction=LT, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/lt" stack_frame_id=543} + %convert_element_type.2196 = s32[41,41]{1,0} convert(%lt.832), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/convert_element_type" stack_frame_id=547} + %broadcast_in_dim.2018 = s32[1,41,41]{2,1,0} reshape(%convert_element_type.2196), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/broadcast_in_dim" stack_frame_id=551} + %min.124 = s32[1,41,41]{2,1,0} minimum(%convert_element_type.2197, %broadcast_in_dim.2018), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/min" stack_frame_id=551} + %broadcast_in_dim.2025 = s32[1,1,41,41]{3,2,1,0} reshape(%min.124), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/broadcast_in_dim" stack_frame_id=728} + %convert_element_type.2205 = pred[1,1,41,41]{3,2,1,0} compare(%broadcast_in_dim.2025, %broadcast.49), direction=NE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/convert_element_type" stack_frame_id=732} + %broadcast_in_dim.2026 = pred[1,1,1,41,41]{4,3,2,1,0} reshape(%convert_element_type.2205), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/broadcast_in_dim" stack_frame_id=34} + %and.385 = pred[1,1,1,41,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.2026), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/vmap()/and" stack_frame_id=34} + %and.386 = pred[1,1,41,41]{3,2,1,0} reshape(%and.385), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/vmap()/and" stack_frame_id=34} + %and.387 = pred[1,1,32,41,41]{4,3,2,1,0} broadcast(%and.386), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/vmap()/and" stack_frame_id=34} + %state_1__263_.1 = bf16[4096,32,128]{2,1,0} parameter(266), metadata={op_name="state[1][263]"} + %dot_general.1369 = bf16[1,41,32,128]{3,2,1,0} dot(%convert_element_type.2200, %state_1__263_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=577} + %iota.747 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/iota" stack_frame_id=598} + %broadcast_in_dim.2021 = f32[1,41]{1,0} reshape(%iota.747), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/broadcast_in_dim" stack_frame_id=601} + %iota.746 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/iota" stack_frame_id=588} + %mul.4127 = f32[64]{0} multiply(%iota.746, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=588} + %div.1260 = f32[64]{0} divide(%mul.4127, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/div" stack_frame_id=589} + %neg.506 = f32[64]{0} negate(%div.1260), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/neg" stack_frame_id=590} + %pow.501 = f32[64]{0} power(%broadcast.51, %neg.506), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/pow" stack_frame_id=593} + %div.1261 = f32[64]{0} divide(%pow.501, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/div" stack_frame_id=594} + %dot_general.1370 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.2021, %div.1261), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/bi,j->bij/dot_general" stack_frame_id=604} + %stack.1968 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1370), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/stack" stack_frame_id=607} + %stack.1969 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1370), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/stack" stack_frame_id=607} + %stack.1970 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1968, %stack.1969), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/stack" stack_frame_id=607} + %reshape.930 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1970), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/reshape"} + %cos.249 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.930), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/cos" stack_frame_id=611} + %convert_element_type.2201 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.249), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/convert_element_type" stack_frame_id=615} + %mul.4128 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2201), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=635} + %mul.4129 = bf16[1,41,128]{2,1,0} reshape(%mul.4128), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=635} + %mul.4130 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.4129), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=635} + %mul.4131 = bf16[1,41,32,128]{3,2,1,0} multiply(%dot_general.1369, %mul.4130), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=635} + %split.499 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1369), slice={[0:1], [0:41], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/split" stack_frame_id=627} + %neg.507 = bf16[1,41,32,64]{3,2,1,0} negate(%split.499), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/neg" stack_frame_id=628} + %stack.1971 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%neg.507), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/stack" stack_frame_id=631} + %split.498 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1369), slice={[0:1], [0:41], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/split" stack_frame_id=627} + %stack.1972 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%split.498), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/stack" stack_frame_id=631} + %stack.1973 = bf16[1,41,32,2,64]{4,3,2,1,0} concatenate(%stack.1971, %stack.1972), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/stack" stack_frame_id=631} + %reshape.931 = bf16[1,41,32,128]{3,2,1,0} reshape(%stack.1973), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/reshape" stack_frame_id=634} + %sin.249 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.930), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/sin" stack_frame_id=619} + %convert_element_type.2202 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.249), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/convert_element_type" stack_frame_id=623} + %mul.4132 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2202), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=636} + %mul.4133 = bf16[1,41,128]{2,1,0} reshape(%mul.4132), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=636} + %mul.4134 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.4133), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=636} + %mul.4135 = bf16[1,41,32,128]{3,2,1,0} multiply(%reshape.931, %mul.4134), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/mul" stack_frame_id=636} + %add.1328 = bf16[1,41,32,128]{3,2,1,0} add(%mul.4131, %mul.4135), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/rotary_embedding_29/add" stack_frame_id=637} + %reshape.936 = bf16[1,41,32,1,128]{4,3,2,1,0} reshape(%add.1328), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/reshape" stack_frame_id=34} + %broadcast_in_dim.2023 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%add.1329), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/broadcast_in_dim" stack_frame_id=722} + %reshape.934 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.2023), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/reshape" stack_frame_id=722} + %dot_general.1374 = f32[1,32,41,1,41]{4,3,2,1,0} dot(%reshape.936, %reshape.934), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=34} + %mul.4145 = f32[1,32,41,1,41]{4,3,2,1,0} multiply(%dot_general.1374, %broadcast.48), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/vmap()/mul" stack_frame_id=34} + %vmap_jit__where__.124 = f32[1,1,32,41,41]{4,3,2,1,0} call(%and.387, %mul.4145, %constant.153), to_apply=%_where_0.5, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/vmap(jit(_where))" stack_frame_id=34} + %reduce_max.510 = f32[1,1,32,41]{3,2,1,0} reduce(%vmap_jit__where__.124, %constant.152), dimensions={4}, to_apply=%region_118.123, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/vmap()/reduce_max" stack_frame_id=34} + %max.134 = f32[1,1,32,41]{3,2,1,0} maximum(%reduce_max.510, %broadcast.47), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/vmap()/max" stack_frame_id=34} + %broadcast_in_dim.2027 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%max.134), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %sub.520 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.2027), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/vmap()/sub" stack_frame_id=34} + %sub.521 = f32[1,1,32,41]{3,2,1,0} reshape(%sub.520), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/vmap()/sub" stack_frame_id=34} + %sub.522 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%sub.521), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/vmap()/sub" stack_frame_id=34} + %sub.523 = f32[1,1,32,41,41]{4,3,2,1,0} subtract(%vmap_jit__where__.124, %sub.522), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/vmap()/sub" stack_frame_id=34} + %exp.130 = f32[1,1,32,41,41]{4,3,2,1,0} exponential(%sub.523), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/vmap()/exp" stack_frame_id=34} + %reduce_sum.1537 = f32[1,1,32,41]{3,2,1,0} reduce(%exp.130, %constant.155), dimensions={4}, to_apply=%region_119.124, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/vmap()/reduce_sum" stack_frame_id=34} + %broadcast_in_dim.2028 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%reduce_sum.1537), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %div.1264 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.2028), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/vmap()/div" stack_frame_id=34} + %div.1265 = f32[1,1,32,41]{3,2,1,0} reshape(%div.1264), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/vmap()/div" stack_frame_id=34} + %div.1266 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%div.1265), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/vmap()/div" stack_frame_id=34} + %div.1267 = f32[1,1,32,41,41]{4,3,2,1,0} divide(%exp.130, %div.1266), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/vmap()/div" stack_frame_id=34} + %convert_element_type.2206 = bf16[1,1,32,41,41]{4,3,2,1,0} convert(%div.1267), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/vmap()/convert_element_type" stack_frame_id=34} + %dot_general.1375 = bf16[1,32,128,1,41]{4,3,2,1,0} dot(%reshape.935, %convert_element_type.2206), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=34} + %transpose.128 = bf16[1,41,32,1,128]{1,3,4,2,0} transpose(%dot_general.1375), dimensions={0,4,1,3,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/vmap()/transpose;jit(compiled_generate_function)/transformer_layer_29/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %reshape.937 = bf16[1,41,32,128]{3,2,1,0} reshape(%transpose.128), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/reshape" stack_frame_id=34} + %state_1__266_.1 = bf16[32,128,4096]{2,1,0} parameter(269), metadata={op_name="state[1][266]"} + %dot_general.1376 = bf16[1,41,4096]{2,1,0} dot(%reshape.937, %state_1__266_.1), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=741} + %add.1330 = bf16[1,41,4096]{2,1,0} add(%dot_general.1376, %add.1325), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/add" stack_frame_id=742} + %convert_element_type.2207 = f32[1,41,4096]{2,1,0} convert(%add.1330), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/feedforward_layernorm/convert_element_type" stack_frame_id=746} + %pow.503 = f32[1,41,4096]{2,1,0} power(%convert_element_type.2207, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/feedforward_layernorm/pow" stack_frame_id=749} + %reduce_sum.1538 = f32[1,41]{1,0} reduce(%pow.503, %constant.155), dimensions={2}, to_apply=%region_120.125, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/feedforward_layernorm/reduce_sum" stack_frame_id=43} + %broadcast_in_dim.2029 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1538), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/feedforward_layernorm/broadcast_in_dim" stack_frame_id=43} + %div.1268 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.2029, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/feedforward_layernorm/div" stack_frame_id=43} + %add.1331 = f32[1,41,1]{2,1,0} add(%div.1268, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/feedforward_layernorm/add" stack_frame_id=750} + %rsqrt.252 = f32[1,41,1]{2,1,0} rsqrt(%add.1331), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/feedforward_layernorm/rsqrt" stack_frame_id=753} + %mul.4146 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.252), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/feedforward_layernorm/mul" stack_frame_id=754} + %mul.4147 = f32[1,41]{1,0} reshape(%mul.4146), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/feedforward_layernorm/mul" stack_frame_id=754} + %mul.4148 = f32[1,41,4096]{2,1,0} broadcast(%mul.4147), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/feedforward_layernorm/mul" stack_frame_id=754} + %mul.4149 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.2207, %mul.4148), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/feedforward_layernorm/mul" stack_frame_id=754} + %state_1__271_.1 = bf16[4096]{0} parameter(274), metadata={op_name="state[1][271]"} + %convert_element_type.2208 = f32[4096]{0} convert(%state_1__271_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/feedforward_layernorm/convert_element_type" stack_frame_id=755} + %broadcast_in_dim.2030 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.2208), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/feedforward_layernorm/broadcast_in_dim" stack_frame_id=755} + %mul.4150 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.2030), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/feedforward_layernorm/mul" stack_frame_id=755} + %mul.4151 = f32[1,4096]{1,0} reshape(%mul.4150), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/feedforward_layernorm/mul" stack_frame_id=755} + %mul.4152 = f32[1,41,4096]{2,1,0} broadcast(%mul.4151), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/feedforward_layernorm/mul" stack_frame_id=755} + %mul.4153 = f32[1,41,4096]{2,1,0} multiply(%mul.4149, %mul.4152), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/feedforward_layernorm/mul" stack_frame_id=755} + %convert_element_type.2209 = bf16[1,41,4096]{2,1,0} convert(%mul.4153), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/feedforward_layernorm/convert_element_type" stack_frame_id=759} + %state_1__268_.1 = bf16[4096,14336]{1,0} parameter(271), metadata={op_name="state[1][268]"} + %dot_general.1378 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.2209, %state_1__268_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/feedforward_intermediate_dense/dot_general" stack_frame_id=788} + %state_1__269_.1 = bf16[4096,14336]{1,0} parameter(272), metadata={op_name="state[1][269]"} + %dot_general.1377 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.2209, %state_1__269_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/feedforward_gate_dense/dot_general" stack_frame_id=768} + %convert_element_type.2210 = f32[1,41,14336]{2,1,0} convert(%dot_general.1377), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/convert_element_type" stack_frame_id=772} + %jit_silu_.124 = f32[1,41,14336]{2,1,0} call(%convert_element_type.2210), to_apply=%silu.9, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/jit(silu)" stack_frame_id=775} + %convert_element_type.2211 = bf16[1,41,14336]{2,1,0} convert(%jit_silu_.124), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/convert_element_type" stack_frame_id=779} + %mul.4154 = bf16[1,41,14336]{2,1,0} multiply(%dot_general.1378, %convert_element_type.2211), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/mul" stack_frame_id=791} + %state_1__270_.1 = bf16[14336,4096]{1,0} parameter(273), metadata={op_name="state[1][270]"} + %dot_general.1379 = bf16[1,41,4096]{2,1,0} dot(%mul.4154, %state_1__270_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/feedforward_output_dense/dot_general" stack_frame_id=800} + %add.1332 = bf16[1,41,4096]{2,1,0} add(%dot_general.1379, %add.1330), metadata={op_name="jit(compiled_generate_function)/transformer_layer_29/add" stack_frame_id=801} + %convert_element_type.2214 = f32[1,41,4096]{2,1,0} convert(%add.1332), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention_layernorm/convert_element_type" stack_frame_id=555} + %pow.504 = f32[1,41,4096]{2,1,0} power(%convert_element_type.2214, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention_layernorm/pow" stack_frame_id=558} + %reduce_sum.1539 = f32[1,41]{1,0} reduce(%pow.504, %constant.155), dimensions={2}, to_apply=%region_121.126, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention_layernorm/reduce_sum" stack_frame_id=24} + %broadcast_in_dim.2037 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1539), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention_layernorm/broadcast_in_dim" stack_frame_id=24} + %div.1269 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.2037, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention_layernorm/div" stack_frame_id=24} + %add.1334 = f32[1,41,1]{2,1,0} add(%div.1269, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention_layernorm/add" stack_frame_id=559} + %rsqrt.253 = f32[1,41,1]{2,1,0} rsqrt(%add.1334), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention_layernorm/rsqrt" stack_frame_id=562} + %mul.4155 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.253), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention_layernorm/mul" stack_frame_id=563} + %mul.4156 = f32[1,41]{1,0} reshape(%mul.4155), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention_layernorm/mul" stack_frame_id=563} + %mul.4157 = f32[1,41,4096]{2,1,0} broadcast(%mul.4156), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention_layernorm/mul" stack_frame_id=563} + %mul.4158 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.2214, %mul.4157), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention_layernorm/mul" stack_frame_id=563} + %state_1__276_.1 = bf16[4096]{0} parameter(279), metadata={op_name="state[1][276]"} + %convert_element_type.2215 = f32[4096]{0} convert(%state_1__276_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention_layernorm/convert_element_type" stack_frame_id=564} + %broadcast_in_dim.2038 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.2215), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention_layernorm/broadcast_in_dim" stack_frame_id=564} + %mul.4159 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.2038), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention_layernorm/mul" stack_frame_id=564} + %mul.4160 = f32[1,4096]{1,0} reshape(%mul.4159), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention_layernorm/mul" stack_frame_id=564} + %mul.4161 = f32[1,41,4096]{2,1,0} broadcast(%mul.4160), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention_layernorm/mul" stack_frame_id=564} + %mul.4162 = f32[1,41,4096]{2,1,0} multiply(%mul.4158, %mul.4161), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention_layernorm/mul" stack_frame_id=564} + %convert_element_type.2216 = bf16[1,41,4096]{2,1,0} convert(%mul.4162), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention_layernorm/convert_element_type" stack_frame_id=568} + %state_1__273_.1 = bf16[4096,8,128]{2,1,0} parameter(276), metadata={op_name="state[1][273]"} + %dot_general.1382 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.2216, %state_1__273_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=647} + %iota.757 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/iota" stack_frame_id=677} + %broadcast_in_dim.2040 = f32[1,41]{1,0} reshape(%iota.757), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/broadcast_in_dim" stack_frame_id=680} + %iota.756 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/iota" stack_frame_id=667} + %mul.4172 = f32[64]{0} multiply(%iota.756, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=667} + %div.1272 = f32[64]{0} divide(%mul.4172, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/div" stack_frame_id=668} + %neg.512 = f32[64]{0} negate(%div.1272), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/neg" stack_frame_id=669} + %pow.506 = f32[64]{0} power(%broadcast.51, %neg.512), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/pow" stack_frame_id=672} + %div.1273 = f32[64]{0} divide(%pow.506, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/div" stack_frame_id=673} + %dot_general.1384 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.2040, %div.1273), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/bi,j->bij/dot_general" stack_frame_id=683} + %stack.1989 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1384), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/stack" stack_frame_id=686} + %stack.1990 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1384), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/stack" stack_frame_id=686} + %stack.1991 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1989, %stack.1990), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/stack" stack_frame_id=686} + %reshape.940 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1991), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/reshape"} + %cos.252 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.940), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/cos" stack_frame_id=690} + %convert_element_type.2219 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.252), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/convert_element_type" stack_frame_id=694} + %mul.4173 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2219), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=714} + %mul.4174 = bf16[1,41,128]{2,1,0} reshape(%mul.4173), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=714} + %mul.4175 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.4174), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=714} + %mul.4176 = bf16[1,41,8,128]{3,2,1,0} multiply(%dot_general.1382, %mul.4175), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=714} + %split.505 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1382), slice={[0:1], [0:41], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/split" stack_frame_id=706} + %neg.513 = bf16[1,41,8,64]{3,2,1,0} negate(%split.505), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/neg" stack_frame_id=707} + %stack.1992 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%neg.513), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/stack" stack_frame_id=710} + %split.504 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1382), slice={[0:1], [0:41], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/split" stack_frame_id=706} + %stack.1993 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%split.504), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/stack" stack_frame_id=710} + %stack.1994 = bf16[1,41,8,2,64]{4,3,2,1,0} concatenate(%stack.1992, %stack.1993), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/stack" stack_frame_id=710} + %reshape.941 = bf16[1,41,8,128]{3,2,1,0} reshape(%stack.1994), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/reshape" stack_frame_id=713} + %sin.252 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.940), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/sin" stack_frame_id=698} + %convert_element_type.2220 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.252), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/convert_element_type" stack_frame_id=702} + %mul.4177 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2220), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=715} + %mul.4178 = bf16[1,41,128]{2,1,0} reshape(%mul.4177), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=715} + %mul.4179 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.4178), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=715} + %mul.4180 = bf16[1,41,8,128]{3,2,1,0} multiply(%reshape.941, %mul.4179), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=715} + %add.1336 = bf16[1,41,8,128]{3,2,1,0} add(%mul.4176, %mul.4180), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/add" stack_frame_id=716} + %stack.1995 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%add.1336), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/stack" stack_frame_id=719} + %state_1__274_.1 = bf16[4096,8,128]{2,1,0} parameter(277), metadata={op_name="state[1][274]"} + %dot_general.1383 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.2216, %state_1__274_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=656} + %stack.1996 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dot_general.1383), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/stack" stack_frame_id=719} + %stack.1997 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.1995, %stack.1996), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/stack" stack_frame_id=719} + %stack.2037 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.1997), metadata={op_name="jit(compiled_generate_function)/stack" stack_frame_id=804} + %broadcast_in_dim.2042 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%dot_general.1383), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/broadcast_in_dim" stack_frame_id=725} + %reshape.943 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.2042), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/reshape" stack_frame_id=725} + %iota.750 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/iota" stack_frame_id=525} + %broadcast_in_dim.2031 = f32[41,1]{1,0} reshape(%iota.750), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/broadcast_in_dim" stack_frame_id=528} + %ge.691 = f32[41,1]{1,0} broadcast(%broadcast_in_dim.2031), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/ge" stack_frame_id=532} + %ge.692 = f32[41]{0} reshape(%ge.691), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/ge" stack_frame_id=532} + %ge.693 = f32[41,41]{1,0} broadcast(%ge.692), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/ge" stack_frame_id=532} + %iota.751 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/iota" stack_frame_id=531} + %broadcast_in_dim.2032 = f32[1,41]{1,0} reshape(%iota.751), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/broadcast_in_dim" stack_frame_id=532} + %ge.694 = f32[1,41]{1,0} broadcast(%broadcast_in_dim.2032), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/ge" stack_frame_id=532} + %ge.695 = f32[41]{0} reshape(%ge.694), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/ge" stack_frame_id=532} + %ge.696 = f32[41,41]{1,0} broadcast(%ge.695), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/ge" stack_frame_id=532} + %ge.697 = pred[41,41]{1,0} compare(%ge.693, %ge.696), direction=GE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/ge" stack_frame_id=532} + %broadcast_in_dim.2033 = pred[1,41,41]{2,1,0} reshape(%ge.697), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/broadcast_in_dim" stack_frame_id=535} + %convert_element_type.2213 = s32[1,41,41]{2,1,0} convert(%broadcast_in_dim.2033), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/convert_element_type" stack_frame_id=551} + %iota.752 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/iota" stack_frame_id=538} + %broadcast_in_dim.2034 = s32[41,1]{1,0} reshape(%iota.752), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/broadcast_in_dim" stack_frame_id=536} + %lt.833 = s32[41,1]{1,0} broadcast(%broadcast_in_dim.2034), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/lt" stack_frame_id=543} + %lt.834 = s32[41]{0} reshape(%lt.833), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/lt" stack_frame_id=543} + %lt.835 = s32[41,41]{1,0} broadcast(%lt.834), dimensions={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/lt" stack_frame_id=543} + %iota.753 = s32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/iota" stack_frame_id=541} + %broadcast_in_dim.2035 = s32[1,41]{1,0} reshape(%iota.753), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/broadcast_in_dim" stack_frame_id=539} + %add.1333 = s32[1,41]{1,0} add(%broadcast_in_dim.2035, %broadcast.57), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/add" stack_frame_id=542} + %lt.836 = s32[1,41]{1,0} broadcast(%add.1333), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/lt" stack_frame_id=543} + %lt.837 = s32[41]{0} reshape(%lt.836), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/lt" stack_frame_id=543} + %lt.838 = s32[41,41]{1,0} broadcast(%lt.837), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/lt" stack_frame_id=543} + %lt.839 = pred[41,41]{1,0} compare(%lt.835, %lt.838), direction=LT, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/lt" stack_frame_id=543} + %convert_element_type.2212 = s32[41,41]{1,0} convert(%lt.839), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/convert_element_type" stack_frame_id=547} + %broadcast_in_dim.2036 = s32[1,41,41]{2,1,0} reshape(%convert_element_type.2212), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/broadcast_in_dim" stack_frame_id=551} + %min.125 = s32[1,41,41]{2,1,0} minimum(%convert_element_type.2213, %broadcast_in_dim.2036), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/min" stack_frame_id=551} + %broadcast_in_dim.2043 = s32[1,1,41,41]{3,2,1,0} reshape(%min.125), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/broadcast_in_dim" stack_frame_id=728} + %convert_element_type.2221 = pred[1,1,41,41]{3,2,1,0} compare(%broadcast_in_dim.2043, %broadcast.49), direction=NE, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/convert_element_type" stack_frame_id=732} + %broadcast_in_dim.2044 = pred[1,1,1,41,41]{4,3,2,1,0} reshape(%convert_element_type.2221), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/broadcast_in_dim" stack_frame_id=34} + %and.388 = pred[1,1,1,41,41]{4,3,2,1,0} broadcast(%broadcast_in_dim.2044), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/vmap()/and" stack_frame_id=34} + %and.389 = pred[1,1,41,41]{3,2,1,0} reshape(%and.388), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/vmap()/and" stack_frame_id=34} + %and.390 = pred[1,1,32,41,41]{4,3,2,1,0} broadcast(%and.389), dimensions={0,1,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/vmap()/and" stack_frame_id=34} + %state_1__272_.1 = bf16[4096,32,128]{2,1,0} parameter(275), metadata={op_name="state[1][272]"} + %dot_general.1380 = bf16[1,41,32,128]{3,2,1,0} dot(%convert_element_type.2216, %state_1__272_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/query/bqm,muh->bquh/dot_general" stack_frame_id=577} + %iota.755 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/iota" stack_frame_id=598} + %broadcast_in_dim.2039 = f32[1,41]{1,0} reshape(%iota.755), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/broadcast_in_dim" stack_frame_id=601} + %iota.754 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/iota" stack_frame_id=588} + %mul.4163 = f32[64]{0} multiply(%iota.754, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=588} + %div.1270 = f32[64]{0} divide(%mul.4163, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/div" stack_frame_id=589} + %neg.510 = f32[64]{0} negate(%div.1270), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/neg" stack_frame_id=590} + %pow.505 = f32[64]{0} power(%broadcast.51, %neg.510), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/pow" stack_frame_id=593} + %div.1271 = f32[64]{0} divide(%pow.505, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/div" stack_frame_id=594} + %dot_general.1381 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.2039, %div.1271), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/bi,j->bij/dot_general" stack_frame_id=604} + %stack.1983 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1381), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/stack" stack_frame_id=607} + %stack.1984 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1381), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/stack" stack_frame_id=607} + %stack.1985 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1983, %stack.1984), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/stack" stack_frame_id=607} + %reshape.938 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.1985), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/reshape"} + %cos.251 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.938), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/cos" stack_frame_id=611} + %convert_element_type.2217 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.251), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/convert_element_type" stack_frame_id=615} + %mul.4164 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2217), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=635} + %mul.4165 = bf16[1,41,128]{2,1,0} reshape(%mul.4164), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=635} + %mul.4166 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.4165), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=635} + %mul.4167 = bf16[1,41,32,128]{3,2,1,0} multiply(%dot_general.1380, %mul.4166), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=635} + %split.503 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1380), slice={[0:1], [0:41], [0:32], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/split" stack_frame_id=627} + %neg.511 = bf16[1,41,32,64]{3,2,1,0} negate(%split.503), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/neg" stack_frame_id=628} + %stack.1986 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%neg.511), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/stack" stack_frame_id=631} + %split.502 = bf16[1,41,32,64]{3,2,1,0} slice(%dot_general.1380), slice={[0:1], [0:41], [0:32], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/split" stack_frame_id=627} + %stack.1987 = bf16[1,41,32,1,64]{4,3,2,1,0} reshape(%split.502), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/stack" stack_frame_id=631} + %stack.1988 = bf16[1,41,32,2,64]{4,3,2,1,0} concatenate(%stack.1986, %stack.1987), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/stack" stack_frame_id=631} + %reshape.939 = bf16[1,41,32,128]{3,2,1,0} reshape(%stack.1988), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/reshape" stack_frame_id=634} + %sin.251 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.938), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/sin" stack_frame_id=619} + %convert_element_type.2218 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.251), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/convert_element_type" stack_frame_id=623} + %mul.4168 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2218), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=636} + %mul.4169 = bf16[1,41,128]{2,1,0} reshape(%mul.4168), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=636} + %mul.4170 = bf16[1,41,32,128]{3,2,1,0} broadcast(%mul.4169), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=636} + %mul.4171 = bf16[1,41,32,128]{3,2,1,0} multiply(%reshape.939, %mul.4170), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/mul" stack_frame_id=636} + %add.1335 = bf16[1,41,32,128]{3,2,1,0} add(%mul.4167, %mul.4171), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/rotary_embedding_30/add" stack_frame_id=637} + %reshape.944 = bf16[1,41,32,1,128]{4,3,2,1,0} reshape(%add.1335), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/reshape" stack_frame_id=34} + %broadcast_in_dim.2041 = bf16[1,41,8,4,128]{4,3,2,1,0} broadcast(%add.1336), dimensions={0,1,2,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/broadcast_in_dim" stack_frame_id=722} + %reshape.942 = bf16[1,41,32,128]{3,2,1,0} reshape(%broadcast_in_dim.2041), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/reshape" stack_frame_id=722} + %dot_general.1385 = f32[1,32,41,1,41]{4,3,2,1,0} dot(%reshape.944, %reshape.942), lhs_batch_dims={0,2}, lhs_contracting_dims={4}, rhs_batch_dims={0,2}, rhs_contracting_dims={3}, algorithm=dot_bf16_bf16_f32, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/vmap(BTNH,BSNH->BNTS)/dot_general" stack_frame_id=34} + %mul.4181 = f32[1,32,41,1,41]{4,3,2,1,0} multiply(%dot_general.1385, %broadcast.48), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/vmap()/mul" stack_frame_id=34} + %vmap_jit__where__.125 = f32[1,1,32,41,41]{4,3,2,1,0} call(%and.390, %mul.4181, %constant.153), to_apply=%_where_0.5, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/vmap(jit(_where))" stack_frame_id=34} + %reduce_max.511 = f32[1,1,32,41]{3,2,1,0} reduce(%vmap_jit__where__.125, %constant.152), dimensions={4}, to_apply=%region_122.127, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/vmap()/reduce_max" stack_frame_id=34} + %max.135 = f32[1,1,32,41]{3,2,1,0} maximum(%reduce_max.511, %broadcast.47), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/vmap()/max" stack_frame_id=34} + %broadcast_in_dim.2045 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%max.135), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %sub.524 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.2045), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/vmap()/sub" stack_frame_id=34} + %sub.525 = f32[1,1,32,41]{3,2,1,0} reshape(%sub.524), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/vmap()/sub" stack_frame_id=34} + %sub.526 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%sub.525), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/vmap()/sub" stack_frame_id=34} + %sub.527 = f32[1,1,32,41,41]{4,3,2,1,0} subtract(%vmap_jit__where__.125, %sub.526), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/vmap()/sub" stack_frame_id=34} + %exp.131 = f32[1,1,32,41,41]{4,3,2,1,0} exponential(%sub.527), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/vmap()/exp" stack_frame_id=34} + %reduce_sum.1540 = f32[1,1,32,41]{3,2,1,0} reduce(%exp.131, %constant.155), dimensions={4}, to_apply=%region_123.128, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/vmap()/reduce_sum" stack_frame_id=34} + %broadcast_in_dim.2046 = f32[1,1,32,41,1]{4,3,2,1,0} reshape(%reduce_sum.1540), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/vmap()/broadcast_in_dim" stack_frame_id=34} + %div.1274 = f32[1,1,32,41,1]{4,3,2,1,0} broadcast(%broadcast_in_dim.2046), dimensions={0,1,2,3,4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/vmap()/div" stack_frame_id=34} + %div.1275 = f32[1,1,32,41]{3,2,1,0} reshape(%div.1274), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/vmap()/div" stack_frame_id=34} + %div.1276 = f32[1,1,32,41,41]{4,3,2,1,0} broadcast(%div.1275), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/vmap()/div" stack_frame_id=34} + %div.1277 = f32[1,1,32,41,41]{4,3,2,1,0} divide(%exp.131, %div.1276), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/vmap()/div" stack_frame_id=34} + %convert_element_type.2222 = bf16[1,1,32,41,41]{4,3,2,1,0} convert(%div.1277), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/vmap()/convert_element_type" stack_frame_id=34} + %dot_general.1386 = bf16[1,32,128,1,41]{4,3,2,1,0} dot(%reshape.943, %convert_element_type.2222), lhs_batch_dims={0,2}, lhs_contracting_dims={1}, rhs_batch_dims={1,2}, rhs_contracting_dims={4}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/vmap(BNTS,BSNH->BTNH)/dot_general" stack_frame_id=34} + %transpose.129 = bf16[1,41,32,1,128]{1,3,4,2,0} transpose(%dot_general.1386), dimensions={0,4,1,3,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/vmap()/transpose;jit(compiled_generate_function)/transformer_layer_30/self_attention/vmap(BNTS,BSNH->BTNH)/transpose"} + %reshape.945 = bf16[1,41,32,128]{3,2,1,0} reshape(%transpose.129), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/reshape" stack_frame_id=34} + %state_1__275_.1 = bf16[32,128,4096]{2,1,0} parameter(278), metadata={op_name="state[1][275]"} + %dot_general.1387 = bf16[1,41,4096]{2,1,0} dot(%reshape.945, %state_1__275_.1), lhs_contracting_dims={3,2}, rhs_contracting_dims={1,0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/self_attention/attention_output/bquh,uhm->bqm/dot_general" stack_frame_id=741} + %add.1337 = bf16[1,41,4096]{2,1,0} add(%dot_general.1387, %add.1332), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/add" stack_frame_id=742} + %convert_element_type.2223 = f32[1,41,4096]{2,1,0} convert(%add.1337), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/feedforward_layernorm/convert_element_type" stack_frame_id=746} + %pow.507 = f32[1,41,4096]{2,1,0} power(%convert_element_type.2223, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/feedforward_layernorm/pow" stack_frame_id=749} + %reduce_sum.1541 = f32[1,41]{1,0} reduce(%pow.507, %constant.155), dimensions={2}, to_apply=%region_124.129, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/feedforward_layernorm/reduce_sum" stack_frame_id=43} + %broadcast_in_dim.2047 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1541), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/feedforward_layernorm/broadcast_in_dim" stack_frame_id=43} + %div.1278 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.2047, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/feedforward_layernorm/div" stack_frame_id=43} + %add.1338 = f32[1,41,1]{2,1,0} add(%div.1278, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/feedforward_layernorm/add" stack_frame_id=750} + %rsqrt.254 = f32[1,41,1]{2,1,0} rsqrt(%add.1338), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/feedforward_layernorm/rsqrt" stack_frame_id=753} + %mul.4182 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.254), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/feedforward_layernorm/mul" stack_frame_id=754} + %mul.4183 = f32[1,41]{1,0} reshape(%mul.4182), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/feedforward_layernorm/mul" stack_frame_id=754} + %mul.4184 = f32[1,41,4096]{2,1,0} broadcast(%mul.4183), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/feedforward_layernorm/mul" stack_frame_id=754} + %mul.4185 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.2223, %mul.4184), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/feedforward_layernorm/mul" stack_frame_id=754} + %state_1__280_.1 = bf16[4096]{0} parameter(283), metadata={op_name="state[1][280]"} + %convert_element_type.2224 = f32[4096]{0} convert(%state_1__280_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/feedforward_layernorm/convert_element_type" stack_frame_id=755} + %broadcast_in_dim.2048 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.2224), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/feedforward_layernorm/broadcast_in_dim" stack_frame_id=755} + %mul.4186 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.2048), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/feedforward_layernorm/mul" stack_frame_id=755} + %mul.4187 = f32[1,4096]{1,0} reshape(%mul.4186), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/feedforward_layernorm/mul" stack_frame_id=755} + %mul.4188 = f32[1,41,4096]{2,1,0} broadcast(%mul.4187), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/feedforward_layernorm/mul" stack_frame_id=755} + %mul.4189 = f32[1,41,4096]{2,1,0} multiply(%mul.4185, %mul.4188), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/feedforward_layernorm/mul" stack_frame_id=755} + %convert_element_type.2225 = bf16[1,41,4096]{2,1,0} convert(%mul.4189), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/feedforward_layernorm/convert_element_type" stack_frame_id=759} + %state_1__277_.1 = bf16[4096,14336]{1,0} parameter(280), metadata={op_name="state[1][277]"} + %dot_general.1389 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.2225, %state_1__277_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/feedforward_intermediate_dense/dot_general" stack_frame_id=788} + %state_1__278_.1 = bf16[4096,14336]{1,0} parameter(281), metadata={op_name="state[1][278]"} + %dot_general.1388 = bf16[1,41,14336]{2,1,0} dot(%convert_element_type.2225, %state_1__278_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/feedforward_gate_dense/dot_general" stack_frame_id=768} + %convert_element_type.2226 = f32[1,41,14336]{2,1,0} convert(%dot_general.1388), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/convert_element_type" stack_frame_id=772} + %jit_silu_.125 = f32[1,41,14336]{2,1,0} call(%convert_element_type.2226), to_apply=%silu.9, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/jit(silu)" stack_frame_id=775} + %convert_element_type.2227 = bf16[1,41,14336]{2,1,0} convert(%jit_silu_.125), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/convert_element_type" stack_frame_id=779} + %mul.4190 = bf16[1,41,14336]{2,1,0} multiply(%dot_general.1389, %convert_element_type.2227), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/mul" stack_frame_id=791} + %state_1__279_.1 = bf16[14336,4096]{1,0} parameter(282), metadata={op_name="state[1][279]"} + %dot_general.1390 = bf16[1,41,4096]{2,1,0} dot(%mul.4190, %state_1__279_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/feedforward_output_dense/dot_general" stack_frame_id=800} + %add.1339 = bf16[1,41,4096]{2,1,0} add(%dot_general.1390, %add.1337), metadata={op_name="jit(compiled_generate_function)/transformer_layer_30/add" stack_frame_id=801} + %convert_element_type.2228 = f32[1,41,4096]{2,1,0} convert(%add.1339), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention_layernorm/convert_element_type" stack_frame_id=555} + %pow.508 = f32[1,41,4096]{2,1,0} power(%convert_element_type.2228, %broadcast.56), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention_layernorm/pow" stack_frame_id=558} + %reduce_sum.1542 = f32[1,41]{1,0} reduce(%pow.508, %constant.155), dimensions={2}, to_apply=%region_125.130, metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention_layernorm/reduce_sum" stack_frame_id=24} + %broadcast_in_dim.2049 = f32[1,41,1]{2,1,0} reshape(%reduce_sum.1542), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention_layernorm/broadcast_in_dim" stack_frame_id=24} + %div.1279 = f32[1,41,1]{2,1,0} divide(%broadcast_in_dim.2049, %broadcast.55), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention_layernorm/div" stack_frame_id=24} + %add.1340 = f32[1,41,1]{2,1,0} add(%div.1279, %broadcast.54), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention_layernorm/add" stack_frame_id=559} + %rsqrt.255 = f32[1,41,1]{2,1,0} rsqrt(%add.1340), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention_layernorm/rsqrt" stack_frame_id=562} + %mul.4191 = f32[1,41,1]{2,1,0} broadcast(%rsqrt.255), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention_layernorm/mul" stack_frame_id=563} + %mul.4192 = f32[1,41]{1,0} reshape(%mul.4191), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention_layernorm/mul" stack_frame_id=563} + %mul.4193 = f32[1,41,4096]{2,1,0} broadcast(%mul.4192), dimensions={0,1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention_layernorm/mul" stack_frame_id=563} + %mul.4194 = f32[1,41,4096]{2,1,0} multiply(%convert_element_type.2228, %mul.4193), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention_layernorm/mul" stack_frame_id=563} + %state_1__285_.1 = bf16[4096]{0} parameter(288), metadata={op_name="state[1][285]"} + %convert_element_type.2229 = f32[4096]{0} convert(%state_1__285_.1), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention_layernorm/convert_element_type" stack_frame_id=564} + %broadcast_in_dim.2050 = f32[1,1,4096]{2,1,0} reshape(%convert_element_type.2229), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention_layernorm/broadcast_in_dim" stack_frame_id=564} + %mul.4195 = f32[1,1,4096]{2,1,0} broadcast(%broadcast_in_dim.2050), dimensions={0,1,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention_layernorm/mul" stack_frame_id=564} + %mul.4196 = f32[1,4096]{1,0} reshape(%mul.4195), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention_layernorm/mul" stack_frame_id=564} + %mul.4197 = f32[1,41,4096]{2,1,0} broadcast(%mul.4196), dimensions={0,2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention_layernorm/mul" stack_frame_id=564} + %mul.4198 = f32[1,41,4096]{2,1,0} multiply(%mul.4194, %mul.4197), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention_layernorm/mul" stack_frame_id=564} + %convert_element_type.2230 = bf16[1,41,4096]{2,1,0} convert(%mul.4198), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention_layernorm/convert_element_type" stack_frame_id=568} + %state_1__282_.1 = bf16[4096,8,128]{2,1,0} parameter(285), metadata={op_name="state[1][282]"} + %dot_general.1391 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.2230, %state_1__282_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/key/bkm,mvh->bkvh/dot_general" stack_frame_id=647} + %iota.759 = f32[41]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/rotary_embedding_31/iota" stack_frame_id=677} + %broadcast_in_dim.2051 = f32[1,41]{1,0} reshape(%iota.759), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/rotary_embedding_31/broadcast_in_dim" stack_frame_id=680} + %iota.758 = f32[64]{0} iota(), iota_dimension=0, metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/rotary_embedding_31/iota" stack_frame_id=667} + %mul.4199 = f32[64]{0} multiply(%iota.758, %broadcast.53), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/rotary_embedding_31/mul" stack_frame_id=667} + %div.1280 = f32[64]{0} divide(%mul.4199, %broadcast.52), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/rotary_embedding_31/div" stack_frame_id=668} + %neg.514 = f32[64]{0} negate(%div.1280), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/rotary_embedding_31/neg" stack_frame_id=669} + %pow.509 = f32[64]{0} power(%broadcast.51, %neg.514), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/rotary_embedding_31/pow" stack_frame_id=672} + %div.1281 = f32[64]{0} divide(%pow.509, %broadcast.50), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/rotary_embedding_31/div" stack_frame_id=673} + %dot_general.1393 = f32[1,41,64]{2,1,0} dot(%broadcast_in_dim.2051, %div.1281), lhs_contracting_dims={}, rhs_contracting_dims={}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/rotary_embedding_31/bi,j->bij/dot_general" stack_frame_id=683} + %stack.1998 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1393), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/rotary_embedding_31/stack" stack_frame_id=686} + %stack.1999 = f32[1,41,1,64]{3,2,1,0} reshape(%dot_general.1393), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/rotary_embedding_31/stack" stack_frame_id=686} + %stack.2000 = f32[1,41,2,64]{3,2,1,0} concatenate(%stack.1998, %stack.1999), dimensions={2}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/rotary_embedding_31/stack" stack_frame_id=686} + %reshape.946 = f32[1,41,1,128]{3,2,1,0} reshape(%stack.2000), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/rotary_embedding_31/broadcast_in_dim;jit(compiled_generate_function)/transformer_layer_31/self_attention/rotary_embedding_31/reshape"} + %cos.253 = f32[1,41,1,128]{3,2,1,0} cosine(%reshape.946), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/rotary_embedding_31/cos" stack_frame_id=690} + %convert_element_type.2231 = bf16[1,41,1,128]{3,2,1,0} convert(%cos.253), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/rotary_embedding_31/convert_element_type" stack_frame_id=694} + %mul.4200 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2231), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/rotary_embedding_31/mul" stack_frame_id=714} + %mul.4201 = bf16[1,41,128]{2,1,0} reshape(%mul.4200), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/rotary_embedding_31/mul" stack_frame_id=714} + %mul.4202 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.4201), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/rotary_embedding_31/mul" stack_frame_id=714} + %mul.4203 = bf16[1,41,8,128]{3,2,1,0} multiply(%dot_general.1391, %mul.4202), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/rotary_embedding_31/mul" stack_frame_id=714} + %split.507 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1391), slice={[0:1], [0:41], [0:8], [64:128]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/rotary_embedding_31/split" stack_frame_id=706} + %neg.515 = bf16[1,41,8,64]{3,2,1,0} negate(%split.507), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/rotary_embedding_31/neg" stack_frame_id=707} + %stack.2001 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%neg.515), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/rotary_embedding_31/stack" stack_frame_id=710} + %split.506 = bf16[1,41,8,64]{3,2,1,0} slice(%dot_general.1391), slice={[0:1], [0:41], [0:8], [0:64]}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/rotary_embedding_31/split" stack_frame_id=706} + %stack.2002 = bf16[1,41,8,1,64]{4,3,2,1,0} reshape(%split.506), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/rotary_embedding_31/stack" stack_frame_id=710} + %stack.2003 = bf16[1,41,8,2,64]{4,3,2,1,0} concatenate(%stack.2001, %stack.2002), dimensions={3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/rotary_embedding_31/stack" stack_frame_id=710} + %reshape.947 = bf16[1,41,8,128]{3,2,1,0} reshape(%stack.2003), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/rotary_embedding_31/reshape" stack_frame_id=713} + %sin.253 = f32[1,41,1,128]{3,2,1,0} sine(%reshape.946), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/rotary_embedding_31/sin" stack_frame_id=698} + %convert_element_type.2232 = bf16[1,41,1,128]{3,2,1,0} convert(%sin.253), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/rotary_embedding_31/convert_element_type" stack_frame_id=702} + %mul.4204 = bf16[1,41,1,128]{3,2,1,0} broadcast(%convert_element_type.2232), dimensions={0,1,2,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/rotary_embedding_31/mul" stack_frame_id=715} + %mul.4205 = bf16[1,41,128]{2,1,0} reshape(%mul.4204), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/rotary_embedding_31/mul" stack_frame_id=715} + %mul.4206 = bf16[1,41,8,128]{3,2,1,0} broadcast(%mul.4205), dimensions={0,1,3}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/rotary_embedding_31/mul" stack_frame_id=715} + %mul.4207 = bf16[1,41,8,128]{3,2,1,0} multiply(%reshape.947, %mul.4206), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/rotary_embedding_31/mul" stack_frame_id=715} + %add.1341 = bf16[1,41,8,128]{3,2,1,0} add(%mul.4203, %mul.4207), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/rotary_embedding_31/add" stack_frame_id=716} + %stack.2004 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%add.1341), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/stack" stack_frame_id=719} + %state_1__283_.1 = bf16[4096,8,128]{2,1,0} parameter(286), metadata={op_name="state[1][283]"} + %dot_general.1392 = bf16[1,41,8,128]{3,2,1,0} dot(%convert_element_type.2230, %state_1__283_.1), lhs_contracting_dims={2}, rhs_contracting_dims={0}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/value/bkm,mvh->bkvh/dot_general" stack_frame_id=656} + %stack.2005 = bf16[1,1,41,8,128]{4,3,2,1,0} reshape(%dot_general.1392), metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/stack" stack_frame_id=719} + %stack.2006 = bf16[1,2,41,8,128]{4,3,2,1,0} concatenate(%stack.2004, %stack.2005), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/transformer_layer_31/self_attention/stack" stack_frame_id=719} + %stack.2038 = bf16[1,1,2,41,8,128]{5,4,3,2,1,0} reshape(%stack.2006), metadata={op_name="jit(compiled_generate_function)/stack" stack_frame_id=804} + %stack.2039 = bf16[1,32,2,41,8,128]{5,4,3,2,1,0} concatenate(%stack.2007, %stack.2008, %stack.2009, %stack.2010, %stack.2011, /*index=5*/%stack.2012, %stack.2013, %stack.2014, %stack.2015, %stack.2016, /*index=10*/%stack.2017, %stack.2018, %stack.2019, %stack.2020, %stack.2021, /*index=15*/%stack.2022, %stack.2023, %stack.2024, %stack.2025, %stack.2026, /*index=20*/%stack.2027, %stack.2028, %stack.2029, %stack.2030, %stack.2031, /*index=25*/%stack.2032, %stack.2033, %stack.2034, %stack.2035, %stack.2036, /*index=30*/%stack.2037, %stack.2038), dimensions={1}, metadata={op_name="jit(compiled_generate_function)/stack" stack_frame_id=804} + %inputs__padding_mask__.1 = pred[1,41]{1,0} parameter(0), metadata={op_name="inputs[\'padding_mask\']"} + %convert_element_type.2233 = s32[1,41]{1,0} convert(%inputs__padding_mask__.1), metadata={op_name="jit(compiled_generate_function)/convert_element_type" stack_frame_id=808} + %constant.154 = s32[] constant(0) + %reduce_sum.1543 = s32[1]{0} reduce(%convert_element_type.2233, %constant.154), dimensions={1}, to_apply=%region_126.131, metadata={op_name="jit(compiled_generate_function)/reduce_sum" stack_frame_id=46} + %constant.151 = s32[] constant(2147483647) + %reduce_min.7 = s32[] reduce(%reduce_sum.1543, %constant.151), dimensions={0}, to_apply=%region_127.132, metadata={op_name="jit(compiled_generate_function)/reduce_min" stack_frame_id=49} + %constant.150 = s32[] constant(41) + %sub.528 = s32[] subtract(%constant.150, %reduce_min.7), metadata={op_name="jit(compiled_generate_function)/sub" stack_frame_id=809} + %state_1__281_.1 = bf16[4096,32,128]{2,1,0} parameter(284), metadata={op_name="state[1][281]"} + %state_1__284_.1 = bf16[32,128,4096]{2,1,0} parameter(287), metadata={op_name="state[1][284]"} + %state_1__289_.1 = bf16[4096]{0} parameter(292), metadata={op_name="state[1][289]"} + %state_1__287_.1 = bf16[4096,14336]{1,0} parameter(290), metadata={op_name="state[1][287]"} + %state_1__286_.1 = bf16[4096,14336]{1,0} parameter(289), metadata={op_name="state[1][286]"} + %state_1__288_.1 = bf16[14336,4096]{1,0} parameter(291), metadata={op_name="state[1][288]"} + %state_1__290_.1 = bf16[4096]{0} parameter(293), metadata={op_name="state[1][290]"} + %state_1__1_.1 = bf16[4096,32000]{1,0} parameter(4), metadata={op_name="state[1][1]"} + %while.29 = (u32[2]{0}, s32[1,41]{1,0}, bf16[1,32,2,41,8,128]{5,4,3,2,1,0}, s32[], s32[], /*index=5*/pred[1,41]{1,0}, s32[], bf16[32000,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=10*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=15*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=20*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=25*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=30*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=35*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=40*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=45*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=50*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=55*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=60*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=65*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=70*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=75*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=80*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=85*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=90*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=95*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=100*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=105*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=110*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=115*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=120*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=125*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=130*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=135*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=140*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=145*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=150*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=155*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=160*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=165*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=170*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=175*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=180*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=185*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=190*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=195*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=200*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=205*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=210*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=215*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=220*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=225*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=230*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=235*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=240*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=245*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=250*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=255*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=260*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=265*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=270*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=275*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=280*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=285*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=290*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=295*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32000]{1,0}) tuple(%state_0__0_.1, %inputs__token_ids__.1, %stack.2039, %reduce_min.7, %constant.154, /*index=5*/%inputs__padding_mask__.1, %sub.528, %state_1__0_.1, %state_1__6_.1, %state_1__2_.1, /*index=10*/%state_1__3_.1, %state_1__4_.1, %state_1__5_.1, %state_1__10_.1, %state_1__8_.1, /*index=15*/%state_1__7_.1, %state_1__9_.1, %state_1__15_.1, %state_1__11_.1, %state_1__12_.1, /*index=20*/%state_1__13_.1, %state_1__14_.1, %state_1__19_.1, %state_1__17_.1, %state_1__16_.1, /*index=25*/%state_1__18_.1, %state_1__24_.1, %state_1__20_.1, %state_1__21_.1, %state_1__22_.1, /*index=30*/%state_1__23_.1, %state_1__28_.1, %state_1__26_.1, %state_1__25_.1, %state_1__27_.1, /*index=35*/%state_1__33_.1, %state_1__29_.1, %state_1__30_.1, %state_1__31_.1, %state_1__32_.1, /*index=40*/%state_1__37_.1, %state_1__35_.1, %state_1__34_.1, %state_1__36_.1, %state_1__42_.1, /*index=45*/%state_1__38_.1, %state_1__39_.1, %state_1__40_.1, %state_1__41_.1, %state_1__46_.1, /*index=50*/%state_1__44_.1, %state_1__43_.1, %state_1__45_.1, %state_1__51_.1, %state_1__47_.1, /*index=55*/%state_1__48_.1, %state_1__49_.1, %state_1__50_.1, %state_1__55_.1, %state_1__53_.1, /*index=60*/%state_1__52_.1, %state_1__54_.1, %state_1__60_.1, %state_1__56_.1, %state_1__57_.1, /*index=65*/%state_1__58_.1, %state_1__59_.1, %state_1__64_.1, %state_1__62_.1, %state_1__61_.1, /*index=70*/%state_1__63_.1, %state_1__69_.1, %state_1__65_.1, %state_1__66_.1, %state_1__67_.1, /*index=75*/%state_1__68_.1, %state_1__73_.1, %state_1__71_.1, %state_1__70_.1, %state_1__72_.1, /*index=80*/%state_1__78_.1, %state_1__74_.1, %state_1__75_.1, %state_1__76_.1, %state_1__77_.1, /*index=85*/%state_1__82_.1, %state_1__80_.1, %state_1__79_.1, %state_1__81_.1, %state_1__87_.1, /*index=90*/%state_1__83_.1, %state_1__84_.1, %state_1__85_.1, %state_1__86_.1, %state_1__91_.1, /*index=95*/%state_1__89_.1, %state_1__88_.1, %state_1__90_.1, %state_1__96_.1, %state_1__92_.1, /*index=100*/%state_1__93_.1, %state_1__94_.1, %state_1__95_.1, %state_1__100_.1, %state_1__98_.1, /*index=105*/%state_1__97_.1, %state_1__99_.1, %state_1__105_.1, %state_1__101_.1, %state_1__102_.1, /*index=110*/%state_1__103_.1, %state_1__104_.1, %state_1__109_.1, %state_1__107_.1, %state_1__106_.1, /*index=115*/%state_1__108_.1, %state_1__114_.1, %state_1__110_.1, %state_1__111_.1, %state_1__112_.1, /*index=120*/%state_1__113_.1, %state_1__118_.1, %state_1__116_.1, %state_1__115_.1, %state_1__117_.1, /*index=125*/%state_1__123_.1, %state_1__119_.1, %state_1__120_.1, %state_1__121_.1, %state_1__122_.1, /*index=130*/%state_1__127_.1, %state_1__125_.1, %state_1__124_.1, %state_1__126_.1, %state_1__132_.1, /*index=135*/%state_1__128_.1, %state_1__129_.1, %state_1__130_.1, %state_1__131_.1, %state_1__136_.1, /*index=140*/%state_1__134_.1, %state_1__133_.1, %state_1__135_.1, %state_1__141_.1, %state_1__137_.1, /*index=145*/%state_1__138_.1, %state_1__139_.1, %state_1__140_.1, %state_1__145_.1, %state_1__143_.1, /*index=150*/%state_1__142_.1, %state_1__144_.1, %state_1__150_.1, %state_1__146_.1, %state_1__147_.1, /*index=155*/%state_1__148_.1, %state_1__149_.1, %state_1__154_.1, %state_1__152_.1, %state_1__151_.1, /*index=160*/%state_1__153_.1, %state_1__159_.1, %state_1__155_.1, %state_1__156_.1, %state_1__157_.1, /*index=165*/%state_1__158_.1, %state_1__163_.1, %state_1__161_.1, %state_1__160_.1, %state_1__162_.1, /*index=170*/%state_1__168_.1, %state_1__164_.1, %state_1__165_.1, %state_1__166_.1, %state_1__167_.1, /*index=175*/%state_1__172_.1, %state_1__170_.1, %state_1__169_.1, %state_1__171_.1, %state_1__177_.1, /*index=180*/%state_1__173_.1, %state_1__174_.1, %state_1__175_.1, %state_1__176_.1, %state_1__181_.1, /*index=185*/%state_1__179_.1, %state_1__178_.1, %state_1__180_.1, %state_1__186_.1, %state_1__182_.1, /*index=190*/%state_1__183_.1, %state_1__184_.1, %state_1__185_.1, %state_1__190_.1, %state_1__188_.1, /*index=195*/%state_1__187_.1, %state_1__189_.1, %state_1__195_.1, %state_1__191_.1, %state_1__192_.1, /*index=200*/%state_1__193_.1, %state_1__194_.1, %state_1__199_.1, %state_1__197_.1, %state_1__196_.1, /*index=205*/%state_1__198_.1, %state_1__204_.1, %state_1__200_.1, %state_1__201_.1, %state_1__202_.1, /*index=210*/%state_1__203_.1, %state_1__208_.1, %state_1__206_.1, %state_1__205_.1, %state_1__207_.1, /*index=215*/%state_1__213_.1, %state_1__209_.1, %state_1__210_.1, %state_1__211_.1, %state_1__212_.1, /*index=220*/%state_1__217_.1, %state_1__215_.1, %state_1__214_.1, %state_1__216_.1, %state_1__222_.1, /*index=225*/%state_1__218_.1, %state_1__219_.1, %state_1__220_.1, %state_1__221_.1, %state_1__226_.1, /*index=230*/%state_1__224_.1, %state_1__223_.1, %state_1__225_.1, %state_1__231_.1, %state_1__227_.1, /*index=235*/%state_1__228_.1, %state_1__229_.1, %state_1__230_.1, %state_1__235_.1, %state_1__233_.1, /*index=240*/%state_1__232_.1, %state_1__234_.1, %state_1__240_.1, %state_1__236_.1, %state_1__237_.1, /*index=245*/%state_1__238_.1, %state_1__239_.1, %state_1__244_.1, %state_1__242_.1, %state_1__241_.1, /*index=250*/%state_1__243_.1, %state_1__249_.1, %state_1__245_.1, %state_1__246_.1, %state_1__247_.1, /*index=255*/%state_1__248_.1, %state_1__253_.1, %state_1__251_.1, %state_1__250_.1, %state_1__252_.1, /*index=260*/%state_1__258_.1, %state_1__254_.1, %state_1__255_.1, %state_1__256_.1, %state_1__257_.1, /*index=265*/%state_1__262_.1, %state_1__260_.1, %state_1__259_.1, %state_1__261_.1, %state_1__267_.1, /*index=270*/%state_1__263_.1, %state_1__264_.1, %state_1__265_.1, %state_1__266_.1, %state_1__271_.1, /*index=275*/%state_1__269_.1, %state_1__268_.1, %state_1__270_.1, %state_1__276_.1, %state_1__272_.1, /*index=280*/%state_1__273_.1, %state_1__274_.1, %state_1__275_.1, %state_1__280_.1, %state_1__278_.1, /*index=285*/%state_1__277_.1, %state_1__279_.1, %state_1__285_.1, %state_1__281_.1, %state_1__282_.1, /*index=290*/%state_1__283_.1, %state_1__284_.1, %state_1__289_.1, %state_1__287_.1, %state_1__286_.1, /*index=295*/%state_1__288_.1, %state_1__290_.1, %state_1__1_.1), metadata={op_name="jit(compiled_generate_function)/while" stack_frame_id=54} + %while.30 = (u32[2]{0}, s32[1,41]{1,0}, bf16[1,32,2,41,8,128]{5,4,3,2,1,0}, s32[], s32[], /*index=5*/pred[1,41]{1,0}, s32[], bf16[32000,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=10*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=15*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=20*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=25*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=30*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=35*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=40*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=45*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=50*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=55*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=60*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=65*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=70*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=75*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=80*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=85*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=90*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=95*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=100*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=105*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=110*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=115*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=120*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=125*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=130*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=135*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=140*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=145*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=150*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=155*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=160*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=165*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=170*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=175*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=180*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=185*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=190*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=195*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=200*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=205*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=210*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=215*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=220*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=225*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=230*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=235*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=240*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=245*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=250*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=255*/bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, /*index=260*/bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, /*index=265*/bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, /*index=270*/bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, /*index=275*/bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, /*index=280*/bf16[4096,8,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, /*index=285*/bf16[4096,14336]{1,0}, bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32,128]{2,1,0}, bf16[4096,8,128]{2,1,0}, /*index=290*/bf16[4096,8,128]{2,1,0}, bf16[32,128,4096]{2,1,0}, bf16[4096]{0}, bf16[4096,14336]{1,0}, bf16[4096,14336]{1,0}, /*index=295*/bf16[14336,4096]{1,0}, bf16[4096]{0}, bf16[4096,32000]{1,0}) while(%while.29), condition=%region_265.283, body=%region_128.280, metadata={op_name="jit(compiled_generate_function)/while" stack_frame_id=54} + %constant.138 = s32[] constant(2) + %broadcast.46 = s32[1,41]{1,0} broadcast(%constant.138), dimensions={} + %eq.5 = pred[1,41]{1,0} compare(%while.30#1, %broadcast.46), direction=EQ, metadata={op_name="jit(compiled_generate_function)/eq" stack_frame_id=816} + %not.6 = pred[1,41]{1,0} not(%inputs__padding_mask__.1), metadata={op_name="jit(compiled_generate_function)/not" stack_frame_id=812} + %and.391 = pred[1,41]{1,0} and(%eq.5, %not.6), metadata={op_name="jit(compiled_generate_function)/and" stack_frame_id=819} + %convert_element_type.2234 = s32[1,41]{1,0} convert(%and.391), metadata={op_name="jit(compiled_generate_function)/convert_element_type" stack_frame_id=823} + %jit_cumsum_.1 = s32[1,41]{1,0} call(%convert_element_type.2234), to_apply=%cumsum.286, metadata={op_name="jit(compiled_generate_function)/jit(cumsum)" stack_frame_id=826} + %sub.529 = s32[1,41]{1,0} subtract(%jit_cumsum_.1, %convert_element_type.2234), metadata={op_name="jit(compiled_generate_function)/sub" stack_frame_id=827} + %constant.137 = s32[] constant(0) + %convert_element_type.1731 = s32[1,41]{1,0} broadcast(%constant.137), dimensions={}, metadata={op_name="jit(compiled_generate_function)/convert_element_type" stack_frame_id=510} + %convert_element_type.2235 = pred[1,41]{1,0} compare(%sub.529, %convert_element_type.1731), direction=NE, metadata={op_name="jit(compiled_generate_function)/convert_element_type" stack_frame_id=510} + %not.7 = pred[1,41]{1,0} not(%convert_element_type.2235), metadata={op_name="jit(compiled_generate_function)/not" stack_frame_id=830} + ROOT %tuple.11 = (pred[1,41]{1,0}, s32[1,41]{1,0}, u32[2]{0}) tuple(%not.7, %while.30#1, %while.30#0) +} + diff --git a/third_party/xla/xla/backends/gpu/autotuner/BUILD b/third_party/xla/xla/backends/gpu/autotuner/BUILD index 2a33d0eb5faa39..17dfa6f4381a48 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/BUILD +++ b/third_party/xla/xla/backends/gpu/autotuner/BUILD @@ -590,6 +590,7 @@ cc_library( "//xla/tsl/platform:errors", "//xla/tsl/platform:statusor", "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/container:inlined_vector", "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_macros", @@ -933,6 +934,7 @@ xla_cc_binary( "//xla/backends/autotuner", "//xla/backends/autotuner:autotune_cache_store", "//xla/backends/autotuner:autotuner_cache_interface", + "//xla/backends/autotuner:backends_proto_cc", "//xla/backends/autotuner:codegen_backend", "//xla/backends/autotuner:codegen_orchestrator", "//xla/backends/autotuner:config_assigner", @@ -1012,6 +1014,7 @@ xla_test( "//xla:xla_proto_cc", "//xla/backends/autotuner", "//xla/backends/autotuner:autotuner_cache_interface", + "//xla/backends/autotuner:backends_proto_cc", "//xla/backends/autotuner:codegen_backend", "//xla/backends/autotuner:codegen_orchestrator", "//xla/backends/autotuner:directory_store", diff --git a/third_party/xla/xla/backends/gpu/autotuner/autotuner_main.cc b/third_party/xla/xla/backends/gpu/autotuner/autotuner_main.cc index d8c4de4f3dd45f..28aa98ee00ddd5 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/autotuner_main.cc +++ b/third_party/xla/xla/backends/gpu/autotuner/autotuner_main.cc @@ -39,6 +39,7 @@ limitations under the License. #include "mlir/IR/MLIRContext.h" #include "xla/backends/autotuner/autotuner.h" #include "xla/backends/autotuner/autotuner_cache_interface.h" +#include "xla/backends/autotuner/backends.pb.h" #include "xla/backends/autotuner/codegen_backend.h" #include "xla/backends/autotuner/codegen_orchestrator.h" #include "xla/backends/autotuner/directory_store.h" diff --git a/third_party/xla/xla/backends/gpu/autotuner/fission_backend.cc b/third_party/xla/xla/backends/gpu/autotuner/fission_backend.cc index 96a020c95a6b4e..835579afb51bd4 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/fission_backend.cc +++ b/third_party/xla/xla/backends/gpu/autotuner/fission_backend.cc @@ -40,6 +40,7 @@ limitations under the License. #include "xla/shape_util.h" #include "xla/status_macros.h" #include "xla/tools/hlo_decomposer.h" +#include "xla/xla.pb.h" namespace xla { diff --git a/third_party/xla/xla/backends/gpu/autotuner/hipblaslt.cc b/third_party/xla/xla/backends/gpu/autotuner/hipblaslt.cc index ed41dcfdb0ba01..9a0ec963f79d8e 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/hipblaslt.cc +++ b/third_party/xla/xla/backends/gpu/autotuner/hipblaslt.cc @@ -22,6 +22,7 @@ limitations under the License. #include #include "absl/algorithm/container.h" +#include "absl/container/inlined_vector.h" #include "absl/log/log.h" #include "absl/status/status.h" #include "absl/status/status_macros.h" diff --git a/third_party/xla/xla/backends/gpu/autotuner/native_emitter_test.cc b/third_party/xla/xla/backends/gpu/autotuner/native_emitter_test.cc index 3bfadc6793afc0..25eb45f1bfd3f3 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/native_emitter_test.cc +++ b/third_party/xla/xla/backends/gpu/autotuner/native_emitter_test.cc @@ -15,6 +15,7 @@ limitations under the License. #include "xla/backends/gpu/autotuner/native_emitter.h" +#include #include #include diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/transforms/promote_shuffle_to_dpp.cc b/third_party/xla/xla/backends/gpu/codegen/emitters/transforms/promote_shuffle_to_dpp.cc index 74400a7a7b3165..90ac6b35bbaa69 100644 --- a/third_party/xla/xla/backends/gpu/codegen/emitters/transforms/promote_shuffle_to_dpp.cc +++ b/third_party/xla/xla/backends/gpu/codegen/emitters/transforms/promote_shuffle_to_dpp.cc @@ -14,7 +14,6 @@ limitations under the License. ==============================================================================*/ #include -#include #include #include diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc index fab420a6d06a74..f30e4a414fd273 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc @@ -47,6 +47,7 @@ limitations under the License. #include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/IR/Attributes.h" #include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinOps.h" #include "mlir/IR/BuiltinTypeInterfaces.h" #include "mlir/IR/BuiltinTypes.h" #include "mlir/IR/Location.h" @@ -1052,8 +1053,22 @@ absl::StatusOr GetCollectiveBlockLevelFusionConfig( block_level_config.set_num_ctas(1); // No block-level clustering. block_level_config.set_num_stages(1); // No pipelining of loops. xtile::Tile* output_tile = block_level_config.add_output_tiles(); - const llvm::SmallVector tile_sizes = + llvm::SmallVector tile_sizes = GreedyPowerOfTwoTiles(output_shape, launch_dims.num_blocks()); + // For AllGather, the tile on the gather dimension must not exceed the + // per-rank size. Otherwise a single tile would span multiple replicas, + // but the tiling framework assigns a single replica_id per tile. + if (instr->opcode() == HloOpcode::kAllGather) { + const auto* all_gather = Cast(instr); + const int64_t gather_dim = all_gather->all_gather_dimension(); + const int64_t num_devices = + std::get(collective_info).num_devices; + const int64_t per_rank_size = + output_shape.dimensions(gather_dim) / num_devices; + tile_sizes[gather_dim] = std::min( + tile_sizes[gather_dim], static_cast(llvm::bit_floor( + static_cast(per_rank_size)))); + } output_tile->mutable_sizes()->Assign(tile_sizes.begin(), tile_sizes.end()); ABSL_RETURN_IF_ERROR( ValidateBlockLevelFusionConfig(block_level_config, collective_info)); @@ -1087,6 +1102,25 @@ absl::StatusOr> GetCollectiveUnmanagedKernelArguments( case HloOpcode::kAllReduce: return GetAllReduceUnmanagedKernelArguments( computation, Cast(root)); + case HloOpcode::kAllGather: { + // AllGather only needs the 3 metadata args (rank, signal_value, + // signal_buffers). No per-parameter scratch buffer args because the + // input parameter already maps to the symmetric scratch buffer via + // the pointer table mechanism in RequiredReplicaIdBounds. + const int32_t num_devices = Cast(root) + ->device_list() + ->num_devices_per_group(); + std::vector unmanaged_arguments; + unmanaged_arguments.reserve(kNumCollectiveMetadataArgs); + // rank and signal_value + unmanaged_arguments.push_back(ShapeUtil::MakeShape(S32, {})); + unmanaged_arguments.push_back(ShapeUtil::MakeShape(S32, {})); + // signal_buffers: pointer-to-pointer table + static constexpr int32_t kMaxBlocksPerGrid = 32; + unmanaged_arguments.push_back( + ShapeUtil::MakeShape(S32, {num_devices, kMaxBlocksPerGrid})); + return unmanaged_arguments; + } default: return std::vector(); } @@ -1103,6 +1137,16 @@ absl::StatusOr AddCollectiveMetadataArguments( fn_arg_types.push_back(ttir::PointerType::get( ttir::PointerType::get(b.getI32Type(), kGlobalAddressSpace), kGlobalAddressSpace)); + + // For AllGather, the input parameter already maps to the symmetric scratch + // buffer via RequiredReplicaIdBounds/SelectBufferOp, so we don't add + // per-parameter scratch buffer opaque args. Only AllReduce (and future ops + // that need explicit remote buffer pointers) add them. + const HloInstruction* root = hlo_computation->root_instruction(); + if (root->opcode() == HloOpcode::kAllGather) { + return kNumCollectiveMetadataArgs; + } + for (HloInstruction* p : hlo_computation->parameter_instructions()) { PrimitiveType type = p->shape().element_type(); mlir::Type ir_type; @@ -1160,4 +1204,77 @@ absl::StatusOr CreateCollectiveKernelSpec( } } +CollectiveCodegenConfig CreateCollectiveCodegenConfig( + const HloInstruction* instr) { + const HloInstruction* collective = instr; + if (instr->opcode() == HloOpcode::kFusion) { + collective = instr->fused_instructions_computation()->root_instruction(); + } + CollectiveCodegenConfig config; + auto gpu_config = collective->backend_config(); + if (!gpu_config.ok()) { + return config; + } + const auto strategy = + gpu_config->collective_backend_config().kernel_strategy(); + // One-shot AllGather: the runtime copies input to symmetric scratch and the + // kernel needs a barrier before reading peers' data. + if (collective->opcode() == HloOpcode::kAllGather && + strategy == CollectiveBackendConfig::KERNEL_STRATEGY_TRITON_ONE_SHOT) { + config.copy_input_to_scratch = true; + config.emit_entry_barrier = true; + } + return config; +} + +absl::Status EmitCollectiveEntryBarrier(mlir::ModuleOp module, + int32_t world_size) { + // Find the xtile::EntryFuncOp in the module. + xtile::EntryFuncOp entry_func = nullptr; + for (auto fn : module.getOps()) { + entry_func = fn; + break; + } + if (!entry_func) { + return absl::InternalError( + "No xtile::EntryFuncOp found in module for barrier insertion."); + } + + // The opaque args are appended after the regular input/output args. + // Their count is stored in the "num_opaque_args" attribute. + auto num_opaque_attr = + entry_func->getAttrOfType("num_opaque_args"); + if (!num_opaque_attr || + num_opaque_attr.getInt() < kNumCollectiveMetadataArgs) { + return absl::InternalError( + absl::StrCat("Expected at least ", kNumCollectiveMetadataArgs, + " opaque args for collective entry barrier, got ", + num_opaque_attr ? num_opaque_attr.getInt() : 0)); + } + + // The opaque args start at (total_args - num_opaque_args). + int32_t total_args = entry_func.getNumArguments(); + int32_t opaque_start = + total_args - num_opaque_attr.getInt() - kNumTileIndexArgs; + // Layout: opaque[0]=rank, opaque[1]=signal_value, opaque[2]=signal_buffers + mlir::Value rank_arg = entry_func.getArgument(opaque_start); + mlir::Value signal_value_arg = entry_func.getArgument(opaque_start + 1); + mlir::Value signal_buffers_arg = entry_func.getArgument(opaque_start + 2); + + // Insert at the beginning of the entry block, right after program_id + // extraction (which is the first op). We want the barrier before any + // scf.for tile loop. + mlir::Block& entry_block = entry_func.front(); + auto loc = entry_func.getLoc(); + mlir::ImplicitLocOpBuilder builder(loc, &entry_block, entry_block.begin()); + + // Inter-block barrier via signal flags. This blocks until all + // remote ranks have also signaled. + mtx::BlockBarrierOp::create(builder, signal_buffers_arg, rank_arg, + signal_value_arg, + builder.getI32IntegerAttr(world_size)); + + return absl::OkStatus(); +} + } // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.h b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.h index 03de7aedab2c8b..094e529a0ded1b 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.h +++ b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.h @@ -24,6 +24,7 @@ limitations under the License. #include "absl/status/statusor.h" #include "llvm/ADT/SmallVector.h" #include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinOps.h" #include "mlir/IR/PatternMatch.h" #include "mlir/IR/Types.h" #include "mlir/Support/LLVM.h" @@ -100,9 +101,23 @@ absl::StatusOr> GetCollectiveUnmanagedKernelArguments( mlir::LogicalResult RewriteAllReduce(mlir::stablehlo::AllReduceOp op, mlir::PatternRewriter& rewriter); +// Creates a lightweight codegen config from the collective HLO instruction. +// The returned config drives codegen decisions (e.g. whether the runtime must +// copy the input to scratch before kernel launch). +CollectiveCodegenConfig CreateCollectiveCodegenConfig( + const HloInstruction* instr); + // Creates a CollectiveKernelSpec for a given collective or fusion instruction. absl::StatusOr CreateCollectiveKernelSpec( const HloInstruction* instr, const LaunchDimensions& launch_dimensions); +// Emits a collective entry barrier at the start of the entry function in +// |module|. The barrier ensures all ranks have completed their D2D copies +// before any rank starts reading from the symmetric scratch buffers. +// Uses the opaque metadata args (rank, signal_value, signal_buffers) already +// present in the EntryFuncOp. +absl::Status EmitCollectiveEntryBarrier(mlir::ModuleOp module, + int32_t world_size); + } // namespace xla::gpu #endif // XLA_BACKENDS_GPU_CODEGEN_TRITON_COLLECTIVE_EMITTER_H_ diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/support.cc b/third_party/xla/xla/backends/gpu/codegen/triton/support.cc index fa6aba5d7463dc..ddf31f482d22a1 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/support.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/support.cc @@ -879,16 +879,8 @@ CodegenDecision IsTritonSupportedInstructionImpl( return IsTritonSupportedAllReduce(*Cast(&instr), gpu_version); case HloOpcode::kAllGather: - if (instr.shape().element_type() == S4) { - return CodegenDecision::Forbid("S4 is not supported."); - } - return instr.GetModule() - ->config() - .debug_options() - .xla_gpu_experimental_enable_tiling_propagation() - ? CodegenDecision::Allow() - : CodegenDecision::Forbid(absl::StrCat( - HloOpcodeString(instr.opcode()), " is not supported")); + return CodegenDecision(instr.shape().element_type() != S4, + "S4 is not supported."); default: // Not all instructions have a special handling. break; diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_gather.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_gather.hlo index e10ecd9dfb2f69..fbe165429e82bb 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_gather.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_gather.hlo @@ -1,3 +1,4 @@ +// RUN: fusion_to_triton %s | FileCheck %s // RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s f { %param0 = f32[128,128]{1,0} parameter(0) @@ -35,7 +36,6 @@ ENTRY entry { // CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: i32, // CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: i32, // CHECK-SAME: %[[ARG4:[a-zA-Z0-9_]+]]: !tt.ptr>, -// CHECK-SAME: %[[ARG5:[a-zA-Z0-9_]+]]: !tt.ptr>, // CHECK-SAME: %[[PID:[a-zA-Z0-9_]+]]: index) // CHECK: %[[IDX0:.*]] = xla.apply_indexing #[[MAP0]](%[[PID]]) diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_deviceless_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_deviceless_test.cc index ac71a1c2873be6..c7937b9ceb9ce8 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_deviceless_test.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_deviceless_test.cc @@ -473,5 +473,57 @@ ENTRY entry { triple, data_layout, mlir_context)); } +TEST_F(TritonEmitterDevicelessTest, + AllGatherFusionUsesTiledHloComputationWhenTilingPropagationDisabled) { + constexpr absl::string_view kHloText = R"( +f { + param0 = f32[128,128]{1,0} parameter(0) + ROOT result = f32[256,128]{1,0} all-gather(param0), + replica_groups={{0,1}}, dimensions={0} +} + +ENTRY entry { + p0 = f32[128,128]{1,0} parameter(0) + ROOT fusion = f32[256,128]{1,0} fusion(p0), + kind=kCustom, calls=f, + backend_config={ + "fusion_backend_config": { + "kind": "__triton_collective", + "block_level_fusion_config": { + "num_warps": "4", + "output_tiles": [{sizes: [16,16]}], + "num_ctas": 1, + "num_stages": 1, + "is_tma_allowed": false, + "is_warp_specialization_allowed": false + } + } + } +} +)"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr hlo_module, + ParseAndReturnVerifiedModule(kHloText)); + // Explicitly ensure xla_gpu_experimental_enable_tiling_propagation is false. + hlo_module->mutable_config() + .mutable_debug_options() + .set_xla_gpu_experimental_enable_tiling_propagation(false); + + const auto* fusion = Cast( + hlo_module->entry_computation()->root_instruction()); + ASSERT_NE(fusion, nullptr); + + const se::DeviceDescription dev_info = TestGpuDeviceInfo::H100SXMDeviceInfo(); + mlir::MLIRContext mlir_context; + RegisterSymbolicExprStorage(&mlir_context); + + ASSERT_OK_AND_ASSIGN(const auto gpu_backend_config, + fusion->backend_config()); + EXPECT_OK(CreateTritonModule("test_fn", *fusion, dev_info, + BlockLevelParameters::FromBlockLevelFusionConfig( + gpu_backend_config.fusion_backend_config() + .block_level_fusion_config()), + mlir_context)); +} + } // namespace } // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/xtile_compiler.cc b/third_party/xla/xla/backends/gpu/codegen/triton/xtile_compiler.cc index d3f4c7c00e1e0c..d52453328a497e 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/xtile_compiler.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/xtile_compiler.cc @@ -227,6 +227,16 @@ absl::Status ValidateComplexUseInTritonFusion( } return absl::OkStatus(); } + +bool IsAllGatherFusion(const HloFusionInstruction& fusion) { + const HloComputation* computation = fusion.fused_instructions_computation(); + if (computation == nullptr) { + return false; + } + return absl::c_any_of(computation->instructions(), + HloPredicateIsOp); +} + } // namespace namespace ttir = ::mlir::triton; @@ -300,7 +310,7 @@ absl::StatusOr> TileAndEmitXTileModule( bool use_experimental_tiling, bool enable_same_shape_multi_output_fusion) { const HloComputation* computation = fusion.fused_instructions_computation(); - if (use_experimental_tiling) { + if (use_experimental_tiling || IsAllGatherFusion(fusion)) { using experimental::TiledHloComputation; using experimental::TilingSpace; @@ -376,7 +386,8 @@ absl::StatusOr CreateTritonModule( const DebugOptions& debug_options = fusion.GetModule()->config().debug_options(); bool use_experimental_tiling = - debug_options.xla_gpu_experimental_enable_tiling_propagation(); + debug_options.xla_gpu_experimental_enable_tiling_propagation() || + IsAllGatherFusion(fusion); bool enable_same_shape_multi_output_fusion = debug_options .xla_gpu_experimental_enable_same_shape_multi_output_fusion(); @@ -436,6 +447,14 @@ absl::StatusOr CreateTritonModule( absl::MakeSpan(opaque_args_types), mlir_context, use_experimental_tiling, enable_same_shape_multi_output_fusion)); + if (fusion_kind == kTritonCollectiveFusionKind && + CreateCollectiveCodegenConfig(&fusion).emit_entry_barrier) { + const HloInstruction* root = hlo_computation->root_instruction(); + int32_t world_size = root->replica_groups()[0].replica_ids_size(); + ABSL_RETURN_IF_ERROR( + EmitCollectiveEntryBarrier(triton_module.get(), world_size)); + } + if (DumpingEnabledForHloModule(*hlo_computation->parent()) && DumpingEnabledForEmitter("triton-fusion", debug_options)) { auto suffix = absl::StrCat(fusion.name(), ".before_validation.ttir.txt"); diff --git a/third_party/xla/xla/backends/gpu/collectives/BUILD b/third_party/xla/xla/backends/gpu/collectives/BUILD index 6e6aadc1c6ff5e..b21fc137166589 100644 --- a/third_party/xla/xla/backends/gpu/collectives/BUILD +++ b/third_party/xla/xla/backends/gpu/collectives/BUILD @@ -1,4 +1,8 @@ -load("@local_config_rocm//rocm:build_defs.bzl", "if_rocm_is_configured") +load( + "@local_config_rocm//rocm:build_defs.bzl", + "if_rocm_is_configured", + "rocm_library", +) load("@local_config_sycl//sycl:build_defs.bzl", "if_sycl_is_configured") load("//xla:xla.default.bzl", "xla_cc_test") load("//xla/stream_executor:build_defs.bzl", "if_cuda_or_rocm_is_configured") @@ -1267,6 +1271,26 @@ cc_library( ], ) +rocm_library( + name = "mori_kernels", + srcs = [ + "mori_kernels.cu.cc", + ], + hdrs = [ + "mori_kernels.h", + "mori_stub.h", + ], + # copybara:uncomment compatible_with = ["//buildenv/target:non_prod"], + copts = ["-U__HIP_DISABLE_CPP_FUNCTIONS__"], # <-- only if needed + linkstatic = True, + tags = [ + "gpu", + "no-oneapi", + "rocm-only", + ], + deps = [], +) + cc_library( name = "mori_collectives", srcs = [ @@ -1276,10 +1300,13 @@ cc_library( hdrs = [ "mori_collectives.h", "mori_communicator.h", + "mori_kernels.h", "mori_stub.h", ], tags = [ "gpu", + "no-oneapi", + "rocm-only", ], visibility = ["//visibility:public"], deps = [ @@ -1299,6 +1326,7 @@ cc_library( "//xla/core/collectives:communicator", "//xla/core/collectives:rank_id", "//xla/core/collectives:reduction_kind", + "//xla/core/collectives:symmetric_memory", "//xla/pjrt/distributed:key_value_store_interface", "//xla/runtime:device_id", "//xla/runtime:process_id", @@ -1307,6 +1335,7 @@ cc_library( "//xla/stream_executor:platform_manager", "//xla/stream_executor:stream", "//xla/stream_executor:stream_executor_h", + "//xla/stream_executor/rocm:rocm_status", "//xla/tsl/platform:env", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/base", @@ -1329,7 +1358,10 @@ cc_library( "@com_google_absl//absl/types:span", "@tsl//tsl/platform:casts", "@tsl//tsl/platform:numbers", - ], + ] + if_rocm_is_configured([ + ":mori_kernels", + "@local_config_rocm//rocm:rocm_headers", + ]), alwayslink = True, ) diff --git a/third_party/xla/xla/backends/gpu/collectives/mori_communicator.cc b/third_party/xla/xla/backends/gpu/collectives/mori_communicator.cc index d1d71dc7fd55c8..8dfb3bfa4e80e9 100644 --- a/third_party/xla/xla/backends/gpu/collectives/mori_communicator.cc +++ b/third_party/xla/xla/backends/gpu/collectives/mori_communicator.cc @@ -36,13 +36,14 @@ limitations under the License. #include "xla/backends/gpu/collectives/cancellation_token.h" #include "xla/backends/gpu/collectives/gpu_collectives.h" #include "xla/backends/gpu/collectives/mori_collectives.h" -#include "xla/backends/gpu/collectives/mori_stub.h" +#include "xla/backends/gpu/collectives/mori_kernels.h" #include "xla/core/collectives/communicator.h" #include "xla/core/collectives/rank_id.h" #include "xla/core/collectives/reduction_kind.h" #include "xla/future.h" #include "xla/primitive_util.h" #include "xla/stream_executor/device_address.h" +#include "xla/stream_executor/rocm/rocm_status.h" #include "xla/stream_executor/stream.h" #include "xla/util.h" #include "xla/xla_data.pb.h" @@ -51,31 +52,104 @@ limitations under the License. namespace shmem = ::mori::shmem; namespace xla::gpu { -static auto AsRocmStream(se::Stream* stream) { - return reinterpret_cast( +using ::mori::collective::CollectivesFacade; +namespace { + +hipStream_t AsHipStream(se::Stream* stream) { + return reinterpret_cast( stream->platform_specific_handle().stream); } -static size_t ToMoriByteCount(PrimitiveType dtype, size_t count) { +size_t ToMoriByteCount(PrimitiveType dtype, size_t count) { if (primitive_util::IsComplexType(dtype)) { count *= 2; } return count * primitive_util::BitWidth(dtype) / 8; } +absl::StatusOr<::mori::collective::DataType> ToMoriDataType( + PrimitiveType dtype) { +#define MORI_TYPE_DISPATCH(x) \ + case x: \ + return ::mori::collective::DataType::x; + switch (dtype) { + MORI_TYPE_DISPATCH(F8E5M2) + MORI_TYPE_DISPATCH(F8E4M3FN) + MORI_TYPE_DISPATCH(F16) + MORI_TYPE_DISPATCH(BF16) + MORI_TYPE_DISPATCH(S8) + MORI_TYPE_DISPATCH(U8) + MORI_TYPE_DISPATCH(S32) + MORI_TYPE_DISPATCH(U32) + MORI_TYPE_DISPATCH(S64) + MORI_TYPE_DISPATCH(U64) + MORI_TYPE_DISPATCH(F32) + MORI_TYPE_DISPATCH(F64) + default: + return absl::UnimplementedError(absl::StrFormat( + "MORI: unsupported dtype: %d", static_cast(dtype))); + } +#undef MORI_TYPE_DISPATCH +} + +// Translate an XLA ReductionKind to the facade's reduction-op enum. +absl::StatusOr<::mori::collective::ReduceOpKind> ToMoriReduceOp( + ReductionKind r) { +#define MORI_OP_DISPATCH(x) \ + case ReductionKind::x: \ + return ::mori::collective::ReduceOpKind::x; + switch (r) { + MORI_OP_DISPATCH(SUM) + MORI_OP_DISPATCH(PRODUCT) + MORI_OP_DISPATCH(MIN) + MORI_OP_DISPATCH(MAX) + default: + return absl::UnimplementedError(absl::StrFormat( + "MORI: unsupported reduction op: %d", static_cast(r))); + } +#undef MORI_OP_DISPATCH +} + +absl::StatusOr ToStream(const Communicator::Executor& executor) { + if (auto* gpu_executor = + absl::down_cast(&executor)) { + return gpu_executor->stream(); + } + return InvalidArgument("Communicator executor is not a GPU executor"); +} +} // namespace + absl::StatusOr> MoriCommunicator::Create( MoriCollectives* coll, std::shared_ptr cancel, int rank, absl::Span rank_to_pe) { auto comm = absl::WrapUnique(new MoriCommunicator(coll, cancel)); const int num_ranks = static_cast(rank_to_pe.size()); + if (num_ranks <= 0) { + return absl::InvalidArgumentError(absl::StrFormat( + "MoriCommunicator: unsupported number of ranks %d", num_ranks)); + } comm->rank_ = rank; comm->num_ranks_ = num_ranks; + + // The CollectivesFacade owns this communicator's symmetric-heap staging + // buffer and the push reduce-scatter group counters. It records the rank + // identity (rank/num_ranks) and allocates the ~2GB staging; the unique_ptr + // frees it (before ShmemFinalize) when the communicator is destroyed. + const size_t buffer_size = 2UL << 30; // 2GB + comm->facade_ = CollectivesFacade::Create(rank, num_ranks, buffer_size); + if (comm->facade_ == nullptr) { + return absl::InternalError("CollectivesFacade::Create failed"); + } VLOG(1) << "Created " << *comm << " with participants: " << num_ranks; return comm; } -MoriCommunicator::~MoriCommunicator() {} +MoriCommunicator::~MoriCommunicator() { + // facade_ (unique_ptr) releases this communicator's staging + counters via + // the CollectivesFacade dtor here, before MoriCollectives::Finalize() -> + // ShmemFinalize. +} #define CHECK_CANCELLED() \ if (cancel_->IsCancelled()) { \ @@ -97,41 +171,26 @@ absl::Status MoriCommunicator::Abort() { return absl::OkStatus(); } -absl::Status MoriCommunicator::Barrier(const Communicator::Executor& executor) { +absl::Status MoriCommunicator::Barrier(const Executor& executor) { VLOG(1) << "Barrier: " << ToString(); CHECK_CANCELLED() ABSL_ASSIGN_OR_RETURN(se::Stream * stream, ToStream(executor)); - (void)stream; - // return xla_mori::BarrierOnStream(AsRocmStream(stream)); - return absl::OkStatus(); + return se::gpu::ToStatus(facade_->RunBarrier(AsHipStream(stream))); } absl::StatusOr MoriCommunicator::NumRanks() const { - VLOG(5) << "Get the number of ranks in MORI communicator: " << ToString(); CHECK_CANCELLED() - return static_cast(num_ranks_); } absl::StatusOr MoriCommunicator::CurrentRank() { - VLOG(5) << "Get current rank in MORI communicator: " << ToString(); CHECK_CANCELLED() - return static_cast(rank_); } std::string MoriCommunicator::ToString() const { - return absl::StrFormat("MoriCommunicator(rank=%d, num_ranks=%d, my_pe=%d)", - rank_, num_ranks_, shmem::ShmemMyPe()); -} - -absl::StatusOr MoriCommunicator::ToStream( - const Executor& executor) { - if (auto* gpu_executor = - absl::down_cast(&executor)) { - return gpu_executor->stream(); - } - return InvalidArgument("Communicator executor is not a GPU executor"); + return absl::StrFormat("MoriCommunicator(rank=%d, num_ranks=%d)", rank_, + num_ranks_); } Future<> MoriCommunicator::AllReduce(se::DeviceAddressBase send_buffer, @@ -201,20 +260,20 @@ Future<> MoriCommunicator::CollectivePermute( }); } -Future<> MoriCommunicator::Send(se::DeviceAddressBase recv_buffer, - se::DeviceAddressBase send_buffer, +Future<> MoriCommunicator::Send(se::DeviceAddressBase send_buffer, PrimitiveType dtype, size_t count, RankId peer, const Executor& executor) { - return P2P(P2PType::Send, dtype, recv_buffer, send_buffer, count, peer, - executor); + return Execute([send_buffer, dtype, count, peer, &executor, this]() { + return LaunchSend(send_buffer, dtype, count, peer, executor); + }); } Future<> MoriCommunicator::Recv(se::DeviceAddressBase recv_buffer, - se::DeviceAddressBase send_buffer, PrimitiveType dtype, size_t count, RankId peer, const Executor& executor) { - return P2P(P2PType::Recv, dtype, recv_buffer, send_buffer, count, peer, - executor); + return Execute([recv_buffer, dtype, count, peer, &executor, this]() { + return LaunchRecv(recv_buffer, dtype, count, peer, executor); + }); } absl::Status MoriCommunicator::LaunchAllGather( @@ -222,11 +281,14 @@ absl::Status MoriCommunicator::LaunchAllGather( PrimitiveType dtype, size_t count, const Executor& executor) { CHECK_CANCELLED() ABSL_ASSIGN_OR_RETURN(se::Stream * stream, ToStream(executor)); - VLOG(3) << "LaunchAllGather: send_buffer=" << send_buffer.opaque() + + VLOG(3) << "Launch AllGather: send_buffer=" << send_buffer.opaque() << " recv_buffer=" << recv_buffer.opaque() << " count=" << count << " dtype=" << primitive_util::LowercasePrimitiveTypeName(dtype) - << " stream=" << AsRocmStream(stream); - return absl::UnimplementedError("Not implemented"); + << " stream=" << AsHipStream(stream); + return se::gpu::ToStatus(facade_->RunAllGather( + send_buffer.opaque(), recv_buffer.opaque(), ToMoriByteCount(dtype, count), + AsHipStream(stream))); } absl::Status MoriCommunicator::LaunchAllReduce( @@ -234,25 +296,20 @@ absl::Status MoriCommunicator::LaunchAllReduce( PrimitiveType dtype, size_t count, ReductionKind reduction_kind, const Executor& executor) { CHECK_CANCELLED() - ABSL_ASSIGN_OR_RETURN(se::Stream * stream, ToStream(executor)); - auto gpu_stream = AsRocmStream(stream); - (void)gpu_stream; - void* source_ptr = send_buffer.opaque(); - void* dest_ptr = recv_buffer.opaque(); - (void)source_ptr; - (void)dest_ptr; - if (primitive_util::IsComplexType(dtype)) { - count *= 2; - } VLOG(3) << absl::StreamFormat( - "Launch MORI AllReduce send_buffer=%p; recv_buffer=%p; dtype=%s; " - "count=%d; reduction_kind=%v; device_ordinal=%d", + "Launch AllReduce: send_buffer=%p; recv_buffer=%p; dtype=%s; count=%d; " + "reduction_kind=%v; stream=%p", send_buffer.opaque(), recv_buffer.opaque(), primitive_util::LowercasePrimitiveTypeName(dtype), count, reduction_kind, - stream->parent()->device_ordinal()); - return absl::UnimplementedError("Not implemented"); + stream); + + ABSL_ASSIGN_OR_RETURN(auto dt, ToMoriDataType(dtype)); + ABSL_ASSIGN_OR_RETURN(auto op, ToMoriReduceOp(reduction_kind)); + return se::gpu::ToStatus(facade_->RunAllReduce(send_buffer.opaque(), + recv_buffer.opaque(), count, + dt, op, AsHipStream(stream))); } absl::Status MoriCommunicator::LaunchReduceScatter( @@ -265,8 +322,78 @@ absl::Status MoriCommunicator::LaunchReduceScatter( VLOG(3) << "LaunchReduceScatter: send_buffer=" << send_buffer.opaque() << " recv_buffer=" << recv_buffer.opaque() << " count=" << count << " dtype=" << primitive_util::LowercasePrimitiveTypeName(dtype) - << " stream=" << AsRocmStream(stream); - return absl::UnimplementedError("Not implemented"); + << " stream=" << AsHipStream(stream); + + ABSL_ASSIGN_OR_RETURN(auto dt, ToMoriDataType(dtype)); + ABSL_ASSIGN_OR_RETURN(auto op, ToMoriReduceOp(kind)); + return se::gpu::ToStatus( + facade_->RunReduceScatter(send_buffer.opaque(), recv_buffer.opaque(), + count, dt, op, AsHipStream(stream))); +} + +absl::Status MoriCommunicator::LaunchAllToAll( + absl::InlinedVector send_buffers, + absl::InlinedVector recv_buffers, + PrimitiveType dtype, size_t count, const Executor& executor) { + CHECK_CANCELLED() + ABSL_ASSIGN_OR_RETURN(se::Stream * stream, ToStream(executor)); + + auto format_addr = [](std::string* out, se::DeviceAddressBase buf) { + absl::StrAppendFormat(out, "%p", buf.opaque()); + }; + VLOG(3) << absl::StreamFormat( + "Launch MORI AllToAll operation; send_buffers=[%s]; recv_buffers=[%s]; " + "dtype=%s; count=%d; stream=%p", + absl::StrJoin(send_buffers, ", ", format_addr), + absl::StrJoin(recv_buffers, ", ", format_addr), + primitive_util::LowercasePrimitiveTypeName(dtype), count, + AsHipStream(stream)); + + if (send_buffers.size() != recv_buffers.size() || + send_buffers.size() != static_cast(num_ranks_)) { + return InvalidArgument( + "Number of send/recv buffers and number of ranks mismatch"); + } + + CollectivesFacade::AddressVector addrs; + addrs.reserve(num_ranks_); + for (int p = 0; p < num_ranks_; ++p) { + addrs.emplace_back(send_buffers[p].opaque(), recv_buffers[p].opaque()); + } + return se::gpu::ToStatus(facade_->RunAllToAll( + addrs, ToMoriByteCount(dtype, count), AsHipStream(stream))); +} + +absl::Status MoriCommunicator::LaunchSend(se::DeviceAddressBase send_buffer, + PrimitiveType dtype, size_t count, + RankId peer, + const Executor& executor) { + CHECK_CANCELLED() + ABSL_ASSIGN_OR_RETURN(se::Stream * stream, ToStream(executor)); + VLOG(3) << absl::StreamFormat( + "Launch MORI Send operation; send_buffer=%p; dtype=%s; count=%d; " + "peer=%d; stream=%p", + send_buffer.opaque(), primitive_util::LowercasePrimitiveTypeName(dtype), + count, peer.value(), AsHipStream(stream)); + return se::gpu::ToStatus( + facade_->RunSend(send_buffer.opaque(), ToMoriByteCount(dtype, count), + static_cast(peer.value()), AsHipStream(stream))); +} + +absl::Status MoriCommunicator::LaunchRecv(se::DeviceAddressBase recv_buffer, + PrimitiveType dtype, size_t count, + RankId peer, + const Executor& executor) { + CHECK_CANCELLED() + ABSL_ASSIGN_OR_RETURN(se::Stream * stream, ToStream(executor)); + VLOG(3) << absl::StreamFormat( + "Launch MORI Recv operation; recv_buffer=%p; dtype=%s; count=%d; " + "peer=%d; stream=%p", + recv_buffer.opaque(), primitive_util::LowercasePrimitiveTypeName(dtype), + count, peer.value(), AsHipStream(stream)); + return se::gpu::ToStatus( + facade_->RunRecv(recv_buffer.opaque(), ToMoriByteCount(dtype, count), + static_cast(peer.value()), AsHipStream(stream))); } absl::Status MoriCommunicator::LaunchCollectivePermute( @@ -275,13 +402,11 @@ absl::Status MoriCommunicator::LaunchCollectivePermute( absl::Span target_ranks, const Executor& executor) { CHECK_CANCELLED() ABSL_ASSIGN_OR_RETURN(se::Stream * stream, ToStream(executor)); - size_t bytes = ToMoriByteCount(dtype, count); - (void)bytes; auto rank_formatter = [](std::string* out, RankId rank) { absl::StrAppendFormat(out, "%d", rank.value()); }; VLOG(3) << absl::StreamFormat( - "[%d] Launch MORI CollectivePermute operation; send_buffer=%p; " + "[%d] Launch CollectivePermute: send_buffer=%p; " "recv_buffer=%p; dtype=%s; source_rank=%s; target_[ranks=%s]; count=%d; " "stream=%p", stream->parent()->device_ordinal(), send_buffer.opaque(), @@ -289,41 +414,15 @@ absl::Status MoriCommunicator::LaunchCollectivePermute( source_rank ? absl::StrCat(source_rank->value()) : "", absl::StrJoin(target_ranks, ", ", rank_formatter), count, stream); - return absl::UnimplementedError("Not implemented"); -} - -// Performs point-to-point communication between two ranks using MORI. -// Send: launches a single GPU kernel that copies data to the peer via P2P -// and sets a completion flag on the peer. -// Recv: launches a single-thread GPU kernel that waits for the flag. -absl::Status MoriCommunicator::P2P(P2PType p2p_type, PrimitiveType dtype, - se::DeviceAddressBase recv_buffer, - se::DeviceAddressBase send_buffer, - size_t count, RankId peer, - const Executor& executor) { - const char* stype = (p2p_type == P2PType::Send ? " Send" : " Recv"); - VLOG(1) << CurrentRank().value() << stype << " to " << peer.value() - << " count " << count << " MORI communicator: " << ToString(); - CHECK_CANCELLED() - - void* source_ptr = send_buffer.opaque(); - void* dest_ptr = recv_buffer.opaque(); - - ABSL_ASSIGN_OR_RETURN(se::Stream * stream, ToStream(executor)); - auto gpu_stream = AsRocmStream(stream); - size_t bytes = ToMoriByteCount(dtype, count); - int res = 0; - (void)bytes; - (void)res; - (void)gpu_stream; - (void)source_ptr; - (void)dest_ptr; - (void)peer; - (void)stream; - (void)dtype; - (void)count; - (void)p2p_type; - return absl::UnimplementedError("Not implemented"); + std::vector dstPes; + dstPes.reserve(target_ranks.size()); + for (RankId rank : target_ranks) { + dstPes.push_back(static_cast(rank.value())); + } + const int srcPe = source_rank ? static_cast(source_rank->value()) : -1; + return se::gpu::ToStatus(facade_->RunCollectivePermute( + send_buffer.opaque(), recv_buffer.opaque(), ToMoriByteCount(dtype, count), + srcPe, dstPes, AsHipStream(stream))); } Future<> MoriCommunicator::GroupExecute( @@ -339,19 +438,16 @@ absl::Status MoriCommunicator::GroupLaunch( } absl::Status MoriCommunicator::Quiet(const Executor& executor) { - VLOG(1) << "Quiet MORI communicator: " << ToString(); + VLOG(1) << "Quiet: " << ToString(); CHECK_CANCELLED() ABSL_ASSIGN_OR_RETURN(se::Stream * stream, ToStream(executor)); - auto gpu_stream = AsRocmStream(stream); - (void)gpu_stream; - return absl::UnimplementedError("Not implemented"); + return se::gpu::ToStatus(facade_->RunQuiet(AsHipStream(stream))); } absl::Status MoriCommunicator::Fence() { - VLOG(1) << "Fence MORI communicator: " << ToString(); + VLOG(1) << "Fence: " << ToString(); CHECK_CANCELLED() - // rocm_mori_fence(); - return absl::UnimplementedError("Not implemented"); + return se::gpu::ToStatus(facade_->RunFence()); } absl::Status MoriCommunicator::PollUntilDone() const { diff --git a/third_party/xla/xla/backends/gpu/collectives/mori_communicator.h b/third_party/xla/xla/backends/gpu/collectives/mori_communicator.h index ac4b2ffb0e88b7..61b7f14a12c394 100644 --- a/third_party/xla/xla/backends/gpu/collectives/mori_communicator.h +++ b/third_party/xla/xla/backends/gpu/collectives/mori_communicator.h @@ -25,6 +25,7 @@ limitations under the License. #include "absl/functional/function_ref.h" #include "absl/status/status.h" #include "absl/status/statusor.h" +#include "absl/strings/str_format.h" #include "absl/strings/string_view.h" #include "absl/types/span.h" #include "xla/backends/gpu/collectives/cancellation_token.h" @@ -32,20 +33,44 @@ limitations under the License. #include "xla/core/collectives/communicator.h" #include "xla/core/collectives/rank_id.h" #include "xla/core/collectives/reduction_kind.h" +#include "xla/core/collectives/symmetric_memory.h" #include "xla/future.h" #include "xla/stream_executor/device_address.h" #include "xla/stream_executor/stream.h" #include "xla/xla_data.pb.h" +namespace mori::collective { +class CollectivesFacade; +} // namespace mori::collective + namespace xla::gpu { class MoriCollectives; +// Dummy symmetric memory for the MORI backend. MORI collective buffers are +// allocated directly from the symmetric shmem heap, so the local device +// address doubles as the symmetric handle and no separate registration is +// required. This simply returns the address it was created with. +class MoriSymmetricMemory : public SymmetricMemory { + public: + explicit MoriSymmetricMemory(se::DeviceAddressBase addr) : addr_(addr) {} + + se::DeviceAddressBase addr() const final { return addr_; } + + std::string ToString() const final { + return absl::StrFormat("MoriSymmetricMemory(addr=%p, size=%d)", + addr_.opaque(), addr_.size()); + } + + PackedKernelArg PackKernelArg() const final { return addr_.opaque(); } + + private: + se::DeviceAddressBase addr_; +}; + // XLA collectives communicator wrapping a MORI communicator. class MoriCommunicator : public GpuCommunicator { public: - constexpr static uint32_t kMaxTeams = 24; - friend class MoriCollectives; ~MoriCommunicator() override; @@ -63,6 +88,14 @@ class MoriCommunicator : public GpuCommunicator { absl::StatusOr NumRanks() const final; absl::StatusOr CurrentRank() final; + absl::StatusOr> CreateSymmetricMemory( + se::DeviceAddressBase addr) final { + // Dummy implementation: MORI buffers are already allocated from the + // symmetric shmem heap, so the local device address is the symmetric + // handle. Just wrap and return it unchanged. + return std::make_unique(addr); + } + absl::Status Barrier(const Executor& executor) final; Future<> GroupExecute(absl::AnyInvocable group) final; @@ -98,21 +131,9 @@ class MoriCommunicator : public GpuCommunicator { const Executor& executor) final; Future<> Send(se::DeviceAddressBase send_buffer, PrimitiveType dtype, - size_t count, RankId peer, const Executor& executor) final { - return absl::UnimplementedError("Not implemented"); - } - - Future<> Recv(se::DeviceAddressBase recv_buffer, PrimitiveType dtype, - size_t count, RankId peer, const Executor& executor) final { - return absl::UnimplementedError("Not implemented"); - } - - Future<> Send(se::DeviceAddressBase recv_buffer, - se::DeviceAddressBase send_buffer, PrimitiveType dtype, size_t count, RankId peer, const Executor& executor) final; - Future<> Recv(se::DeviceAddressBase recv_buffer, - se::DeviceAddressBase send_buffer, PrimitiveType dtype, + Future<> Recv(se::DeviceAddressBase recv_buffer, PrimitiveType dtype, size_t count, RankId peer, const Executor& executor) final; // Polls the communicator until any pending non-blocking operations are done @@ -149,9 +170,7 @@ class MoriCommunicator : public GpuCommunicator { absl::Status LaunchAllToAll( absl::InlinedVector send_buffers, absl::InlinedVector recv_buffers, - PrimitiveType dtype, size_t count, const Executor& executor) final { - return absl::UnimplementedError("Not implemented"); - } + PrimitiveType dtype, size_t count, const Executor& executor) final; absl::Status LaunchCollectivePermute(se::DeviceAddressBase send_buffer, se::DeviceAddressBase recv_buffer, @@ -162,15 +181,11 @@ class MoriCommunicator : public GpuCommunicator { absl::Status LaunchSend(se::DeviceAddressBase send_buffer, PrimitiveType dtype, size_t count, RankId peer, - const Executor& executor) final { - return absl::UnimplementedError("Not implemented"); - } + const Executor& executor) final; absl::Status LaunchRecv(se::DeviceAddressBase recv_buffer, PrimitiveType dtype, size_t count, RankId peer, - const Executor& executor) final { - return absl::UnimplementedError("Not implemented"); - } + const Executor& executor) final; absl::Status Quiet(const Executor& executor) final; @@ -186,22 +201,16 @@ class MoriCommunicator : public GpuCommunicator { std::shared_ptr cancel) : collectives_(coll), cancel_(std::move(cancel)) {} - enum class P2PType : int32_t { Send, Recv }; - - absl::Status P2P(P2PType p2p_type, PrimitiveType type, - se::DeviceAddressBase recv_buffer, - se::DeviceAddressBase send_buffer, size_t count, RankId peer, - const Executor& executor); - - static absl::StatusOr ToStream(const Executor& executor); - MoriCollectives* collectives_; // Parent MoriCollectives instance // This communicator's participant set (NOT the global MORI clique). `rank_` - // is this rank within the collective, `num_ranks_` the participant count, and - // `rank_to_pe_dev_` a device array mapping collective rank -> global MORI PE. + // is this rank within the collective, `num_ranks_` the participant count. int rank_ = 0; int num_ranks_ = 0; + // Owns this communicator's staging buffer + group counters (created in + // Create(), freed by the facade dtor before ShmemFinalize). Header-only + // facade. + std::unique_ptr<::mori::collective::CollectivesFacade> facade_; // Should all pending collectives cancel? std::shared_ptr cancel_; bool aborted_ = false; // Has Abort() been called? diff --git a/third_party/xla/xla/backends/gpu/collectives/mori_kernels.cu.cc b/third_party/xla/xla/backends/gpu/collectives/mori_kernels.cu.cc new file mode 100644 index 00000000000000..76b616a1942d22 --- /dev/null +++ b/third_party/xla/xla/backends/gpu/collectives/mori_kernels.cu.cc @@ -0,0 +1,19 @@ +/* Copyright 2026 The OpenXLA Authors. +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + http://www.apache.org/licenses/LICENSE-2.0 +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +// This is the single device translation unit for the MORI XLA collectives. It +// is compiled as HIP and defines MORI_KERNELS_IMPL before including the facade, +// so the facade's device path (kernels + non-templated Run* definitions) is +// compiled here exactly once. The host mori_communicator.cc includes the same +// header without MORI_KERNELS_IMPL (decl-only) and links against these symbols. +#define MORI_KERNELS_IMPL +#include "xla/backends/gpu/collectives/mori_kernels.h" diff --git a/third_party/xla/xla/backends/gpu/collectives/mori_kernels.h b/third_party/xla/xla/backends/gpu/collectives/mori_kernels.h new file mode 100644 index 00000000000000..bff6bc1860d5e8 --- /dev/null +++ b/third_party/xla/xla/backends/gpu/collectives/mori_kernels.h @@ -0,0 +1,27 @@ +/* Copyright 2026 The OpenXLA Authors. +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + http://www.apache.org/licenses/LICENSE-2.0 +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef XLA_BACKENDS_GPU_COLLECTIVES_MORI_KERNELS_H_ +#define XLA_BACKENDS_GPU_COLLECTIVES_MORI_KERNELS_H_ + +#include +#include + +// The CollectivesFacade owns the per-device staging + Run* entry points, which +// are non-templated and take mori::collective::DataType / ReduceOpKind enums. +// Host includers (mori_communicator.cc) see decl-only Run* methods; the device +// TU (mori_kernels.cu.cc, compiled as HIP with MORI_KERNELS_IMPL) pulls in the +// full device path and emits the definitions that resolve the host's +// references. +#include "xla/backends/gpu/collectives/mori_stub.h" + +#endif // XLA_BACKENDS_GPU_COLLECTIVES_MORI_KERNELS_H_ diff --git a/third_party/xla/xla/backends/gpu/collectives/mori_stub.h b/third_party/xla/xla/backends/gpu/collectives/mori_stub.h index 267cdf3a35fbe2..2a18155d71dc17 100644 --- a/third_party/xla/xla/backends/gpu/collectives/mori_stub.h +++ b/third_party/xla/xla/backends/gpu/collectives/mori_stub.h @@ -16,9 +16,14 @@ limitations under the License. #ifndef XLA_BACKENDS_GPU_COLLECTIVES_MORI_STUB_H_ #define XLA_BACKENDS_GPU_COLLECTIVES_MORI_STUB_H_ +#include + #include #include #include +#include +#include +#include // Inert stand-in for the subset of the MORI shmem host API used by the MORI // collectives/communicator backbone. These placeholders let the backbone @@ -72,4 +77,74 @@ inline void ShmemFree(void* /*ptr*/) {} } // namespace shmem } // namespace mori +namespace mori { +namespace collective { + +// Element type + reduction op enums mirror the real facade's non-templated API +// (mori/collective/collectives_facade.hpp), so the communicator's enum dispatch +// compiles against either the stub or the real facade. +enum class DataType { + F8E5M2, + F8E4M3FN, + F16, + BF16, + S8, + U8, + S32, + U32, + S64, + U64, + F32, + F64 +}; +enum class ReduceOpKind { SUM, PRODUCT, MIN, MAX }; + +// Inert stand-in for the real MORI CollectivesFacade. Header-only, all Run* are +// no-ops returning hipSuccess. Lets the collectives/communicator wiring compile +// and link without @roc_mori. +class CollectivesFacade { + CollectivesFacade() = default; + + public: + using AddressVector = std::vector>; + + CollectivesFacade(const CollectivesFacade&) = delete; + CollectivesFacade& operator=(const CollectivesFacade&) = delete; + + static std::unique_ptr Create(int /*myPe*/, int /*nPes*/, + size_t /*maxStagingBytes*/) { + return std::unique_ptr(new CollectivesFacade()); + } + ~CollectivesFacade() = default; + + hipError_t RunReduceScatter(const void*, void*, size_t, DataType, + ReduceOpKind, hipStream_t) { + return hipSuccess; + } + hipError_t RunAllReduce(const void*, void*, size_t, DataType, ReduceOpKind, + hipStream_t) { + return hipSuccess; + } + hipError_t RunAllGather(const void*, void*, size_t, hipStream_t) { + return hipSuccess; + } + hipError_t RunAllToAll(const AddressVector&, size_t, hipStream_t) { + return hipSuccess; + } + hipError_t RunBarrier(hipStream_t) { return hipSuccess; } + hipError_t RunSend(const void*, size_t, int, hipStream_t) { + return hipSuccess; + } + hipError_t RunRecv(void*, size_t, int, hipStream_t) { return hipSuccess; } + hipError_t RunCollectivePermute(const void*, void*, size_t, int, + const std::vector&, hipStream_t) { + return hipSuccess; + } + hipError_t RunQuiet(hipStream_t) { return hipSuccess; } + hipError_t RunFence() { return hipSuccess; } +}; + +} // namespace collective +} // namespace mori + #endif // XLA_BACKENDS_GPU_COLLECTIVES_MORI_STUB_H_ diff --git a/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/write_value_thunk_folded/BUILD b/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/write_value_thunk_folded/BUILD index bdbd81d8765242..48804f49136719 100644 --- a/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/write_value_thunk_folded/BUILD +++ b/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/write_value_thunk_folded/BUILD @@ -92,6 +92,7 @@ xla_cc_binary( "//xla/service:gpu_topology", "//xla/service/gpu:nvptx_compiler", "//xla/service/gpu:nvptx_compiler_impl", # buildcleaner: keep + "//xla/stream_executor:device_description_proto_cc", "//xla/stream_executor/cuda:all_runtime", "//xla/stream_executor/cuda:cuda_platform", "//xla/stream_executor/cuda:cuda_platform_id", @@ -133,6 +134,7 @@ xla_test( "//xla:executable_run_options", "//xla:literal", "//xla:literal_util", + "//xla:xla_proto_cc", "//xla/service:compiled_module", "//xla/service:compiler", "//xla/service:executable", diff --git a/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/write_value_thunk_folded/write_value_thunk_folded_compile_main.cc b/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/write_value_thunk_folded/write_value_thunk_folded_compile_main.cc index 0a1d253f0145db..1f684945fc0920 100644 --- a/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/write_value_thunk_folded/write_value_thunk_folded_compile_main.cc +++ b/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/write_value_thunk_folded/write_value_thunk_folded_compile_main.cc @@ -36,6 +36,7 @@ limitations under the License. #include "xla/service/compiler.h" #include "xla/service/gpu_topology.h" #include "xla/stream_executor/cuda/cuda_platform_id.h" +#include "xla/stream_executor/device_description.pb.h" #include "xla/tsl/platform/env.h" #include "tsl/platform/init_main.h" diff --git a/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/write_value_thunk_folded/write_value_thunk_folded_test.cc b/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/write_value_thunk_folded/write_value_thunk_folded_test.cc index a72dd1ac49504e..8a93ac091b7e07 100644 --- a/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/write_value_thunk_folded/write_value_thunk_folded_test.cc +++ b/third_party/xla/xla/backends/gpu/libraries/native_custom_call_thunks/write_value_thunk_folded/write_value_thunk_folded_test.cc @@ -42,6 +42,7 @@ limitations under the License. #include "xla/tsl/platform/env.h" #include "xla/tsl/platform/resource_loader.h" #include "xla/tsl/platform/test.h" +#include "xla/xla.pb.h" namespace xla::gpu { namespace { diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD index 53b24c32bff7be..128145f56e2905 100644 --- a/third_party/xla/xla/backends/gpu/runtime/BUILD +++ b/third_party/xla/xla/backends/gpu/runtime/BUILD @@ -1119,6 +1119,7 @@ cc_library( "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/base:nullability", "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/container:inlined_vector", "@com_google_absl//absl/debugging:symbolize", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", @@ -2542,9 +2543,11 @@ cc_library( "@com_google_absl//absl/base", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings:str_format", "@com_google_absl//absl/strings:string_view", "@com_google_absl//absl/synchronization", "@com_google_absl//absl/types:span", @@ -5494,6 +5497,7 @@ xla_test( deps = [ ":record_ffi", "//xla:status_macros", + "//xla:xla_data_proto_cc", "//xla/ffi", "//xla/ffi:api", "//xla/ffi:call_frame", diff --git a/third_party/xla/xla/backends/gpu/runtime/all_gather.cc b/third_party/xla/xla/backends/gpu/runtime/all_gather.cc index 390ffbc5bdac1b..18ef5534a123fb 100644 --- a/third_party/xla/xla/backends/gpu/runtime/all_gather.cc +++ b/third_party/xla/xla/backends/gpu/runtime/all_gather.cc @@ -30,6 +30,7 @@ limitations under the License. #include "xla/backends/gpu/transforms/collectives/collective_ops_utils.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_instructions.h" +#include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/primitive_util.h" #include "xla/service/collective_ops_utils.h" @@ -59,6 +60,13 @@ absl::Status IsAllGatherKernelSupported(int64_t num_elements, primitive_util::LowercasePrimitiveTypeName(element_type))); } + const int64_t byte_size = + num_elements * primitive_util::ByteWidth(element_type); + if (byte_size > kMaxAllGatherSizeBytes) { + return absl::UnimplementedError( + "Custom all-gather strategy is only supported for small inputs."); + } + // The total transfer size in bits must be aligned to // kBitsPerMemoryTransaction (128 bits = 16 bytes) so each thread can // load/store a complete transaction. @@ -142,11 +150,9 @@ absl::StatusOr BuildAllGatherInfo( "Collective kernels are only supported on devices with NVLink/UALink " "support."); } - ABSL_ASSIGN_OR_RETURN(const CollectiveOpGroupMode group_mode, - GetCollectiveOpGroupMode(all_gather)); - const bool is_local = IsAllReplicasLocal( - gpu_topology.num_devices_per_process(), all_gather->replica_groups(), - group_mode, device_assignment); + ABSL_ASSIGN_OR_RETURN( + const bool is_local, + IsAllReplicasLocal(gpu_topology, *all_gather, device_assignment)); ABSL_RETURN_IF_ERROR(IsAllGatherKernelSupported( is_collective_kernel_enabled, device_info, num_operands, num_devices, num_elements, element_type, is_local, all_gather->replica_groups())); @@ -205,28 +211,37 @@ absl::StatusOr CreateAllGatherKernelSpec( const int64_t remote_size = xla::RoundUpTo(input_size_bytes, kXlaAllocatedBufferAlignBytes); + const DebugOptions& debug_options = + instr->GetModule()->config().debug_options(); + const SymmetricMemoryType sym_mem_type = + IsCrossHostOneShotKernelEnabled(debug_options, DebugOptions::ALLGATHER) + ? SymmetricMemoryType::kLoadStoreAccessible + : SymmetricMemoryType::kXlaRendezvous; + CollectiveKernelSpec kernel_spec = { - /* .input_buffer_specs= */ { - {/*requires_multimem=*/false, SymmetricMemoryType::kNone}}, - /* .output_buffer_specs= */ - {{/*requires_multimem=*/false, SymmetricMemoryType::kNone}}, + /* .codegen_config= */ { + /* .copy_input_to_scratch= */ true, + /* .emit_entry_barrier= */ true, + /* .input_buffer_specs= */ + {{/*requires_multimem=*/false, SymmetricMemoryType::kNone}}, + /* .output_buffer_specs= */ + {{/*requires_multimem=*/false, SymmetricMemoryType::kNone}}, + /* .argument_descriptors= */ + {{KernelArgType::kScratchBuffer, + /*index=*/1}, // scratch buffer as input + {KernelArgType::kOutputBuffer, /*index=*/0}, + {KernelArgType::kRuntimeRank}, + {KernelArgType::kInvocationCount}, + {KernelArgType::kScratchBuffer, + /*index=*/0}}, // signal buffers only + /* .sync_count_increment= */ 1u}, /* .scratch_buffers= */ - {{signal_size, /*requires_multimem=*/false, // Signal flags - SymmetricMemoryType::kXlaRendezvous, + {{signal_size, /*requires_multimem=*/false, sym_mem_type, /*should_memzero=*/true, /*should_double_buffer=*/true}, - {remote_size, /*requires_multimem=*/false, // Symmetric remote buffer - SymmetricMemoryType::kXlaRendezvous, + {remote_size, /*requires_multimem=*/false, sym_mem_type, /*should_memzero=*/false, - /*should_double_buffer=*/true}}, - /* .argument_descriptors= */ - {{KernelArgType::kInputBuffer, /*index=*/0}, // buffers[0].source_buffer - {KernelArgType::kOutputBuffer, /*index=*/0}, // buffers[0].dst_buffer - {KernelArgType::kRuntimeRank}, - {KernelArgType::kInvocationCount}, - {KernelArgType::kScratchBuffer, /*index=*/0}, // signal buffers - {KernelArgType::kScratchBuffer, /*index=*/1}}, // remote buffers - /* .sync_count_increment= */ 1u}; // AllGather is always one-shot + /*should_double_buffer=*/true}}}; return kernel_spec; } diff --git a/third_party/xla/xla/backends/gpu/runtime/all_gather.h b/third_party/xla/xla/backends/gpu/runtime/all_gather.h index c58e61841af7f2..82e6ebca9486aa 100644 --- a/third_party/xla/xla/backends/gpu/runtime/all_gather.h +++ b/third_party/xla/xla/backends/gpu/runtime/all_gather.h @@ -51,6 +51,10 @@ inline constexpr auto kSupportedAllGatherTypes = // is always sized to match the actual grid. inline constexpr int64_t kAllGatherMaxBlocksPerGrid = 32; +// Optimal threshold for one-shot all-gather in bytes for the collective kernel. +// Base on the experimental results. +inline constexpr int64_t kMaxAllGatherSizeBytes = 512 * 1024; // 512 KB + // Encapsulates the information needed to perform an all-gather via the Triton // collective kernel backend. struct AllGatherInfo { @@ -93,14 +97,14 @@ LaunchDimensions AllGatherLaunchDimensions( const se::DeviceDescription& device_info); // Creates a CollectiveKernelSpec describing the resource requirements of a -// Triton all-gather kernel. The returned spec uses the same 6-argument layout -// as the all-reduce kernel: -// [0] input buffer (per-rank source slice) -// [1] output buffer (full gathered destination) +// Triton all-gather kernel. The kernel argument layout is: +// [0] input/scratch buffer pointer table (kScratchBuffer, index 1) +// [1] output buffer (kOutputBuffer, index 0) // [2] runtime rank (kRuntimeRank) // [3] invocation count (kInvocationCount) -// [4] scratch index 0: signal flags (kScratchBuffer) -// [5] scratch index 1: symmetric remote buffer (kScratchBuffer) +// [4] signal flags (kScratchBuffer, index 0) +// The runtime performs a D2D copy from the input buffer to the local rank's +// scratch buffer before kernel launch (copy_input_to_scratch=true). absl::StatusOr CreateAllGatherKernelSpec( const HloInstruction* instr, const LaunchDimensions& launch_dimensions); diff --git a/third_party/xla/xla/backends/gpu/runtime/all_gather_build_info_test.cc b/third_party/xla/xla/backends/gpu/runtime/all_gather_build_info_test.cc index e68d6f7f9b3c6b..7e926b4c85d3af 100644 --- a/third_party/xla/xla/backends/gpu/runtime/all_gather_build_info_test.cc +++ b/third_party/xla/xla/backends/gpu/runtime/all_gather_build_info_test.cc @@ -213,5 +213,14 @@ TEST_F(BuildAllGatherInfoTest, FailsWithoutNvlink) { StatusIs(absl::StatusCode::kUnimplemented, HasSubstr("NVLink/UALink"))); } +TEST_F(BuildAllGatherInfoTest, FailsForLargeInputs) { + // 2 * 1024 * 1024 F32 elements = 8 MB > 4 MB limit -> unimplemented. + EXPECT_THAT( + BuildInfo(CollectiveKernelEnabled(true), F32, + /*num_elements=*/2 * 1024 * 1024, /*replica_groups=*/{0, 1}), + StatusIs(absl::StatusCode::kUnimplemented, + HasSubstr("only supported for small inputs"))); +} + } // namespace } // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/runtime/all_reduce.cc b/third_party/xla/xla/backends/gpu/runtime/all_reduce.cc index 8a4faab4fa0de0..fcb1fc6a581928 100644 --- a/third_party/xla/xla/backends/gpu/runtime/all_reduce.cc +++ b/third_party/xla/xla/backends/gpu/runtime/all_reduce.cc @@ -349,11 +349,9 @@ absl::StatusOr BuildAllReduceInfo( num_elements * primitive_util::ByteWidth(element_type); const AllReduceStrategy strategy = GetAllReduceStrategy(byte_size, is_multimem_enabled); - ABSL_ASSIGN_OR_RETURN(const CollectiveOpGroupMode group_mode, - GetCollectiveOpGroupMode(all_reduce)); - const bool is_local = IsAllReplicasLocal( - gpu_topology.num_devices_per_process(), all_reduce->replica_groups(), - group_mode, device_assignment); + ABSL_ASSIGN_OR_RETURN( + const bool is_local, + IsAllReplicasLocal(gpu_topology, *all_reduce, device_assignment)); if (device_info.device_interconnect_info().active_links <= 0) { return absl::UnimplementedError( "Collective kernels are only supported on devices with NVLink/UALink " @@ -457,10 +455,23 @@ absl::StatusOr CreateAllReduceKernelSpec( : SymmetricMemoryType::kXlaRendezvous; CollectiveKernelSpec kernel_spec = { - /* .input_buffer_specs= */ { - {/*requires_multimem=*/false, SymmetricMemoryType::kNone}}, - /* .output_buffer_specs= */ - {{/*requires_multimem=*/false, SymmetricMemoryType::kNone}}, + /* .codegen_config= */ { + /* .copy_input_to_scratch= */ false, + /* .emit_entry_barrier= */ false, + /* .input_buffer_specs= */ + {{/*requires_multimem=*/false, SymmetricMemoryType::kNone}}, + /* .output_buffer_specs= */ + {{/*requires_multimem=*/false, SymmetricMemoryType::kNone}}, + /* .argument_descriptors= */ + {{KernelArgType::kInputBuffer, + /*index=*/0}, // buffers[0].source_buffer + {KernelArgType::kOutputBuffer, + /*index=*/0}, // buffers[0].dst_buffer + {KernelArgType::kRuntimeRank}, + {KernelArgType::kInvocationCount}, + {KernelArgType::kScratchBuffer, /*index=*/0}, // signal buffers + {KernelArgType::kScratchBuffer, /*index=*/1}}, // scratch buffers + /* .sync_count_increment = */ 1 + static_cast(strategy)}, /* .scratch_buffers= */ {{signal_size, /*requires_multimem=*/false, // Signal buffers sym_mem_type, @@ -469,15 +480,7 @@ absl::StatusOr CreateAllReduceKernelSpec( {remote_size, /*requires_multimem=*/false, // Remote buffers sym_mem_type, /*should_memzero=*/false, - /*should_double_buffer=*/true}}, - /* .argument_descriptors= */ - {{KernelArgType::kInputBuffer, /*index=*/0}, // buffers[0].source_buffer - {KernelArgType::kOutputBuffer, /*index=*/0}, // buffers[0].dst_buffer - {KernelArgType::kRuntimeRank}, - {KernelArgType::kInvocationCount}, - {KernelArgType::kScratchBuffer, /*index=*/0}, // signal buffers - {KernelArgType::kScratchBuffer, /*index=*/1}}, // scratch buffers - /* .sync_count_increment = */ 1 + static_cast(strategy)}; + /*should_double_buffer=*/true}}}; return kernel_spec; } diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_broadcast_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/collective_broadcast_thunk.cc index 4c47b30f77ff82..1cd55cf0c24c75 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_broadcast_thunk.cc +++ b/third_party/xla/xla/backends/gpu/runtime/collective_broadcast_thunk.cc @@ -22,9 +22,11 @@ limitations under the License. #include #include "absl/base/casts.h" +#include "absl/log/check.h" #include "absl/status/status.h" #include "absl/status/status_macros.h" #include "absl/status/statusor.h" +#include "absl/strings/str_format.h" #include "absl/synchronization/mutex.h" #include "absl/types/span.h" #include "xla/backends/gpu/collectives/gpu_clique_key.h" diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_broadcast_thunk.h b/third_party/xla/xla/backends/gpu/runtime/collective_broadcast_thunk.h index a899d033cfa1e2..83ec456ffaf06e 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_broadcast_thunk.h +++ b/third_party/xla/xla/backends/gpu/runtime/collective_broadcast_thunk.h @@ -1,6 +1,7 @@ #include #include "absl/strings/string_view.h" +#include "absl/synchronization/mutex.h" #include "xla/backends/gpu/collectives/gpu_clique_key.h" #include "xla/backends/gpu/runtime/thunk.pb.h" #include "xla/runtime/buffer_use.h" diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_broadcast_thunk_test.cc b/third_party/xla/xla/backends/gpu/runtime/collective_broadcast_thunk_test.cc index 525bc0bbb4bca8..f393334e5690e0 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_broadcast_thunk_test.cc +++ b/third_party/xla/xla/backends/gpu/runtime/collective_broadcast_thunk_test.cc @@ -257,6 +257,12 @@ ENTRY test_computation { se::StreamExecutor* executor = backend().default_stream_executor(); + if (executor->GetDeviceDescription() + .gpu_compute_capability() + .oneapi_compute_capability()) { + GTEST_SKIP() << "oneAPI command buffers are not implemented yet."; + } + ASSERT_OK_AND_ASSIGN( std::unique_ptr compiled_module, backend().compiler()->RunHloPasses(module->Clone(), executor, diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_clique_requests_test.cc b/third_party/xla/xla/backends/gpu/runtime/collective_clique_requests_test.cc index b8ce64af5bcf6c..f292eaae560257 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_clique_requests_test.cc +++ b/third_party/xla/xla/backends/gpu/runtime/collective_clique_requests_test.cc @@ -15,7 +15,6 @@ limitations under the License. #include "xla/backends/gpu/runtime/collective_clique_requests.h" -#include #include #include "absl/algorithm/container.h" diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.cc index a14258fc9d372c..cd61489fdfa30f 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.cc +++ b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.cc @@ -172,7 +172,7 @@ absl::StatusOr> BuildKernelArguments( const se::DeviceAddressBase metadata, const GpuCliqueKey& clique_key, int32_t num_parameters) { std::vector kernel_args; - kernel_args.reserve(kernel_spec.argument_descriptors.size()); + kernel_args.reserve(kernel_spec.codegen_config.argument_descriptors.size()); auto get_buffer_index = [](std::optional index, int32_t num_buffers) -> absl::StatusOr { TF_RET_CHECK(index.has_value() && *index >= 0 && *index < num_buffers) @@ -186,10 +186,13 @@ absl::StatusOr> BuildKernelArguments( // allocations. static constexpr auto count_multimem_buffers = [](const IoBufferSpec& spec) -> bool { return spec.requires_multimem; }; - return absl::c_count_if(spec.input_buffer_specs, count_multimem_buffers) + - absl::c_count_if(spec.output_buffer_specs, count_multimem_buffers); + return absl::c_count_if(spec.codegen_config.input_buffer_specs, + count_multimem_buffers) + + absl::c_count_if(spec.codegen_config.output_buffer_specs, + count_multimem_buffers); }(kernel_spec); - for (const KernelArgDescriptor& desc : kernel_spec.argument_descriptors) { + for (const KernelArgDescriptor& desc : + kernel_spec.codegen_config.argument_descriptors) { switch (desc.type) { case KernelArgType::kInputBuffer: { ABSL_ASSIGN_OR_RETURN(const int32_t buffer_index, @@ -231,12 +234,12 @@ absl::StatusOr> BuildKernelArguments( } bool RequiresMultimem(CollectiveKernelSpec kernel_spec) { - for (const auto& op : kernel_spec.input_buffer_specs) { + for (const auto& op : kernel_spec.codegen_config.input_buffer_specs) { if (op.requires_multimem) { return true; } } - for (const auto& res : kernel_spec.output_buffer_specs) { + for (const auto& res : kernel_spec.codegen_config.output_buffer_specs) { if (res.requires_multimem) { return true; } @@ -296,14 +299,18 @@ absl::Status CollectiveKernelThunk::IsSupported( return absl::FailedPreconditionError( absl::StrFormat("Empty kernel name ('%s')", kernel_name_)); } - // Check if peer access is supported for all devices in the clique. + // Check if peer access is supported for all devices in the clique managed by + // the current process. for (const GlobalDeviceId& device : clique_key.devices()) { - ABSL_ASSIGN_OR_RETURN(const int peer_device_id, - GetLocalDeviceId(device, collective_params)); - if (!executor.CanEnablePeerAccessTo(peer_device_id)) { + auto peer_device_id = GetLocalDeviceId(device, collective_params); + if (!peer_device_id.ok()) { + // Device is managed by a different process in multi-process setups. + continue; + } + if (!executor.CanEnablePeerAccessTo(*peer_device_id)) { return absl::FailedPreconditionError(absl::StrFormat( "Peer access is not supported from device %d to device %d", - executor.device_ordinal(), peer_device_id)); + executor.device_ordinal(), *peer_device_id)); } } return absl::OkStatus(); @@ -321,29 +328,34 @@ absl::Status CollectiveKernelThunk::Prepare(const PrepareParams& params) { IsSupported(clique_key, *params.executor, *params.collective_params)); // Validate that the kernel spec is compatible with the thunk buffers. - TF_RET_CHECK(kernel_spec_.input_buffer_specs.size() == buffers_.size()) + TF_RET_CHECK(kernel_spec_.codegen_config.input_buffer_specs.size() == + buffers_.size()) << "Kernel spec input_buffer_specs size (" - << kernel_spec_.input_buffer_specs.size() + << kernel_spec_.codegen_config.input_buffer_specs.size() << ") must equal thunk buffers size (" << buffers_.size() << ")"; - TF_RET_CHECK(kernel_spec_.output_buffer_specs.size() == buffers_.size()) + TF_RET_CHECK(kernel_spec_.codegen_config.output_buffer_specs.size() == + buffers_.size()) << "Kernel spec output_buffer_specs size (" - << kernel_spec_.output_buffer_specs.size() + << kernel_spec_.codegen_config.output_buffer_specs.size() << ") must equal thunk buffers size (" << buffers_.size() << ")"; - for (const KernelArgDescriptor& desc : kernel_spec_.argument_descriptors) { + for (const KernelArgDescriptor& desc : + kernel_spec_.codegen_config.argument_descriptors) { switch (desc.type) { case KernelArgType::kInputBuffer: - TF_RET_CHECK(desc.index.has_value() && - desc.index.value() < - kernel_spec_.input_buffer_specs.size() && - desc.index.value() < buffers_.size()) + TF_RET_CHECK( + desc.index.has_value() && + desc.index.value() < + kernel_spec_.codegen_config.input_buffer_specs.size() && + desc.index.value() < buffers_.size()) << "Invalid input buffer argument index: " << desc.index.value_or(-999); break; case KernelArgType::kOutputBuffer: - TF_RET_CHECK(desc.index.has_value() && - desc.index.value() < - kernel_spec_.output_buffer_specs.size() && - desc.index.value() < buffers_.size()) + TF_RET_CHECK( + desc.index.has_value() && + desc.index.value() < + kernel_spec_.codegen_config.output_buffer_specs.size() && + desc.index.value() < buffers_.size()) << "Invalid output buffer argument index: " << desc.index.value_or(-999); break; @@ -392,16 +404,20 @@ absl::Status CollectiveKernelThunk::Prepare(const PrepareParams& params) { // If we decided to run kernel using multimem strategy we request symmetric // memory for buffers that explicitly requested it. - for (size_t i = 0; i < kernel_spec_.input_buffer_specs.size(); ++i) { - if (kernel_spec_.input_buffer_specs[i].symmetric_memory_type == + for (size_t i = 0; + i < kernel_spec_.codegen_config.input_buffer_specs.size(); ++i) { + if (kernel_spec_.codegen_config.input_buffer_specs[i] + .symmetric_memory_type == SymmetricMemoryType::kLoadStoreAccessible) { ABSL_RETURN_IF_ERROR( params.collective_memory_requests->RequestSymmetricAllocation( clique_key, buffers_[i].source_buffer.slice.index())); } } - for (size_t i = 0; i < kernel_spec_.output_buffer_specs.size(); ++i) { - if (kernel_spec_.output_buffer_specs[i].symmetric_memory_type == + for (size_t i = 0; + i < kernel_spec_.codegen_config.output_buffer_specs.size(); ++i) { + if (kernel_spec_.codegen_config.output_buffer_specs[i] + .symmetric_memory_type == SymmetricMemoryType::kLoadStoreAccessible) { ABSL_RETURN_IF_ERROR( params.collective_memory_requests->RequestSymmetricAllocation( @@ -447,7 +463,8 @@ absl::Status CollectiveKernelThunk::Initialize(const InitializeParams& params) { << "Kernel name must be set for collective kernel thunk."; // Create kernel for execution. std::unique_ptr kernel = nullptr; - const int32_t num_args = kernel_spec_.argument_descriptors.size(); + const int32_t num_args = + kernel_spec_.codegen_config.argument_descriptors.size(); if (cubin_.has_value()) { ABSL_ASSIGN_OR_RETURN(kernel, CreateKernel(kernel_name_, num_args, *cubin_, params.executor, shmem_bytes_)); @@ -478,15 +495,18 @@ absl::Status CollectiveKernelThunk::Initialize(const InitializeParams& params) { if (state != nullptr) { std::vector parameters; - parameters.reserve(kernel_spec_.argument_descriptors.size()); - for (size_t i = 0; i < kernel_spec_.input_buffer_specs.size(); ++i) { - if (kernel_spec_.input_buffer_specs[i].requires_multimem) { + parameters.reserve(kernel_spec_.codegen_config.argument_descriptors.size()); + for (size_t i = 0; + i < kernel_spec_.codegen_config.input_buffer_specs.size(); ++i) { + if (kernel_spec_.codegen_config.input_buffer_specs[i].requires_multimem) { parameters.push_back(params.buffer_allocations->GetDeviceAddress( buffers_[i].source_buffer.slice)); } } - for (size_t i = 0; i < kernel_spec_.output_buffer_specs.size(); ++i) { - if (kernel_spec_.output_buffer_specs[i].requires_multimem) { + for (size_t i = 0; + i < kernel_spec_.codegen_config.output_buffer_specs.size(); ++i) { + if (kernel_spec_.codegen_config.output_buffer_specs[i] + .requires_multimem) { parameters.push_back(params.buffer_allocations->GetDeviceAddress( buffers_[i].destination_buffer.slice)); } @@ -557,9 +577,9 @@ absl::Status CollectiveKernelThunk::Initialize(const InitializeParams& params) { return spec.requires_multimem; }; int32_t scratch_buffers_index = - absl::c_count_if(kernel_spec_.input_buffer_specs, + absl::c_count_if(kernel_spec_.codegen_config.input_buffer_specs, is_multimem_buffer) + - absl::c_count_if(kernel_spec_.output_buffer_specs, + absl::c_count_if(kernel_spec_.codegen_config.output_buffer_specs, is_multimem_buffer); param_to_peers_ptrs.resize(num_parameters * clique_key.num_devices()); for (size_t i = scratch_buffers_index; i < num_parameters; ++i) { @@ -631,7 +651,7 @@ absl::Status CollectiveKernelThunk::ExecuteOnStream( state = it->second.get(); } - state->invocation_count += kernel_spec_.sync_count_increment; + state->invocation_count += kernel_spec_.codegen_config.sync_count_increment; TF_RET_CHECK(state->kernel != nullptr) << "Kernel is not initialized for collective kernel thunk."; @@ -639,8 +659,10 @@ absl::Status CollectiveKernelThunk::ExecuteOnStream( return buffer_spec.requires_multimem; }; int32_t num_parameters = - absl::c_count_if(kernel_spec_.input_buffer_specs, has_multimem) + - absl::c_count_if(kernel_spec_.output_buffer_specs, has_multimem) + + absl::c_count_if(kernel_spec_.codegen_config.input_buffer_specs, + has_multimem) + + absl::c_count_if(kernel_spec_.codegen_config.output_buffer_specs, + has_multimem) + kernel_spec_.scratch_buffers.size(); ABSL_ASSIGN_OR_RETURN( @@ -649,6 +671,32 @@ absl::Status CollectiveKernelThunk::ExecuteOnStream( state->invocation_count, state->metadata, clique_key, num_parameters)); + // For AllGather: copy input buffer → local rank's symmetric scratch buffer + // before kernel launch. All ranks perform this copy, then the entry barrier + // in the kernel waits until every rank has finished its copy. + if (kernel_spec_.codegen_config.copy_input_to_scratch) { + StreamMemory* memory_state = nullptr; + { + absl::MutexLock lock(mutex_); + memory_state = per_stream_memory_.at(stream->parent()).get(); + } + // scratch_allocations[1] is the symmetric remote buffer that holds this + // rank's input slice. scratch_allocations[0] is the signal buffer. + TF_RET_CHECK(memory_state->scratch_allocations.size() >= 2) + << "Expected at least 2 scratch allocations for D2D copy."; + se::DeviceAddressBase input_addr = + params.buffer_allocations->GetDeviceAddress( + buffers_[0].source_buffer.slice); + const int64_t copy_size = GetInputSizeBytes(); + ABSL_RETURN_IF_ERROR( + stream->Memcpy(memory_state->scratch_allocations[1].address_ptr(), + input_addr, copy_size)); + VLOG(3) << "D2D copy: src=" << input_addr.opaque() << " dst=" + << memory_state->scratch_allocations[1].address().opaque() + << " size=" << copy_size + << " invocation=" << state->invocation_count; + } + return ExecuteKernelOnStream(*state->kernel, kernel_args, launch_dimensions_, /*cluster_dim=*/std::nullopt, stream); } @@ -682,24 +730,24 @@ CollectiveKernelThunk::FromProto( CollectiveKernelSpec kernel_spec; if (thunk_proto.has_kernel_spec()) { const CollectiveKernelSpecProto& proto_spec = thunk_proto.kernel_spec(); - kernel_spec.input_buffer_specs.reserve( + kernel_spec.codegen_config.input_buffer_specs.reserve( proto_spec.input_buffer_specs_size()); for (const auto& input : proto_spec.input_buffer_specs()) { TF_RET_CHECK( SymmetricMemoryTypeProto_IsValid(input.symmetric_memory_type())) << "Invalid symmetric_memory_type: " << input.symmetric_memory_type(); - kernel_spec.input_buffer_specs.push_back( + kernel_spec.codegen_config.input_buffer_specs.push_back( {input.requires_multimem(), static_cast(input.symmetric_memory_type())}); } - kernel_spec.output_buffer_specs.reserve( + kernel_spec.codegen_config.output_buffer_specs.reserve( proto_spec.output_buffer_specs_size()); for (const auto& output : proto_spec.output_buffer_specs()) { TF_RET_CHECK( SymmetricMemoryTypeProto_IsValid(output.symmetric_memory_type())) << "Invalid symmetric_memory_type: " << output.symmetric_memory_type(); - kernel_spec.output_buffer_specs.push_back( + kernel_spec.codegen_config.output_buffer_specs.push_back( {output.requires_multimem(), static_cast(output.symmetric_memory_type())}); } @@ -714,7 +762,7 @@ CollectiveKernelThunk::FromProto( static_cast(scratch.symmetric_memory_type()), scratch.should_memzero(), scratch.should_double_buffer()}); } - kernel_spec.argument_descriptors.reserve( + kernel_spec.codegen_config.argument_descriptors.reserve( proto_spec.argument_descriptors_size()); for (const auto& arg : proto_spec.argument_descriptors()) { TF_RET_CHECK(KernelArgTypeProto_IsValid(arg.type())) @@ -726,16 +774,20 @@ CollectiveKernelThunk::FromProto( << "Invalid argument index: " << arg.index(); arg_desc.index = arg.index(); } - kernel_spec.argument_descriptors.push_back(std::move(arg_desc)); + kernel_spec.codegen_config.argument_descriptors.push_back( + std::move(arg_desc)); } - kernel_spec.sync_count_increment = proto_spec.invocation_count_increment(); + kernel_spec.codegen_config.sync_count_increment = + proto_spec.invocation_count_increment(); + kernel_spec.codegen_config.copy_input_to_scratch = + proto_spec.copy_input_to_scratch(); } else { // Backward-compatibility fallback for legacy AOT-compiled kernels without // an explicit CollectiveKernelSpec. // Can be removed in February 2027 (6months backward compatibility window). - kernel_spec.input_buffer_specs.push_back( + kernel_spec.codegen_config.input_buffer_specs.push_back( {/*requires_multimem=*/false, SymmetricMemoryType::kNone}); - kernel_spec.output_buffer_specs.push_back( + kernel_spec.codegen_config.output_buffer_specs.push_back( {/*requires_multimem=*/false, SymmetricMemoryType::kNone}); TF_RET_CHECK(thunk_proto.buffers_size() > 0) << "At least one buffer is required for collective kernel thunk."; @@ -766,14 +818,14 @@ CollectiveKernelThunk::FromProto( /*symmetric_memory_type=*/SymmetricMemoryType::kXlaRendezvous, /*should_memzero=*/false, /*should_double_buffer=*/true}}; - kernel_spec.argument_descriptors = { + kernel_spec.codegen_config.argument_descriptors = { {KernelArgType::kInputBuffer, /*index=*/0}, {KernelArgType::kOutputBuffer, /*index=*/0}, {KernelArgType::kRuntimeRank}, {KernelArgType::kInvocationCount}, {KernelArgType::kScratchBuffer, /*index=*/0}, {KernelArgType::kScratchBuffer, /*index=*/1}}; - kernel_spec.sync_count_increment = + kernel_spec.codegen_config.sync_count_increment = 1 + static_cast(GetAllReduceStrategy( input_size_bytes, /*is_multimem_enabled=*/false)); } @@ -809,16 +861,16 @@ absl::StatusOr CollectiveKernelThunk::ToProto() const { auto* proto_spec = thunk_proto->mutable_kernel_spec(); proto_spec->mutable_input_buffer_specs()->Reserve( - kernel_spec_.input_buffer_specs.size()); - for (const auto& op : kernel_spec_.input_buffer_specs) { + kernel_spec_.codegen_config.input_buffer_specs.size()); + for (const auto& op : kernel_spec_.codegen_config.input_buffer_specs) { auto* io_proto = proto_spec->add_input_buffer_specs(); io_proto->set_requires_multimem(op.requires_multimem); io_proto->set_symmetric_memory_type( static_cast(op.symmetric_memory_type)); } proto_spec->mutable_output_buffer_specs()->Reserve( - kernel_spec_.output_buffer_specs.size()); - for (const auto& res : kernel_spec_.output_buffer_specs) { + kernel_spec_.codegen_config.output_buffer_specs.size()); + for (const auto& res : kernel_spec_.codegen_config.output_buffer_specs) { auto* io_proto = proto_spec->add_output_buffer_specs(); io_proto->set_requires_multimem(res.requires_multimem); io_proto->set_symmetric_memory_type( @@ -836,15 +888,18 @@ absl::StatusOr CollectiveKernelThunk::ToProto() const { scratch_proto->set_should_double_buffer(scratch.should_double_buffer); } proto_spec->mutable_argument_descriptors()->Reserve( - kernel_spec_.argument_descriptors.size()); - for (const auto& arg : kernel_spec_.argument_descriptors) { + kernel_spec_.codegen_config.argument_descriptors.size()); + for (const auto& arg : kernel_spec_.codegen_config.argument_descriptors) { auto* arg_proto = proto_spec->add_argument_descriptors(); arg_proto->set_type(static_cast(arg.type)); if (arg.index.has_value()) { arg_proto->set_index(arg.index.value()); } } - proto_spec->set_invocation_count_increment(kernel_spec_.sync_count_increment); + proto_spec->set_invocation_count_increment( + kernel_spec_.codegen_config.sync_count_increment); + proto_spec->set_copy_input_to_scratch( + kernel_spec_.codegen_config.copy_input_to_scratch); for (const CollectiveThunk::Buffer& buffer : buffers_) { ABSL_ASSIGN_OR_RETURN(*thunk_proto->add_buffers(), buffer.ToProto()); diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.proto b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.proto index b340cdac918324..3d1ae344248142 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.proto +++ b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk.proto @@ -62,6 +62,7 @@ message CollectiveKernelSpecProto { repeated ScratchBufferSpecProto scratch_buffers = 3; repeated KernelArgDescriptorProto argument_descriptors = 4; int32 invocation_count_increment = 5; + bool copy_input_to_scratch = 6; } message CollectiveKernelThunkProto { diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk_test.cc b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk_test.cc index 0270630937bd49..5f16b3a0899220 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk_test.cc +++ b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk_test.cc @@ -187,10 +187,21 @@ CollectiveKernelSpec CreateCollectiveKernelSpec( ? SymmetricMemoryType::kLoadStoreAccessible : SymmetricMemoryType::kXlaRendezvous; return { - /*operand_buffer_specs=*/{ - {/*requires_multimem=*/false, SymmetricMemoryType::kNone}}, - /*result_buffer_specs=*/ - {{/*requires_multimem=*/false, SymmetricMemoryType::kNone}}, + /*codegen_config=*/{ + /*copy_input_to_scratch=*/false, + /*emit_entry_barrier=*/false, + /*input_buffer_specs=*/ + {{/*requires_multimem=*/false, SymmetricMemoryType::kNone}}, + /*output_buffer_specs=*/ + {{/*requires_multimem=*/false, SymmetricMemoryType::kNone}}, + /*argument_descriptors=*/ + {{KernelArgType::kInputBuffer, 0}, + {KernelArgType::kOutputBuffer, 0}, + {KernelArgType::kRuntimeRank}, + {KernelArgType::kInvocationCount}, + {KernelArgType::kScratchBuffer, 0}, + {KernelArgType::kScratchBuffer, 1}}, + }, /*scratch_buffers=*/ {{signal_size, /*requires_multimem=*/false, sym_mem_type, /*should_memzero=*/true, @@ -199,13 +210,6 @@ CollectiveKernelSpec CreateCollectiveKernelSpec( /*requires_multimem=*/is_multimem_enabled, sym_mem_type, /*should_memzero=*/false, /*should_double_buffer=*/true}}, - /*argument_descriptors=*/ - {{KernelArgType::kInputBuffer, 0}, - {KernelArgType::kOutputBuffer, 0}, - {KernelArgType::kRuntimeRank}, - {KernelArgType::kInvocationCount}, - {KernelArgType::kScratchBuffer, 0}, - {KernelArgType::kScratchBuffer, 1}}, }; } diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_params.cc b/third_party/xla/xla/backends/gpu/runtime/collective_params.cc index 430baceec7d214..e48cf1e07f26cf 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_params.cc +++ b/third_party/xla/xla/backends/gpu/runtime/collective_params.cc @@ -19,12 +19,10 @@ limitations under the License. #include #include -#include "absl/base/casts.h" #include "absl/container/flat_hash_map.h" #include "absl/log/check.h" #include "absl/status/status_macros.h" #include "absl/status/statusor.h" -#include "absl/strings/string_view.h" #include "absl/types/span.h" #include "xla/backends/gpu/collectives/gpu_collectives.h" #include "xla/core/collectives/collectives.h" diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_params.h b/third_party/xla/xla/backends/gpu/runtime/collective_params.h index 6d8fb0fe2ec774..b6d21c4d5624a1 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_params.h +++ b/third_party/xla/xla/backends/gpu/runtime/collective_params.h @@ -158,15 +158,20 @@ struct KernelArgDescriptor { std::optional index = std::nullopt; }; -// This structure contains the information required to configure and launch a -// custom collective kernel. -struct CollectiveKernelSpec { +// Lightweight codegen-time configuration for collective kernels. +// Built from HLO instruction properties alone, without launch dimensions. +struct CollectiveCodegenConfig { + // If true, the runtime copies the input buffer to the local rank's scratch + // buffer before kernel launch. The kernel receives the scratch buffer as + // its input argument. + bool copy_input_to_scratch = false; + // If true, a cross-rank barrier is emitted before the tile loop. The barrier + // waits until all ranks have populated the symmetric scratch buffers. + bool emit_entry_barrier = false; // Specs for input operand buffers. std::vector input_buffer_specs; // Specs for output result buffers. std::vector output_buffer_specs; - // Specs for scratch buffers these are allocated by the thunk. - std::vector scratch_buffers; // Argument descriptors that determine how the kernel is invoked. std::vector argument_descriptors; // Each time ExecuteOnStream is called, the invocation count is incremented by @@ -175,6 +180,14 @@ struct CollectiveKernelSpec { uint32_t sync_count_increment = 1; }; +// This structure contains the information required to configure and launch a +// custom collective kernel. +struct CollectiveKernelSpec { + CollectiveCodegenConfig codegen_config; + // Specs for scratch buffers these are allocated by the thunk. + std::vector scratch_buffers; +}; + } // namespace xla::gpu #endif // XLA_BACKENDS_GPU_RUNTIME_COLLECTIVE_PARAMS_H_ diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_permute_thunk_test.cc b/third_party/xla/xla/backends/gpu/runtime/collective_permute_thunk_test.cc index 05de71e24e20d4..d7e8180f48a1d9 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_permute_thunk_test.cc +++ b/third_party/xla/xla/backends/gpu/runtime/collective_permute_thunk_test.cc @@ -272,6 +272,12 @@ ENTRY test_computation { se::StreamExecutor* executor = backend().default_stream_executor(); + if (executor->GetDeviceDescription() + .gpu_compute_capability() + .oneapi_compute_capability()) { + GTEST_SKIP() << "oneAPI command buffers are not implemented yet."; + } + ASSERT_OK_AND_ASSIGN( std::unique_ptr compiled_module, backend().compiler()->RunHloPasses(module->Clone(), executor, diff --git a/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass.cc b/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass.cc index bb48404292e764..35923cdd91ac13 100644 --- a/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass.cc +++ b/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass.cc @@ -123,6 +123,15 @@ CommandBufferConfig GetCommandBufferConfig( std::move(commands), std::move(enabled_collectives), device_info, debug_options.xla_gpu_command_buffer_unroll_loops(), num_local_devices}; + // oneAPI command buffers are not implemented yet. Hence, disable command + // buffer conversion for the oneAPI backend. + // TODO(intel-tf): Remove this fallback once oneAPI command buffers are + // implemented. + if (device_info.gpu_compute_capability().IsOneAPI()) { + config.enabled_commands.clear(); + return config; + } + // Erase command buffer cmd types that are not supported by the gpu runtime. static constexpr auto kRequireConditionals = {DebugOptions::CONDITIONAL, DebugOptions::WHILE}; diff --git a/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk.cc index bafe141e69d927..2a2bc8f7f20a3b 100644 --- a/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk.cc +++ b/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk.cc @@ -17,6 +17,7 @@ limitations under the License. #include #include +#include #include #include #include @@ -27,6 +28,7 @@ limitations under the License. #include "absl/algorithm/container.h" #include "absl/base/nullability.h" #include "absl/container/flat_hash_map.h" +#include "absl/container/inlined_vector.h" #include "absl/debugging/symbolize.h" #include "absl/log/check.h" #include "absl/log/log.h" diff --git a/third_party/xla/xla/backends/gpu/runtime/record_ffi_test.cc b/third_party/xla/xla/backends/gpu/runtime/record_ffi_test.cc index 0d0393ed8a0ce6..eca91c3fc50cde 100644 --- a/third_party/xla/xla/backends/gpu/runtime/record_ffi_test.cc +++ b/third_party/xla/xla/backends/gpu/runtime/record_ffi_test.cc @@ -45,6 +45,7 @@ limitations under the License. #include "xla/stream_executor/stream.h" #include "xla/stream_executor/stream_executor.h" #include "xla/tsl/platform/test.h" +#include "xla/xla_data.pb.h" namespace xla::gpu { namespace { diff --git a/third_party/xla/xla/backends/gpu/tests/BUILD b/third_party/xla/xla/backends/gpu/tests/BUILD index 6fbd1a47698d7c..2d00e528ee647a 100644 --- a/third_party/xla/xla/backends/gpu/tests/BUILD +++ b/third_party/xla/xla/backends/gpu/tests/BUILD @@ -296,6 +296,7 @@ xla_test( deps = [ ":hlo_pjrt_gpu_test_base", "//xla:error_spec", + "//xla:xla_proto_cc", "//xla/hlo/ir:hlo", "//xla/hlo/testlib:filecheck", "//xla/stream_executor:device_description", @@ -1686,6 +1687,144 @@ xla_test( ], ) +xla_test( + name = "all_gather_e2e_test", + srcs = ["all_gather_e2e_test.cc"], + backend_tags = { + "gpu": [ + "multi_gpu", + ], + }, + backends = [ + "gpu", + ], + deps = [ + ":collective_ops_e2e_test_base", + "//xla:literal", + "//xla:literal_util", + "//xla:xla_proto_cc", + "//xla/hlo/ir:hlo", + "//xla/service/gpu:backend_configs_cc", + "//xla/tests:literal_test_util", + "//xla/tsl/platform:test", + "@com_google_absl//absl/strings:string_view", + "@com_google_googletest//:gtest_main", + ], +) + +xla_test( + name = "all_gather_e2e_multiprocess_test", + srcs = ["all_gather_e2e_multiprocess_test.cc"], + backend_tags = { + "gpu": [ + "multi_gpu", + ], + }, + backends = [ + "gpu", + ], + deps = [ + "//xla:debug_options_flags", + "//xla:literal", + "//xla:literal_util", + "//xla:shape_util", + "//xla:status_macros", + "//xla:xla_data_proto_cc", + "//xla:xla_proto_cc", + "//xla/hlo/builder:xla_computation", + "//xla/hlo/ir:hlo", + "//xla/hlo/parser:hlo_parser", + "//xla/pjrt:pjrt_client", + "//xla/pjrt:pjrt_compiler", + "//xla/pjrt:pjrt_executable", + "//xla/pjrt/distributed", + "//xla/pjrt/distributed:client", + "//xla/pjrt/distributed:service", + "//xla/pjrt/plugin/xla_gpu:xla_gpu_allocator_config", + "//xla/pjrt/plugin/xla_gpu:xla_gpu_client_options", + "//xla/pjrt/plugin/xla_gpu:xla_gpu_pjrt_client", + "//xla/service:device_assignment", + "//xla/service:gpu_topology", + "//xla/service:platform_util", + "//xla/service/gpu:backend_configs_cc", + "//xla/stream_executor:device_description", + "//xla/stream_executor:platform", + "//xla/stream_executor:stream_executor_h", + "//xla/stream_executor/cuda:cuda_compute_capability", + "//xla/tests:literal_test_util", + "//xla/tsl/platform:env", + "//xla/tsl/platform:subprocess", + "//xla/tsl/platform:test", + "//xla/tsl/util:command_line_flags", + "@com_google_absl//absl/log", + "@com_google_absl//absl/log:check", + "@com_google_absl//absl/log:flags", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings:str_format", + "@com_google_absl//absl/strings:string_view", + "@com_google_absl//absl/time", + "@com_google_googletest//:gtest", + "@tsl//tsl/platform:path", + ], +) + +xla_test( + name = "all_reduce_multi_process_e2e_test", + srcs = ["all_reduce_multi_process_e2e_test.cc"], + backend_tags = { + "gpu": [ + "multi_gpu", + ], + }, + backends = [ + "gpu", + ], + deps = [ + "//xla:debug_options_flags", + "//xla:literal", + "//xla:literal_util", + "//xla:shape_util", + "//xla:status_macros", + "//xla:xla_data_proto_cc", + "//xla:xla_proto_cc", + "//xla/hlo/builder:xla_computation", + "//xla/hlo/ir:hlo", + "//xla/hlo/parser:hlo_parser", + "//xla/pjrt:pjrt_client", + "//xla/pjrt:pjrt_compiler", + "//xla/pjrt:pjrt_executable", + "//xla/pjrt/distributed", + "//xla/pjrt/distributed:client", + "//xla/pjrt/distributed:service", + "//xla/pjrt/plugin/xla_gpu:xla_gpu_allocator_config", + "//xla/pjrt/plugin/xla_gpu:xla_gpu_client_options", + "//xla/pjrt/plugin/xla_gpu:xla_gpu_pjrt_client", + "//xla/service:device_assignment", + "//xla/service:gpu_topology", + "//xla/service:platform_util", + "//xla/service/gpu:backend_configs_cc", + "//xla/stream_executor:device_description", + "//xla/stream_executor:platform", + "//xla/stream_executor:stream_executor_h", + "//xla/stream_executor/cuda:cuda_compute_capability", + "//xla/tests:literal_test_util", + "//xla/tsl/platform:env", + "//xla/tsl/platform:subprocess", + "//xla/tsl/platform:test", + "//xla/tsl/util:command_line_flags", + "@com_google_absl//absl/log", + "@com_google_absl//absl/log:check", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings:str_format", + "@com_google_absl//absl/strings:string_view", + "@com_google_absl//absl/time", + "@com_google_googletest//:gtest", + "@tsl//tsl/platform:path", + ], +) + xla_cc_test( name = "host_cross_compilation_test", srcs = ["host_cross_compilation_test.cc"], diff --git a/third_party/xla/xla/backends/gpu/tests/all_gather_e2e_multiprocess_test.cc b/third_party/xla/xla/backends/gpu/tests/all_gather_e2e_multiprocess_test.cc new file mode 100644 index 00000000000000..f4c92db7806df3 --- /dev/null +++ b/third_party/xla/xla/backends/gpu/tests/all_gather_e2e_multiprocess_test.cc @@ -0,0 +1,373 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include +#include +#include +#include +#include +#include +#include + +#include +#include "absl/log/check.h" +#include "absl/log/log.h" +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "absl/strings/str_format.h" +#include "absl/strings/string_view.h" +#include "absl/time/time.h" +#include "xla/debug_options_flags.h" +#include "xla/hlo/builder/xla_computation.h" +#include "xla/hlo/ir/hlo_computation.h" +#include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/ir/hlo_module.h" +#include "xla/hlo/ir/hlo_opcode.h" +#include "xla/hlo/parser/hlo_parser.h" +#include "xla/literal.h" +#include "xla/literal_util.h" +#include "xla/pjrt/distributed/client.h" +#include "xla/pjrt/distributed/distributed.h" +#include "xla/pjrt/distributed/service.h" +#include "xla/pjrt/pjrt_client.h" +#include "xla/pjrt/pjrt_compiler.h" +#include "xla/pjrt/pjrt_executable.h" +#include "xla/pjrt/plugin/xla_gpu/xla_gpu_allocator_config.h" +#include "xla/pjrt/plugin/xla_gpu/xla_gpu_client_options.h" +#include "xla/pjrt/plugin/xla_gpu/xla_gpu_pjrt_client.h" +#include "xla/service/device_assignment.h" +#include "xla/service/gpu/backend_configs.pb.h" +#include "xla/service/gpu_topology.h" +#include "xla/service/platform_util.h" +#include "xla/shape_util.h" +#include "xla/status_macros.h" +#include "xla/stream_executor/cuda/cuda_compute_capability.h" +#include "xla/stream_executor/device_description.h" +#include "xla/stream_executor/platform.h" +#include "xla/stream_executor/stream_executor.h" +#include "xla/tests/literal_test_util.h" +#include "xla/tsl/platform/env.h" +#include "xla/tsl/platform/subprocess.h" +#include "xla/tsl/platform/test.h" +#include "xla/tsl/util/command_line_flags.h" +#include "xla/xla.pb.h" +#include "xla/xla_data.pb.h" +#include "tsl/platform/path.h" + +namespace xla { +namespace { + +inline constexpr size_t kMB = 1024LL * 1024LL; +constexpr int kNumNodes = 2; +static const char* test_binary_name; + +struct MultiProcessGpuClientSetup { + std::unique_ptr service; + std::unique_ptr client; +}; + +absl::StatusOr SetUpMultiProcessGpuClient( + int rank_id, int num_nodes, int port, absl::string_view log_prefix) { + MultiProcessGpuClientSetup prepared_test; + std::string coordinator_address = absl::StrFormat("127.0.0.1:%d", port); + + // Rank 0 creates the coordination service. + if (rank_id == 0) { + LOG(INFO) << log_prefix << ": creating coordination service on " + << coordinator_address; + xla::CoordinationServiceImpl::Options service_options; + service_options.num_nodes = num_nodes; + ABSL_ASSIGN_OR_RETURN(prepared_test.service, + xla::GetDistributedRuntimeService( + absl::StrFormat("[::]:%d", port), service_options)); + LOG(INFO) << log_prefix << ": created coordination service"; + } + + // Connect to the coordination service. + xla::DistributedRuntimeClient::Options distributed_options; + distributed_options.node_id = rank_id; + distributed_options.init_timeout = absl::Seconds(120); + auto distributed_client = + GetDistributedRuntimeClient(coordinator_address, distributed_options); + + LOG(INFO) << log_prefix << ": connecting distributed client"; + ABSL_RETURN_IF_ERROR(distributed_client->Connect()); + LOG(INFO) << log_prefix << ": distributed client connected"; + + // Create the GPU client with a single addressable device per process. + GpuClientOptions options; + options.node_id = rank_id; + options.num_nodes = num_nodes; + options.allowed_devices = {rank_id}; + options.kv_store = + GetDistributedKeyValueStore(distributed_client, /*key_prefix=*/"gpu:"); + options.distributed_client = distributed_client; + options.allocator_config.kind = xla::GpuAllocatorConfig::Kind::kBFC; + options.allocator_config.gpu_system_memory_size = 32 * kMB; + options.allocator_config.collective_memory_size = 0; + options.use_tfrt_gpu_client = true; + + LOG(INFO) << log_prefix << ": creating PjRtClient"; + ABSL_ASSIGN_OR_RETURN(prepared_test.client, GetXlaPjrtGpuClient(options)); + LOG(INFO) << log_prefix << ": PjRtClient created"; + + return prepared_test; +} + +absl::Status AllGatherMultiProcessTestBody(int node_id, int port) { + std::string log_prefix = absl::StrFormat("rank_%d", node_id); + ABSL_ASSIGN_OR_RETURN(se::Platform * platform, PlatformUtil::GetPlatform("gpu")); + ABSL_ASSIGN_OR_RETURN(se::StreamExecutor * executor, + platform->ExecutorForDevice(node_id)); + const auto& desc = executor->GetDeviceDescription(); + if (desc.gpu_compute_capability().IsCuda() && !desc.gpu_compute_capability() + .cuda_compute_capability() + ->IsAtLeastAmpere()) { + LOG(INFO) << log_prefix + << ": skipping one-shot all-gather test, requires Ampere+ for " + "CUDA."; + return absl::OkStatus(); + } + + ABSL_ASSIGN_OR_RETURN( + MultiProcessGpuClientSetup setup, + SetUpMultiProcessGpuClient(node_id, kNumNodes, port, log_prefix)); + std::unique_ptr client = std::move(setup.client); + + TF_RET_CHECK(client->addressable_device_count() == 1) + << "Expected exactly 1 local addressable device per process."; + TF_RET_CHECK(client->device_count() == kNumNodes) + << "Expected " << kNumNodes << " global devices."; + + constexpr absl::string_view kModuleStr = R"( + HloModule test + + ENTRY test_computation { + param_0 = f32[128] parameter(0) + ROOT all-gather = f32[256] all-gather(param_0), dimensions={0}, + replica_groups={{0,1}} + } + )"; + + ABSL_ASSIGN_OR_RETURN(auto hlo_module, + ParseAndReturnUnverifiedModule(kModuleStr, /*config=*/{})); + xla::XlaComputation computation(hlo_module->ToProto()); + + xla::CompileOptions compile_options; + compile_options.executable_build_options.set_num_replicas(kNumNodes); + compile_options.executable_build_options.set_num_partitions(1); + DeviceAssignment device_assignment(kNumNodes, 1); + device_assignment(0, 0) = 0; + device_assignment(1, 0) = 1; + compile_options.executable_build_options.set_device_assignment( + device_assignment); + + GpuTopology gpu_topology( + /*platform_version=*/"", + /*num_partitions=*/1, + /*num_hosts_per_partition=*/1, + /*num_devices_per_host=*/kNumNodes, + /*gpu_target_config=*/std::nullopt, + /*host_target_machine_options=*/std::nullopt, + /*num_devices_per_process=*/kNumNodes); + compile_options.executable_build_options.set_gpu_topology(gpu_topology); + + DebugOptions debug_options = GetDebugOptionsFromFlags(); + debug_options.set_xla_gpu_autotune_level(0); + debug_options.add_xla_gpu_experimental_use_collective_kernels( + DebugOptions::COLLECTIVE_KERNEL_ALL_GATHER); + debug_options.set_xla_gpu_experimental_enable_tiling_propagation(true); + debug_options.add_xla_gpu_unsupported_use_cross_host_one_shot_kernel( + DebugOptions::ALLGATHER); + *compile_options.executable_build_options.mutable_debug_options() = + debug_options; + + LOG(INFO) << log_prefix << ": compiling HLO module with one-shot all-gather"; + ABSL_ASSIGN_OR_RETURN(std::unique_ptr executable, + client->CompileAndLoad(computation, compile_options)); + LOG(INFO) << log_prefix << ": compilation succeeded"; + + // Inspect the collective kernel strategy selected in the optimized HLO. + ABSL_ASSIGN_OR_RETURN(auto hlo_modules, + executable->GetExecutable()->GetHloModules()); + TF_RET_CHECK(!hlo_modules.empty()); + const HloModule* optimized_module = hlo_modules.front().get(); + bool found_all_gather = false; + for (const HloComputation* comp : optimized_module->computations()) { + for (const HloInstruction* instr : comp->instructions()) { + if (instr->opcode() == HloOpcode::kAllGather || + instr->opcode() == HloOpcode::kAllGatherStart) { + found_all_gather = true; + auto gpu_config = instr->backend_config(); + TF_RET_CHECK(gpu_config.ok()) + << "Failed to get GpuBackendConfig for " << instr->name() << ": " + << gpu_config.status(); + LOG(INFO) + << log_prefix << ": AllGather instruction " << instr->name() + << " collective kernel strategy: " + << gpu::CollectiveBackendConfig::CollectiveKernelStrategy_Name( + gpu_config->collective_backend_config().kernel_strategy()); + TF_RET_CHECK( + gpu_config->collective_backend_config().kernel_strategy() == + gpu::CollectiveBackendConfig::KERNEL_STRATEGY_TRITON_ONE_SHOT) + << "Expected AllGather to use KERNEL_STRATEGY_TRITON_ONE_SHOT, but " + "got: " + << gpu::CollectiveBackendConfig::CollectiveKernelStrategy_Name( + gpu_config->collective_backend_config().kernel_strategy()); + } + } + } + TF_RET_CHECK(found_all_gather) + << "Expected to find an AllGather instruction in optimized HLO."; + + // Prepare input literal: rank 0 provides 1.0f, rank 1 provides 2.0f. + const float input_val = (node_id == 0) ? 1.0f : 2.0f; + Literal input_literal = + LiteralUtil::CreateR1(std::vector(128, input_val)); + + ABSL_ASSIGN_OR_RETURN(auto* memory_space, + client->addressable_devices()[0]->default_memory_space()); + ABSL_ASSIGN_OR_RETURN(auto input_buffer, + client->BufferFromHostLiteral(input_literal, memory_space)); + + std::vector> input_ptrs = {{input_buffer.get()}}; + LOG(INFO) << log_prefix << ": executing one-shot all-gather"; + ABSL_ASSIGN_OR_RETURN(auto results, + executable->Execute(input_ptrs, ExecuteOptions())); + LOG(INFO) << log_prefix << ": execution finished"; + + TF_RET_CHECK(results.size() == 1 && results[0].size() == 1); + Literal result_literal(ShapeUtil::MakeShape(F32, {256})); + ABSL_RETURN_IF_ERROR(results[0][0]->ToLiteralSync(&result_literal)); + + // Expected result is [1.0f x 128, 2.0f x 128] on all ranks. + std::vector expected_data; + expected_data.reserve(256); + for (int i = 0; i < 128; ++i) { + expected_data.push_back(1.0f); + } + for (int i = 0; i < 128; ++i) { + expected_data.push_back(2.0f); + } + Literal expected_literal = LiteralUtil::CreateR1(expected_data); + + if (!LiteralTestUtil::Equal(expected_literal, result_literal)) { + return absl::InternalError(absl::StrFormat( + "Result literal %s does not match expected %s on rank %d", + result_literal.ToString(), expected_literal.ToString(), node_id)); + } + + LOG(INFO) << log_prefix << ": verified result successfully"; + return absl::OkStatus(); +} + +TEST(AllGatherMultiProcessE2ETest, OneShotAllGather2Processes) { + absl::StatusOr platform = PlatformUtil::GetPlatform("gpu"); + if (!platform.ok() || (*platform)->VisibleDeviceCount() < kNumNodes) { + GTEST_SKIP() << "Test requires at least " << kNumNodes + << " GPU devices, but found " + << (platform.ok() ? (*platform)->VisibleDeviceCount() : 0); + } + + // Re-pass XLA_FLAGS environment variable to child processes. + const char* xla_flags = std::getenv("XLA_FLAGS"); + if (xla_flags != nullptr) { + tsl::setenv("XLA_FLAGS", xla_flags, /*overwrite=*/true); + } + + int port = tsl::testing::PickUnusedPortOrDie(); + tsl::SubProcess child[kNumNodes]; + + for (int node_id = 0; node_id < kNumNodes; ++node_id) { + std::vector argv = { + test_binary_name, + "--test_to_run=AllGatherMultiProcessHelper", + absl::StrFormat("--node_id=%d", node_id), + absl::StrFormat("--port=%d", port), + "--alsologtostderr", + "--vmodule=gpu_executable=1,thunk_executor=1,all_gather_thunk=5," + "collective_kernel_thunk=5,collective_memory=5", + }; + child[node_id].SetProgram(test_binary_name, argv); + child[node_id].SetChannelAction(tsl::CHAN_STDOUT, tsl::ACTION_PIPE); + child[node_id].SetChannelAction(tsl::CHAN_STDERR, tsl::ACTION_PIPE); + ASSERT_TRUE(child[node_id].Start()) << "Failed to start node " << node_id; + } + + for (int node_id = 0; node_id < kNumNodes; ++node_id) { + std::string stdout_str, stderr_str; + int status = child[node_id].Communicate(nullptr, &stdout_str, &stderr_str); + + const char* undeclared_outputs_dir = + std::getenv("TEST_UNDECLARED_OUTPUTS_DIR"); + if (undeclared_outputs_dir != nullptr && + undeclared_outputs_dir[0] != '\0') { + std::string stderr_file = tsl::io::JoinPath( + undeclared_outputs_dir, + absl::StrFormat("subprocess_node_%d_stderr.log", node_id)); + absl::Status write_status = + tsl::WriteStringToFile(tsl::Env::Default(), stderr_file, stderr_str); + if (!write_status.ok()) { + LOG(WARNING) << "Failed to write stderr to " << stderr_file << ": " + << write_status; + } + } + + EXPECT_EQ(status, 0) << "node " << node_id << " failed with status " + << status << "\nstdout:\n" + << stdout_str << "\nstderr:\n" + << stderr_str; + } +} + +} // namespace +} // namespace xla + +int main(int argc, char* argv[]) { + std::string test_to_run; + int node_id = -1; + int port = -1; + xla::test_binary_name = argv[0]; + + std::vector flag_list = { + tsl::Flag("test_to_run", &test_to_run, + "The test to run in the child process."), + tsl::Flag("node_id", &node_id, + "The node id (rank) for the child process."), + tsl::Flag("port", &port, "The coordinator port for distributed runtime."), + }; + + xla::AppendDebugOptionsFlags(&flag_list); + tsl::Flags::Parse(&argc, argv, flag_list); + testing::InitGoogleTest(&argc, argv); + + if (test_to_run.empty()) { + return RUN_ALL_TESTS(); + } + + absl::Status result = absl::OkStatus(); + if (test_to_run == "AllGatherMultiProcessHelper") { + result = xla::AllGatherMultiProcessTestBody(node_id, port); + } else { + result = absl::InvalidArgumentError(absl::StrFormat( + "Unrecognized multiprocess test name: %s", test_to_run)); + } + + if (!result.ok()) { + LOG(ERROR) << "Child process (node_id " << node_id + << ") failed: " << result; + } + return result.raw_code(); +} diff --git a/third_party/xla/xla/backends/gpu/tests/all_gather_e2e_test.cc b/third_party/xla/xla/backends/gpu/tests/all_gather_e2e_test.cc new file mode 100644 index 00000000000000..d0a3428223eaaa --- /dev/null +++ b/third_party/xla/xla/backends/gpu/tests/all_gather_e2e_test.cc @@ -0,0 +1,254 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include +#include +#include +#include + +#include +#include "absl/strings/string_view.h" +#include "xla/backends/gpu/tests/collective_ops_e2e_test_base.h" +#include "xla/hlo/ir/hlo_computation.h" +#include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/ir/hlo_module.h" +#include "xla/hlo/ir/hlo_opcode.h" +#include "xla/literal.h" +#include "xla/literal_util.h" +#include "xla/service/gpu/backend_configs.pb.h" +#include "xla/tests/literal_test_util.h" +#include "xla/tsl/platform/test.h" +#include "xla/xla.pb.h" + +namespace xla { +namespace { + +class AllGatherTest : public CollectiveOpsWithFlagsBase { + public: + AllGatherTest() + : CollectiveOpsWithFlagsBase( + /*enable_async=*/true, + /*enable_p2p_memcpy=*/false, + /*enable_symmetric_buffer=*/true, + /*memory_size=*/32 * kMB, + /*collectives_memory_size=*/32 * kMB) {} + + protected: + void SetUp() override { + CollectiveOpsE2ETestBase::SetUp(); + if (Capability().IsCuda() && !IsHopperAndHigher()) { + GTEST_SKIP() << "Test requires Hopper or newer architecture."; + } + } + + DebugOptions GetDebugOptionsForTest() const override { + DebugOptions opts = CollectiveOpsWithFlagsBase::GetDebugOptionsForTest(); + opts.clear_xla_gpu_experimental_use_collective_kernels(); + opts.add_xla_gpu_experimental_use_collective_kernels( + DebugOptions::COLLECTIVE_KERNEL_ALL_GATHER); + return opts; + } + + bool CheckDeviceCount(int32_t required_device_count) { + [&]() -> void { + const int32_t current_device_count = device_count(); + if (current_device_count < required_device_count) { + ASSERT_GE(current_device_count, 2) + << "Test requires at least 2 devices but only " + << current_device_count << " available"; + if (current_device_count < required_device_count) { + GTEST_SKIP() << "Test requires at least " << required_device_count + << " devices but only " << current_device_count + << " available."; + } + } + }(); + return !IsSkipped() && !HasFatalFailure(); + } + + void VerifyOneShotAllGather(const HloModule* optimized_module) { + ASSERT_NE(optimized_module, nullptr); + bool found_all_gather = false; + for (const HloComputation* comp : optimized_module->computations()) { + for (const HloInstruction* instr : comp->instructions()) { + if (instr->opcode() == HloOpcode::kAllGather || + instr->opcode() == HloOpcode::kAllGatherStart) { + found_all_gather = true; + ASSERT_OK_AND_ASSIGN(gpu::GpuBackendConfig gpu_config, + instr->backend_config()); + EXPECT_EQ( + gpu_config.collective_backend_config().kernel_strategy(), + gpu::CollectiveBackendConfig::KERNEL_STRATEGY_TRITON_ONE_SHOT) + << "Expected AllGather instruction " << instr->name() + << " to use KERNEL_STRATEGY_TRITON_ONE_SHOT, but got: " + << gpu::CollectiveBackendConfig::CollectiveKernelStrategy_Name( + gpu_config.collective_backend_config().kernel_strategy()); + } + } + } + EXPECT_TRUE(found_all_gather) + << "Expected to find an AllGather instruction in optimized HLO."; + } +}; + +// Basic 2-GPU all-gather of f32[128] -> f32[256]. +TEST_F(AllGatherTest, Basic2GpuF32) { + constexpr int32_t kNumReplicas = 2; + if (!CheckDeviceCount(kNumReplicas)) { + return; + } + + constexpr absl::string_view kModuleStr = R"( + HloModule test + ENTRY test_computation { + param_0 = f32[128] parameter(0) + ROOT all-gather = f32[256] all-gather(param_0), dimensions={0}, + replica_groups={{0,1}} + } + )"; + + ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(kModuleStr, kNumReplicas)); + + // Create input: rank 0 gets [1, 1, ...], rank 1 gets [2, 2, ...]. + Literal input_r0 = + LiteralUtil::CreateR1(std::vector(128, 1.0f)); + Literal input_r1 = + LiteralUtil::CreateR1(std::vector(128, 2.0f)); + + std::vector> args = {{&input_r0}, {&input_r1}}; + ASSERT_OK_AND_ASSIGN(ExecutionResult result, + ExecuteReplicated(std::move(module), args)); + + VerifyOneShotAllGather(result.optimized_module); + + ASSERT_EQ(result.results.size(), kNumReplicas); + + // Expected output: [1, 1, ..., 2, 2, ...] (128 ones followed by 128 twos). + std::vector expected_data; + expected_data.reserve(256); + for (int i = 0; i < 128; ++i) { + expected_data.push_back(1.0f); + } + for (int i = 0; i < 128; ++i) { + expected_data.push_back(2.0f); + } + Literal expected = LiteralUtil::CreateR1(expected_data); + + for (int i = 0; i < kNumReplicas; ++i) { + EXPECT_TRUE(LiteralTestUtil::Equal(expected, result.results[i])) + << "Mismatch at replica " << i; + } +} + +// Larger 2-GPU all-gather to test multi-tile behavior: f32[4096] -> f32[8192]. +TEST_F(AllGatherTest, Large2GpuF32) { + constexpr int32_t kNumReplicas = 2; + if (!CheckDeviceCount(kNumReplicas)) { + return; + } + + constexpr absl::string_view kModuleStr = R"( + HloModule test + ENTRY test_computation { + param_0 = f32[4096] parameter(0) + ROOT all-gather = f32[8192] all-gather(param_0), dimensions={0}, + replica_groups={{0,1}} + } + )"; + + ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(kModuleStr, kNumReplicas)); + + // rank 0: incrementing values [0, 1, 2, ..., 4095] + // rank 1: incrementing values [4096, 4097, ..., 8191] + std::vector data_r0(4096), data_r1(4096); + for (int i = 0; i < 4096; ++i) { + data_r0[i] = static_cast(i); + data_r1[i] = static_cast(i + 4096); + } + Literal input_r0 = LiteralUtil::CreateR1(data_r0); + Literal input_r1 = LiteralUtil::CreateR1(data_r1); + + std::vector> args = {{&input_r0}, {&input_r1}}; + ASSERT_OK_AND_ASSIGN(ExecutionResult result, + ExecuteReplicated(std::move(module), args)); + + VerifyOneShotAllGather(result.optimized_module); + + ASSERT_EQ(result.results.size(), kNumReplicas); + + // Expected: [0, 1, ..., 8191] for both replicas. + std::vector expected_data(8192); + for (int i = 0; i < 8192; ++i) { + expected_data[i] = static_cast(i); + } + Literal expected = LiteralUtil::CreateR1(expected_data); + + for (int i = 0; i < kNumReplicas; ++i) { + EXPECT_TRUE(LiteralTestUtil::Equal(expected, result.results[i])) + << "Mismatch at replica " << i; + } +} + +// 2D shape: f32[16, 32] -> f32[32, 32] (gather along dim 0). +TEST_F(AllGatherTest, TwoDimensional2Gpu) { + constexpr int32_t kNumReplicas = 2; + if (!CheckDeviceCount(kNumReplicas)) { + return; + } + + constexpr absl::string_view kModuleStr = R"( + HloModule test + ENTRY test_computation { + param_0 = f32[16,32] parameter(0) + ROOT all-gather = f32[32,32] all-gather(param_0), dimensions={0}, + replica_groups={{0,1}} + } + )"; + + ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(kModuleStr, kNumReplicas)); + + // rank 0: all 1s, rank 1: all 2s + Literal input_r0 = LiteralUtil::CreateFull({16, 32}, 1.0f); + Literal input_r1 = LiteralUtil::CreateFull({16, 32}, 2.0f); + + std::vector> args = {{&input_r0}, {&input_r1}}; + ASSERT_OK_AND_ASSIGN(ExecutionResult result, + ExecuteReplicated(std::move(module), args)); + + VerifyOneShotAllGather(result.optimized_module); + + ASSERT_EQ(result.results.size(), kNumReplicas); + + // Expected: first 16 rows are 1s, next 16 rows are 2s. + // Build expected by checking individual elements. + Literal expected = LiteralUtil::CreateFull({32, 32}, 0.0f); + for (int64_t row = 0; row < 32; ++row) { + float val = (row < 16) ? 1.0f : 2.0f; + for (int64_t col = 0; col < 32; ++col) { + expected.Set({row, col}, val); + } + } + + for (int i = 0; i < kNumReplicas; ++i) { + EXPECT_TRUE(LiteralTestUtil::Equal(expected, result.results[i])) + << "Mismatch at replica " << i; + } +} + +} // namespace +} // namespace xla diff --git a/third_party/xla/xla/backends/gpu/tests/all_reduce_multi_process_e2e_test.cc b/third_party/xla/xla/backends/gpu/tests/all_reduce_multi_process_e2e_test.cc new file mode 100644 index 00000000000000..02f19c04d2b597 --- /dev/null +++ b/third_party/xla/xla/backends/gpu/tests/all_reduce_multi_process_e2e_test.cc @@ -0,0 +1,367 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include +#include +#include +#include +#include +#include + +#include +#include "absl/log/check.h" +#include "absl/log/log.h" +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "absl/strings/str_format.h" +#include "absl/strings/string_view.h" +#include "absl/time/time.h" +#include "xla/debug_options_flags.h" +#include "xla/hlo/builder/xla_computation.h" +#include "xla/hlo/ir/hlo_computation.h" +#include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/ir/hlo_module.h" +#include "xla/hlo/ir/hlo_opcode.h" +#include "xla/hlo/parser/hlo_parser.h" +#include "xla/literal.h" +#include "xla/literal_util.h" +#include "xla/pjrt/distributed/client.h" +#include "xla/pjrt/distributed/distributed.h" +#include "xla/pjrt/distributed/service.h" +#include "xla/pjrt/pjrt_client.h" +#include "xla/pjrt/pjrt_compiler.h" +#include "xla/pjrt/pjrt_executable.h" +#include "xla/pjrt/plugin/xla_gpu/xla_gpu_allocator_config.h" +#include "xla/pjrt/plugin/xla_gpu/xla_gpu_client_options.h" +#include "xla/pjrt/plugin/xla_gpu/xla_gpu_pjrt_client.h" +#include "xla/service/device_assignment.h" +#include "xla/service/gpu/backend_configs.pb.h" +#include "xla/service/gpu_topology.h" +#include "xla/service/platform_util.h" +#include "xla/shape_util.h" +#include "xla/status_macros.h" +#include "xla/stream_executor/cuda/cuda_compute_capability.h" +#include "xla/stream_executor/device_description.h" +#include "xla/stream_executor/platform.h" +#include "xla/stream_executor/stream_executor.h" +#include "xla/tests/literal_test_util.h" +#include "xla/tsl/platform/env.h" +#include "xla/tsl/platform/subprocess.h" +#include "xla/tsl/platform/test.h" +#include "xla/tsl/util/command_line_flags.h" +#include "xla/xla.pb.h" +#include "xla/xla_data.pb.h" +#include "tsl/platform/path.h" + +namespace xla { +namespace { + +inline constexpr size_t kMB = 1024LL * 1024LL; +constexpr int kNumNodes = 2; +static const char* test_binary_name; + +struct MultiProcessGpuClientSetup { + std::unique_ptr service; + std::unique_ptr client; +}; + +absl::StatusOr SetUpMultiProcessGpuClient( + int rank_id, int num_nodes, int port, absl::string_view log_prefix) { + MultiProcessGpuClientSetup prepared_test; + std::string coordinator_address = absl::StrFormat("127.0.0.1:%d", port); + + // Rank 0 creates the coordination service. + if (rank_id == 0) { + LOG(INFO) << log_prefix << ": creating coordination service on " + << coordinator_address; + xla::CoordinationServiceImpl::Options service_options; + service_options.num_nodes = num_nodes; + ABSL_ASSIGN_OR_RETURN(prepared_test.service, + xla::GetDistributedRuntimeService( + absl::StrFormat("[::]:%d", port), service_options)); + LOG(INFO) << log_prefix << ": created coordination service"; + } + + // Connect to the coordination service. + xla::DistributedRuntimeClient::Options distributed_options; + distributed_options.node_id = rank_id; + distributed_options.init_timeout = absl::Seconds(120); + auto distributed_client = + GetDistributedRuntimeClient(coordinator_address, distributed_options); + + LOG(INFO) << log_prefix << ": connecting distributed client"; + ABSL_RETURN_IF_ERROR(distributed_client->Connect()); + LOG(INFO) << log_prefix << ": distributed client connected"; + + // Create the GPU client with a single addressable device per process. + GpuClientOptions options; + options.node_id = rank_id; + options.num_nodes = num_nodes; + options.allowed_devices = {rank_id}; + options.kv_store = + GetDistributedKeyValueStore(distributed_client, /*key_prefix=*/"gpu:"); + options.distributed_client = distributed_client; + options.allocator_config.kind = xla::GpuAllocatorConfig::Kind::kBFC; + options.allocator_config.gpu_system_memory_size = 32 * kMB; + options.allocator_config.collective_memory_size = 0; + options.use_tfrt_gpu_client = true; + + LOG(INFO) << log_prefix << ": creating PjRtClient"; + ABSL_ASSIGN_OR_RETURN(prepared_test.client, GetXlaPjrtGpuClient(options)); + LOG(INFO) << log_prefix << ": PjRtClient created"; + + return prepared_test; +} + +absl::Status AllReduceMultiProcessTestBody(int node_id, int port) { + std::string log_prefix = absl::StrFormat("rank_%d", node_id); + ABSL_ASSIGN_OR_RETURN(se::Platform * platform, PlatformUtil::GetPlatform("gpu")); + ABSL_ASSIGN_OR_RETURN(se::StreamExecutor * executor, + platform->ExecutorForDevice(node_id)); + const auto& desc = executor->GetDeviceDescription(); + if (desc.gpu_compute_capability().IsCuda() && !desc.gpu_compute_capability() + .cuda_compute_capability() + ->IsAtLeastAmpere()) { + LOG(INFO) << log_prefix + << ": skipping one-shot all-reduce test, requires Ampere+ for " + "CUDA."; + return absl::OkStatus(); + } + + ABSL_ASSIGN_OR_RETURN( + MultiProcessGpuClientSetup setup, + SetUpMultiProcessGpuClient(node_id, kNumNodes, port, log_prefix)); + std::unique_ptr client = std::move(setup.client); + + TF_RET_CHECK(client->addressable_device_count() == 1) + << "Expected exactly 1 local addressable device per process."; + TF_RET_CHECK(client->device_count() == kNumNodes) + << "Expected " << kNumNodes << " global devices."; + + constexpr absl::string_view kModuleStr = R"( + HloModule test + + apply_op { + x = f32[] parameter(0) + y = f32[] parameter(1) + ROOT apply_op = f32[] add(x, y) + } + + ENTRY test_computation { + param_0 = f32[128] parameter(0) + ROOT all-reduce = f32[128] all-reduce(param_0), to_apply=apply_op, + replica_groups={{0,1}} + } + )"; + + ABSL_ASSIGN_OR_RETURN(auto hlo_module, + ParseAndReturnUnverifiedModule(kModuleStr, /*config=*/{})); + xla::XlaComputation computation(hlo_module->ToProto()); + + xla::CompileOptions compile_options; + compile_options.executable_build_options.set_num_replicas(kNumNodes); + compile_options.executable_build_options.set_num_partitions(1); + DeviceAssignment device_assignment(kNumNodes, 1); + device_assignment(0, 0) = 0; + device_assignment(1, 0) = 1; + compile_options.executable_build_options.set_device_assignment( + device_assignment); + + GpuTopology gpu_topology( + /*platform_version=*/"", + /*num_partitions=*/1, + /*num_hosts_per_partition=*/1, + /*num_devices_per_host=*/kNumNodes, + /*gpu_target_config=*/std::nullopt, + /*host_target_machine_options=*/std::nullopt, + /*num_devices_per_process=*/kNumNodes); + compile_options.executable_build_options.set_gpu_topology(gpu_topology); + + DebugOptions debug_options = GetDebugOptionsFromFlags(); + debug_options.set_xla_gpu_autotune_level(0); + debug_options.add_xla_gpu_unsupported_use_cross_host_one_shot_kernel( + DebugOptions::ALLREDUCE); + *compile_options.executable_build_options.mutable_debug_options() = + debug_options; + + LOG(INFO) << log_prefix << ": compiling HLO module with one-shot all-reduce"; + ABSL_ASSIGN_OR_RETURN(std::unique_ptr executable, + client->CompileAndLoad(computation, compile_options)); + LOG(INFO) << log_prefix << ": compilation succeeded"; + + // Inspect the collective kernel strategy selected in the optimized HLO. + ABSL_ASSIGN_OR_RETURN(auto hlo_modules, + executable->GetExecutable()->GetHloModules()); + TF_RET_CHECK(!hlo_modules.empty()); + const HloModule* optimized_module = hlo_modules.front().get(); + bool found_all_reduce = false; + for (const HloComputation* comp : optimized_module->computations()) { + for (const HloInstruction* instr : comp->instructions()) { + if (instr->opcode() == HloOpcode::kAllReduce || + instr->opcode() == HloOpcode::kAllReduceStart) { + found_all_reduce = true; + auto gpu_config = instr->backend_config(); + TF_RET_CHECK(gpu_config.ok()) + << "Failed to get GpuBackendConfig for " << instr->name() << ": " + << gpu_config.status(); + LOG(INFO) + << log_prefix << ": AllReduce instruction " << instr->name() + << " collective kernel strategy: " + << gpu::CollectiveBackendConfig::CollectiveKernelStrategy_Name( + gpu_config->collective_backend_config().kernel_strategy()); + TF_RET_CHECK( + gpu_config->collective_backend_config().kernel_strategy() == + gpu::CollectiveBackendConfig::KERNEL_STRATEGY_TRITON_ONE_SHOT) + << "Expected AllReduce to use KERNEL_STRATEGY_TRITON_ONE_SHOT, but " + "got: " + << gpu::CollectiveBackendConfig::CollectiveKernelStrategy_Name( + gpu_config->collective_backend_config().kernel_strategy()); + } + } + } + TF_RET_CHECK(found_all_reduce) + << "Expected to find an AllReduce instruction in optimized HLO."; + + // Prepare input literal: rank 0 provides 10.0f, rank 1 provides 20.0f. + const float input_val = (node_id == 0) ? 10.0f : 20.0f; + Literal input_literal = + LiteralUtil::CreateR1(std::vector(128, input_val)); + + ABSL_ASSIGN_OR_RETURN(auto* memory_space, + client->addressable_devices()[0]->default_memory_space()); + ABSL_ASSIGN_OR_RETURN(auto input_buffer, + client->BufferFromHostLiteral(input_literal, memory_space)); + + std::vector> input_ptrs = {{input_buffer.get()}}; + LOG(INFO) << log_prefix << ": executing one-shot all-reduce"; + ABSL_ASSIGN_OR_RETURN(auto results, + executable->Execute(input_ptrs, ExecuteOptions())); + LOG(INFO) << log_prefix << ": execution finished"; + + TF_RET_CHECK(results.size() == 1 && results[0].size() == 1); + Literal result_literal(ShapeUtil::MakeShape(F32, {128})); + ABSL_RETURN_IF_ERROR(results[0][0]->ToLiteralSync(&result_literal)); + + // Expected result is 10.0f + 20.0f = 30.0f across all ranks. + Literal expected_literal = + LiteralUtil::CreateR1(std::vector(128, 30.0f)); + if (!LiteralTestUtil::Equal(expected_literal, result_literal)) { + return absl::InternalError(absl::StrFormat( + "Result literal %s does not match expected %s on rank %d", + result_literal.ToString(), expected_literal.ToString(), node_id)); + } + + LOG(INFO) << log_prefix << ": verified result successfully"; + return absl::OkStatus(); +} + +TEST(AllReduceMultiProcessE2ETest, OneShotAllReduce2Processes) { + absl::StatusOr platform = PlatformUtil::GetPlatform("gpu"); + if (!platform.ok() || (*platform)->VisibleDeviceCount() < kNumNodes) { + GTEST_SKIP() << "Test requires at least " << kNumNodes + << " GPU devices, but found " + << (platform.ok() ? (*platform)->VisibleDeviceCount() : 0); + } + + // Re-pass XLA_FLAGS environment variable to child processes. + const char* xla_flags = std::getenv("XLA_FLAGS"); + if (xla_flags != nullptr) { + tsl::setenv("XLA_FLAGS", xla_flags, /*overwrite=*/true); + } + + int port = tsl::testing::PickUnusedPortOrDie(); + tsl::SubProcess child[kNumNodes]; + + for (int node_id = 0; node_id < kNumNodes; ++node_id) { + std::vector argv = { + test_binary_name, + "--test_to_run=AllReduceMultiProcessHelper", + absl::StrFormat("--node_id=%d", node_id), + absl::StrFormat("--port=%d", port), + "--alsologtostderr", + "--vmodule=gpu_executable=1,thunk_executor=1,all_reduce_thunk=5," + "collective_kernel_thunk=5,collective_memory=5", + }; + child[node_id].SetProgram(test_binary_name, argv); + child[node_id].SetChannelAction(tsl::CHAN_STDOUT, tsl::ACTION_PIPE); + child[node_id].SetChannelAction(tsl::CHAN_STDERR, tsl::ACTION_PIPE); + ASSERT_TRUE(child[node_id].Start()) << "Failed to start node " << node_id; + } + + for (int node_id = 0; node_id < kNumNodes; ++node_id) { + std::string stdout_str, stderr_str; + int status = child[node_id].Communicate(nullptr, &stdout_str, &stderr_str); + + const char* undeclared_outputs_dir = + std::getenv("TEST_UNDECLARED_OUTPUTS_DIR"); + if (undeclared_outputs_dir != nullptr && + undeclared_outputs_dir[0] != '\0') { + std::string stderr_file = tsl::io::JoinPath( + undeclared_outputs_dir, + absl::StrFormat("subprocess_node_%d_stderr.log", node_id)); + absl::Status write_status = + tsl::WriteStringToFile(tsl::Env::Default(), stderr_file, stderr_str); + if (!write_status.ok()) { + LOG(WARNING) << "Failed to write stderr to " << stderr_file << ": " + << write_status; + } + } + + EXPECT_EQ(status, 0) << "node " << node_id << " failed with status " + << status << "\nstdout:\n" + << stdout_str << "\nstderr:\n" + << stderr_str; + } +} + +} // namespace +} // namespace xla + +int main(int argc, char* argv[]) { + std::string test_to_run; + int node_id = -1; + int port = -1; + xla::test_binary_name = argv[0]; + + std::vector flag_list = { + tsl::Flag("test_to_run", &test_to_run, + "The test to run in the child process."), + tsl::Flag("node_id", &node_id, + "The node id (rank) for the child process."), + tsl::Flag("port", &port, "The coordinator port for distributed runtime."), + }; + + xla::AppendDebugOptionsFlags(&flag_list); + tsl::Flags::Parse(&argc, argv, flag_list); + testing::InitGoogleTest(&argc, argv); + + if (test_to_run.empty()) { + return RUN_ALL_TESTS(); + } + + absl::Status result = absl::OkStatus(); + if (test_to_run == "AllReduceMultiProcessHelper") { + result = xla::AllReduceMultiProcessTestBody(node_id, port); + } else { + result = absl::InvalidArgumentError(absl::StrFormat( + "Unrecognized multiprocess test name: %s", test_to_run)); + } + + if (!result.ok()) { + LOG(ERROR) << "Child process (node_id " << node_id + << ") failed: " << result; + } + return result.raw_code(); +} diff --git a/third_party/xla/xla/backends/gpu/tests/gpu_cub_sort_test.cc b/third_party/xla/xla/backends/gpu/tests/gpu_cub_sort_test.cc index c160271faeaddf..1cbbdce85c00b0 100644 --- a/third_party/xla/xla/backends/gpu/tests/gpu_cub_sort_test.cc +++ b/third_party/xla/xla/backends/gpu/tests/gpu_cub_sort_test.cc @@ -191,7 +191,7 @@ numpy_order_comparator { rhs_is_zero = pred[] compare(rhs, c_zero), direction=EQ rhs_no_neg_zero = bf16[] select(rhs_is_zero, c_zero, rhs) rhs_no_neg_zero_or_nan = bf16[] select(rhs_is_nan, c_nan, rhs_no_neg_zero) - ROOT compare.20017 = pred[] compare(lhs_no_neg_zero_or_nan, rhs_no_neg_zero_or_nan), direction=GT, type=TOTALORDER + ROOT compare.20017 = pred[] compare(lhs_no_neg_zero_or_nan, rhs_no_neg_zero_or_nan), direction=GT, order=TOTAL } ENTRY main { @@ -223,7 +223,7 @@ TEST_P(CubSortKeysSpecialOrderingTest, CompareToReferenceTotalOrderLt) { compare { lhs = f32[] parameter(0) rhs = f32[] parameter(1) - ROOT comp = pred[] compare(lhs, rhs), direction=LT, type=TOTALORDER + ROOT comp = pred[] compare(lhs, rhs), direction=LT, order=TOTAL } ENTRY main { diff --git a/third_party/xla/xla/backends/gpu/tests/swap_conv_operands_test.cc b/third_party/xla/xla/backends/gpu/tests/swap_conv_operands_test.cc index 7ed04346490c4f..941aef2f58280f 100644 --- a/third_party/xla/xla/backends/gpu/tests/swap_conv_operands_test.cc +++ b/third_party/xla/xla/backends/gpu/tests/swap_conv_operands_test.cc @@ -26,6 +26,7 @@ limitations under the License. #include "xla/hlo/testlib/filecheck.h" #include "xla/stream_executor/device_description.h" #include "xla/tests/hlo_pjrt_interpreter_reference_mixin.h" +#include "xla/xla.pb.h" namespace xla::gpu { namespace { @@ -33,6 +34,13 @@ namespace { class SwapConvOperandsTest : public HloInterpreterReferenceMixin { public: + DebugOptions GetDebugOptionsForTest() const override { + DebugOptions debug_options = HloInterpreterReferenceMixin< + HloPjRtGpuTestBase>::GetDebugOptionsForTest(); + debug_options.set_xla_gpu_experimental_enable_conv_fusion(false); + return debug_options; + } + void MatchOptimizedHlo(absl::string_view hlo, absl::string_view expected_hlo) { ASSERT_OK_AND_ASSIGN(std::unique_ptr optimized_module, diff --git a/third_party/xla/xla/backends/gpu/transforms/BUILD b/third_party/xla/xla/backends/gpu/transforms/BUILD index 3a1d75ccf57470..28e572b92b9302 100644 --- a/third_party/xla/xla/backends/gpu/transforms/BUILD +++ b/third_party/xla/xla/backends/gpu/transforms/BUILD @@ -274,6 +274,7 @@ xla_cc_test( "//xla/hlo/testlib:hlo_hardware_independent_test_base", "//xla/hlo/utils:hlo_traversal", "//xla/service:hlo_cost_analysis", + "//xla/service:hlo_module_config", "//xla/service/gpu:backend_configs_cc", "//xla/service/gpu:gpu_device_info_for_tests", "//xla/service/gpu:ir_emission_utils", @@ -3635,6 +3636,8 @@ cc_library( srcs = ["topk_specializer.cc"], hdrs = ["topk_specializer.h"], deps = [ + "//xla:comparison_util", + "//xla:literal_util", "//xla:shape_util", "//xla:status_macros", "//xla:util", @@ -3650,6 +3653,7 @@ cc_library( "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", ], @@ -3658,6 +3662,9 @@ cc_library( xla_test( name = "topk_specializer_test", srcs = ["topk_specializer_test.cc"], + backend_tags = { + "gb300": ["broken"], # TODO(b/505574653): re-enable when fixed. + }, backends = ["gpu"], shard_count = 10, deps = [ @@ -3666,6 +3673,7 @@ xla_test( "//xla/backends/gpu/tests:hlo_pjrt_gpu_test_base", "//xla/hlo/ir:hlo", "//xla/hlo/pass:hlo_pass", + "//xla/hlo/testlib:filecheck", "//xla/service:topk_rewriter", "//xla/stream_executor:device_description", "//xla/tests:hlo_test_base", diff --git a/third_party/xla/xla/backends/gpu/transforms/collectives/collective_fusion.cc b/third_party/xla/xla/backends/gpu/transforms/collectives/collective_fusion.cc index 9aad3aa66a8787..179dcdddadc84e 100644 --- a/third_party/xla/xla/backends/gpu/transforms/collectives/collective_fusion.cc +++ b/third_party/xla/xla/backends/gpu/transforms/collectives/collective_fusion.cc @@ -111,6 +111,10 @@ std::vector GetFusionCandidates( } bool ShouldFlatten(const HloInstruction* instr) { + if (HloPredicateIsNotOp( + instr)) { + return false; + } const int64_t size_bytes = ShapeUtil::ElementsIn(instr->shape()) * primitive_util::ByteWidth(instr->shape().element_type()); diff --git a/third_party/xla/xla/backends/gpu/transforms/collectives/collective_kernel_strategy_annotator_test.cc b/third_party/xla/xla/backends/gpu/transforms/collectives/collective_kernel_strategy_annotator_test.cc index 4f20357adfe598..e36e8e883b30a7 100644 --- a/third_party/xla/xla/backends/gpu/transforms/collectives/collective_kernel_strategy_annotator_test.cc +++ b/third_party/xla/xla/backends/gpu/transforms/collectives/collective_kernel_strategy_annotator_test.cc @@ -249,6 +249,34 @@ TEST_F(CollectiveKernelStrategyAnnotatorTest, EXPECT_EQ(strategy, CollectiveBackendConfig::KERNEL_STRATEGY_DEFAULT); } +// 2 * 1024 * 1024 F32 elements per replica = 8 MB > 4 MB limit +// → ineligible → KERNEL_STRATEGY_DEFAULT (falls back to NCCL). +TEST_F(CollectiveKernelStrategyAnnotatorTest, + LargeAllGatherKeepsDefaultStrategy) { + constexpr int kNumReplicas = 8; + constexpr int64_t kInputElements = 2 * 1024 * 1024; + constexpr int64_t kOutputElements = kInputElements * kNumReplicas; + std::string replica_groups_str = "0,1,2,3,4,5,6,7"; + std::string hlo = absl::StrFormat(kAllGatherHloTemplate, kInputElements, + kOutputElements, replica_groups_str); + + ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(hlo, kNumReplicas)); + module->mutable_config() + .mutable_debug_options() + .add_xla_gpu_experimental_use_collective_kernels( + DebugOptions::COLLECTIVE_KERNEL_ALL_GATHER); + ASSERT_OK_AND_ASSIGN(auto local_topology, MakeLocalGpuTopology(kNumReplicas)); + + CollectiveKernelStrategyAnnotator annotator(*local_topology, + /*is_multimem_enabled=*/false); + ASSERT_OK(annotator.Run(module.get()).status()); + + ASSERT_OK_AND_ASSIGN(auto strategy, + GetKernelStrategy(module.get(), HloOpcode::kAllGather)); + EXPECT_EQ(strategy, CollectiveBackendConfig::KERNEL_STRATEGY_DEFAULT); +} + // Module with both AllReduce and AllGather: both should be annotated in a // single pass. TEST_F(CollectiveKernelStrategyAnnotatorTest, diff --git a/third_party/xla/xla/backends/gpu/transforms/collectives/collective_ops_utils.cc b/third_party/xla/xla/backends/gpu/transforms/collectives/collective_ops_utils.cc index 1838bfe46ff948..babfecbc8046f2 100644 --- a/third_party/xla/xla/backends/gpu/transforms/collectives/collective_ops_utils.cc +++ b/third_party/xla/xla/backends/gpu/transforms/collectives/collective_ops_utils.cc @@ -43,6 +43,7 @@ limitations under the License. #include "xla/service/collective_ops_utils.h" #include "xla/service/device_assignment.h" #include "xla/service/gpu/backend_configs.pb.h" +#include "xla/service/gpu_topology.h" #include "xla/service/hlo_module_config.h" #include "xla/side_effect_util.h" #include "xla/stream_executor/device_description.h" @@ -356,7 +357,7 @@ bool IsAllReplicasLocal(int64_t gpus_per_host, << (device_assignment != nullptr ? device_assignment->ToString() : "nullptr"); } - // functional_hlo_runner assigns 0 for device assigments for multi-host + // functional_hlo_runner assigns 0 for device assignments for multi-host // cases. In this case we ignore the assignment. // See LoadAndCompile in functional_hlo_runner.cc for more details. const bool has_device_assignment = @@ -384,6 +385,50 @@ bool IsAllReplicasLocal(int64_t gpus_per_host, }); } +bool IsAllReplicasLocal(const GpuTopology& gpu_topology, + const DebugOptions& debug_options, + std::optional op_type, + absl::Span replica_groups, + CollectiveOpGroupMode group_mode, + const DeviceAssignment* device_assignment) { + int64_t gpus_per_host = + IsCrossHostOneShotKernelEnabled(debug_options, op_type) + ? gpu_topology.slice_size() + : gpu_topology.num_devices_per_process(); + return IsAllReplicasLocal(gpus_per_host, replica_groups, group_mode, + device_assignment); +} + +absl::StatusOr IsAllReplicasLocal( + const GpuTopology& gpu_topology, const HloInstruction& instruction, + const DeviceAssignment* device_assignment) { + auto collective = DynCast(&instruction); + CHECK(collective != nullptr) + << "Instruction is not a collective instruction: " << instruction.name(); + auto group_mode_status = GetCollectiveOpGroupMode(collective); + if (!group_mode_status.ok()) { + LOG(WARNING) << "Failed to get collective op group mode for " + << instruction.name() << ": " << group_mode_status.status(); + return false; + } + const DebugOptions& debug_options = + instruction.GetModule() != nullptr + ? instruction.GetModule()->config().debug_options() + : DebugOptions::default_instance(); + auto op_type = GetCollectiveOpType(&instruction); + return IsAllReplicasLocal(gpu_topology, debug_options, op_type, + collective->replica_groups(), *group_mode_status, + device_assignment); +} + +bool IsAllReplicasLocal(const GpuTopology& gpu_topology, + absl::Span replica_groups, + CollectiveOpGroupMode group_mode, + const DeviceAssignment* device_assignment) { + return IsAllReplicasLocal(gpu_topology.num_devices_per_process(), + replica_groups, group_mode, device_assignment); +} + bool IsTritonCollectiveKernel( CollectiveBackendConfig::CollectiveKernelStrategy kernel_strategy) { return kernel_strategy == diff --git a/third_party/xla/xla/backends/gpu/transforms/collectives/collective_ops_utils.h b/third_party/xla/xla/backends/gpu/transforms/collectives/collective_ops_utils.h index b0f2aa38379b70..74a4cd98f6998d 100644 --- a/third_party/xla/xla/backends/gpu/transforms/collectives/collective_ops_utils.h +++ b/third_party/xla/xla/backends/gpu/transforms/collectives/collective_ops_utils.h @@ -17,6 +17,7 @@ limitations under the License. #define XLA_BACKENDS_GPU_TRANSFORMS_COLLECTIVES_COLLECTIVE_OPS_UTILS_H_ #include +#include #include "absl/container/flat_hash_set.h" #include "absl/status/statusor.h" @@ -26,8 +27,10 @@ limitations under the License. #include "xla/hlo/ir/hlo_opcode.h" #include "xla/service/device_assignment.h" #include "xla/service/gpu/backend_configs.pb.h" +#include "xla/service/gpu_topology.h" #include "xla/service/hlo_module_config.h" #include "xla/stream_executor/device_description.h" +#include "xla/xla.pb.h" namespace xla { namespace gpu { @@ -101,6 +104,30 @@ bool IsAllReplicasLocal(int64_t gpus_per_host, CollectiveOpGroupMode group_mode, const DeviceAssignment* device_assignment = nullptr); +// Returns true if all replicas in every replica group of the collective +// are located on the same host (or within a single GB cluster host / partition +// if xla_gpu_unsupported_use_cross_host_one_shot_kernel is enabled for the +// collective). +bool IsAllReplicasLocal(const GpuTopology& gpu_topology, + const DebugOptions& debug_options, + std::optional op_type, + absl::Span replica_groups, + CollectiveOpGroupMode group_mode, + const DeviceAssignment* device_assignment = nullptr); + +// Returns true if all replicas in every replica group of the collective +// are located on the same host (or within a single GB cluster host / partition +// if xla_gpu_unsupported_use_cross_host_one_shot_kernel is enabled for the +// collective). +absl::StatusOr IsAllReplicasLocal( + const GpuTopology& gpu_topology, const HloInstruction& instruction, + const DeviceAssignment* device_assignment = nullptr); + +bool IsAllReplicasLocal(const GpuTopology& gpu_topology, + absl::Span replica_groups, + CollectiveOpGroupMode group_mode, + const DeviceAssignment* device_assignment = nullptr); + // Returns true if the instruction was generated by the SPMD partitioner. // Checks both the kSpmdGeneratedAttr frontend attribute and the // is_spmd_generated field in CollectiveBackendConfig (set when collectives diff --git a/third_party/xla/xla/backends/gpu/transforms/collectives/collective_ops_utils_test.cc b/third_party/xla/xla/backends/gpu/transforms/collectives/collective_ops_utils_test.cc index 5d89faa86d390b..07f8d6790c7ab4 100644 --- a/third_party/xla/xla/backends/gpu/transforms/collectives/collective_ops_utils_test.cc +++ b/third_party/xla/xla/backends/gpu/transforms/collectives/collective_ops_utils_test.cc @@ -15,6 +15,8 @@ limitations under the License. #include "xla/backends/gpu/transforms/collectives/collective_ops_utils.h" +#include + #include #include #include "absl/status/status_matchers.h" @@ -23,8 +25,10 @@ limitations under the License. #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_instructions.h" #include "xla/hlo/parser/hlo_parser.h" +#include "xla/service/device_assignment.h" #include "xla/service/gpu/backend_configs.pb.h" #include "xla/service/gpu/gpu_device_info_for_tests.h" +#include "xla/service/gpu_topology.h" #include "xla/stream_executor/cuda/cuda_compute_capability.h" #include "xla/stream_executor/device_description.h" #include "xla/tsl/platform/statusor.h" @@ -704,4 +708,145 @@ TEST(IsSpmdGeneratedTest, ReturnsTrueWhenBackendConfigSet) { EXPECT_TRUE(IsSpmdGenerated(*ar)); } +TEST(IsAllReplicasLocalTest, SingleHostSingleProcess) { + GpuTopology topology( + /*platform_version=*/"", + /*num_partitions=*/1, + /*num_hosts_per_partition=*/1, + /*num_devices_per_host=*/8, + /*gpu_target_config=*/std::nullopt, + /*host_target_machine_options=*/std::nullopt, + /*num_devices_per_process=*/8); + + ReplicaGroup group; + group.add_replica_ids(0); + group.add_replica_ids(1); + + EXPECT_TRUE(IsAllReplicasLocal( + topology, {group}, + CollectiveOpGroupMode::COLLECTIVE_OP_GROUP_MODE_CROSS_REPLICA)); +} + +TEST(IsAllReplicasLocalTest, + SingleGBClusterHostMultiProcess_DefaultDisabled_ReturnsFalse) { + // A single GB cluster host with 2 processes (each having 1 device). + // num_hosts_per_partition = 2, num_devices_per_host = 1, slice_size = 2. + GpuTopology topology( + /*platform_version=*/"", + /*num_partitions=*/1, + /*num_hosts_per_partition=*/2, + /*num_devices_per_host=*/1, + /*gpu_target_config=*/std::nullopt, + /*host_target_machine_options=*/std::nullopt, + /*num_devices_per_process=*/1); + + DeviceAssignment da(2, 1); + da(0, 0) = 0; + da(1, 0) = 1; + + ReplicaGroup group; + group.add_replica_ids(0); + group.add_replica_ids(1); + + // Without the flag enabled, multi-process within a GB cluster host is NOT + // local. + DebugOptions debug_options; + EXPECT_FALSE(IsAllReplicasLocal( + topology, debug_options, DebugOptions::ALLREDUCE, {group}, + CollectiveOpGroupMode::COLLECTIVE_OP_GROUP_MODE_CROSS_REPLICA, &da)); +} + +TEST(IsAllReplicasLocalTest, + SingleGBClusterHostMultiProcess_FlagEnabled_ReturnsTrue) { + GpuTopology topology( + /*platform_version=*/"", + /*num_partitions=*/1, + /*num_hosts_per_partition=*/2, + /*num_devices_per_host=*/1, + /*gpu_target_config=*/std::nullopt, + /*host_target_machine_options=*/std::nullopt, + /*num_devices_per_process=*/1); + + DeviceAssignment da(2, 1); + da(0, 0) = 0; + da(1, 0) = 1; + + ReplicaGroup group; + group.add_replica_ids(0); + group.add_replica_ids(1); + + // With ALLREDUCE enabled in + // xla_gpu_unsupported_use_cross_host_one_shot_kernel: + DebugOptions debug_options; + debug_options.add_xla_gpu_unsupported_use_cross_host_one_shot_kernel( + DebugOptions::ALLREDUCE); + EXPECT_TRUE(IsAllReplicasLocal( + topology, debug_options, DebugOptions::ALLREDUCE, {group}, + CollectiveOpGroupMode::COLLECTIVE_OP_GROUP_MODE_CROSS_REPLICA, &da)); + + // With ALLCOLLECTIVES enabled: + DebugOptions all_collectives_opts; + all_collectives_opts.add_xla_gpu_unsupported_use_cross_host_one_shot_kernel( + DebugOptions::ALLCOLLECTIVES); + EXPECT_TRUE(IsAllReplicasLocal( + topology, all_collectives_opts, DebugOptions::ALLREDUCE, {group}, + CollectiveOpGroupMode::COLLECTIVE_OP_GROUP_MODE_CROSS_REPLICA, &da)); +} + +TEST(IsAllReplicasLocalTest, + SingleGBClusterHostMultiProcess_FlagEnabledForDifferentCollective) { + GpuTopology topology( + /*platform_version=*/"", + /*num_partitions=*/1, + /*num_hosts_per_partition=*/2, + /*num_devices_per_host=*/1, + /*gpu_target_config=*/std::nullopt, + /*host_target_machine_options=*/std::nullopt, + /*num_devices_per_process=*/1); + + DeviceAssignment da(2, 1); + da(0, 0) = 0; + da(1, 0) = 1; + + ReplicaGroup group; + group.add_replica_ids(0); + group.add_replica_ids(1); + + // With only ALLGATHER enabled in the flag, ALLREDUCE should return false: + DebugOptions debug_options; + debug_options.add_xla_gpu_unsupported_use_cross_host_one_shot_kernel( + DebugOptions::ALLGATHER); + EXPECT_FALSE(IsAllReplicasLocal( + topology, debug_options, DebugOptions::ALLREDUCE, {group}, + CollectiveOpGroupMode::COLLECTIVE_OP_GROUP_MODE_CROSS_REPLICA, &da)); +} + +TEST(IsAllReplicasLocalTest, CrossGBClusterPartitionReturnsFalse) { + // 2 GB cluster partitions with 2 devices each (slice_size = 2). + GpuTopology topology( + /*platform_version=*/"", + /*num_partitions=*/2, + /*num_hosts_per_partition=*/2, + /*num_devices_per_host=*/1, + /*gpu_target_config=*/std::nullopt, + /*host_target_machine_options=*/std::nullopt, + /*num_devices_per_process=*/1); + + DeviceAssignment da(2, 1); + da(0, 0) = 0; + da(1, 0) = 2; // Device 2 is in partition 1 + + ReplicaGroup group; + group.add_replica_ids(0); + group.add_replica_ids(1); + + DebugOptions debug_options; + debug_options.add_xla_gpu_unsupported_use_cross_host_one_shot_kernel( + DebugOptions::ALLREDUCE); + + EXPECT_FALSE(IsAllReplicasLocal( + topology, debug_options, DebugOptions::ALLREDUCE, {group}, + CollectiveOpGroupMode::COLLECTIVE_OP_GROUP_MODE_CROSS_REPLICA, &da)); +} + } // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/transforms/fusion_block_level_rewriter.cc b/third_party/xla/xla/backends/gpu/transforms/fusion_block_level_rewriter.cc index 8f0a0476a3c6aa..1304cdeb81d8a7 100644 --- a/third_party/xla/xla/backends/gpu/transforms/fusion_block_level_rewriter.cc +++ b/third_party/xla/xla/backends/gpu/transforms/fusion_block_level_rewriter.cc @@ -178,15 +178,18 @@ absl::StatusOr ShouldTryRewriteFusion( const DebugOptions& debug_options = fusion->GetModule()->config().debug_options(); - const bool can_emit_same_shape_multi_output_fusion = - IsSameShapeMultiOutputFusion(*fusion, + + // Same-shape multi-output fusions require the new tiling propagation + // infrastructure. + if (IsSameShapeMultiOutputFusion(*fusion, Shape::Equal().IgnoreElementType()) && debug_options .xla_gpu_experimental_enable_same_shape_multi_output_fusion() && - debug_options.xla_gpu_experimental_enable_tiling_propagation(); + debug_options.xla_gpu_experimental_enable_tiling_propagation()) { + return true; + } if (fusion->IsMultiOutputFusion() && - !can_emit_same_shape_multi_output_fusion && !debug_options.xla_gpu_unsupported_enable_triton_multi_output_fusion()) { return false; } diff --git a/third_party/xla/xla/backends/gpu/transforms/fusion_block_level_rewriter_test.cc b/third_party/xla/xla/backends/gpu/transforms/fusion_block_level_rewriter_test.cc index 159b4654526786..6b5c098f98a3f8 100644 --- a/third_party/xla/xla/backends/gpu/transforms/fusion_block_level_rewriter_test.cc +++ b/third_party/xla/xla/backends/gpu/transforms/fusion_block_level_rewriter_test.cc @@ -41,6 +41,7 @@ License. #include "xla/service/gpu/gpu_device_info_for_tests.h" #include "xla/service/gpu/ir_emission_utils.h" #include "xla/service/hlo_cost_analysis.h" +#include "xla/service/hlo_module_config.h" #include "xla/stream_executor/cuda/cuda_compute_capability.h" #include "xla/stream_executor/device_description.h" #include "xla/tsl/platform/statusor.h" @@ -460,5 +461,46 @@ ENTRY entry { absl_testing::IsOkAndHolds(false)); } +TEST_F(FusionBlockLevelRewriterTestBase, + RewritesSameShapeMultiOutputFusionWithoutGeneralBlockLevelRewriter) { + const absl::string_view hlo_text = R"hlo( +f { + p0 = f32[10,10] parameter(0) + p1 = f32[10,10] parameter(1) + add = f32[10,10] add(p0, p1) + sub = f32[10,10] subtract(p0, p1) + ROOT multi_output = (f32[10,10], f32[10,10]) tuple(add, sub) +} + +ENTRY entry { + p0 = f32[10,10] parameter(0) + p1 = f32[10,10] parameter(1) + ROOT fusion = (f32[10,10], f32[10,10]) fusion(p0, p1), kind=kLoop, calls=f +})hlo"; + DebugOptions debug_options = + HloHardwareIndependentTestBase::GetDebugOptionsForTest(); + debug_options.set_xla_gpu_experimental_enable_fusion_block_level_rewriter( + false); + debug_options.set_xla_gpu_experimental_enable_tiling_propagation(true); + debug_options.set_xla_gpu_experimental_enable_same_shape_multi_output_fusion( + true); + + HloModuleConfig config; + config.set_debug_options(debug_options); + + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(hlo_text, config)); + + EXPECT_THAT( + FusionBlockLevelRewriter(device_info_, HloCostAnalysis::DefaultShapeSize, + &mlir_context_) + .Run(module.get()), + absl_testing::IsOkAndHolds(true)); + const HloInstruction* root = module->entry_computation()->root_instruction(); + EXPECT_EQ(root->opcode(), HloOpcode::kFusion); + EXPECT_EQ(root->fusion_kind(), HloInstruction::FusionKind::kCustom); + EXPECT_TRUE(HasTritonBlockLevelFusionConfig(root)); +} + } // namespace } // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/transforms/gemm_fusion.cc b/third_party/xla/xla/backends/gpu/transforms/gemm_fusion.cc index cd8591f85691f6..e53bc1324f3f3c 100644 --- a/third_party/xla/xla/backends/gpu/transforms/gemm_fusion.cc +++ b/third_party/xla/xla/backends/gpu/transforms/gemm_fusion.cc @@ -1459,6 +1459,29 @@ FusionDecision ShouldFuseUser(const HloInstruction* user, break; } + int64_t src_operand_index = user->operand_index(fusion); + for (int i = 0; i < user->operand_count(); ++i) { + const HloInstruction* operand = user->operand(i); + // Skip source operand. + if (i == src_operand_index) { + continue; + } + // Currently only + // - effective parameters + // - broadcasts of effective parameters + // - broadcasts of scalars + // are accepted as other inputs of non-unary operations in + // the output fusion. + if ((operand->opcode() == HloOpcode::kBroadcast && + (ShapeUtil::IsScalar(operand->operand(0)->shape()) || + hlo_query::IsEffectiveParameter(*operand->operand(0)))) || + hlo_query::IsEffectiveParameter(*operand)) { + continue; + } + return FusionDecision::Forbid( + "Has multiple inputs - not properly analyzed yet."); + } + if (!triton_fusion::IsOutputWorthFusing(original_user)) { return FusionDecision::Forbid( "Not obviously profitable to fuse as output."); diff --git a/third_party/xla/xla/backends/gpu/transforms/gemm_fusion_test.cc b/third_party/xla/xla/backends/gpu/transforms/gemm_fusion_test.cc index 7022507b428b44..88598eebc95896 100644 --- a/third_party/xla/xla/backends/gpu/transforms/gemm_fusion_test.cc +++ b/third_party/xla/xla/backends/gpu/transforms/gemm_fusion_test.cc @@ -435,10 +435,10 @@ ENTRY e { p1 = f32[32,7] parameter(1) d = f32[6,7] dot(p0, p1), lhs_contracting_dims={0}, rhs_contracting_dims={0} - p2 = s8[3,14] parameter(2) - c1 = f32[3,14] convert(p2) - b1 = f32[6,7] bitcast(c1) - ROOT a = f32[6,7] add(d, b1) + p2 = f32[2,3] parameter(2) + b1 = f32[6] bitcast(p2) + br = f32[6,7] broadcast(b1), dimensions={0} + ROOT a = f32[6,7] add(d, br) })")); ASSERT_THAT(GemmFusion(gpu_version_).Run(module.get()), IsOkAndHolds(true)); EXPECT_THAT(module->entry_computation()->root_instruction(), @@ -456,14 +456,14 @@ ENTRY e { d = f32[6,7] dot(p0, p1), lhs_contracting_dims={0}, rhs_contracting_dims={0} b1 = f32[3,14] bitcast(d) - p2 = s8[3,14] parameter(2) - c1 = f32[3,14] convert(p2) - ROOT a = f32[3,14] add(b1, c1) + p2 = f32[3,14] parameter(2) + ROOT a = f32[3,14] add(b1, p2) })")); ASSERT_THAT(GemmFusion(gpu_version_).Run(module.get()), IsOkAndHolds(true)); - EXPECT_THAT(module->entry_computation()->root_instruction(), - GmockMatch(m::Bitcast(m::Fusion(m::Parameter(), m::Parameter(), - m::Bitcast(m::Parameter()))))); + EXPECT_THAT( + module->entry_computation()->root_instruction(), + GmockMatch(m::Bitcast(m::Fusion(m::Bitcast(m::Parameter()), + m::Parameter(), m::Parameter())))); } TEST_P(GemmFusionTestV2, BitcastIsHoistedAboveConstants) { diff --git a/third_party/xla/xla/backends/gpu/transforms/layout_assignment_a100.hlo b/third_party/xla/xla/backends/gpu/transforms/layout_assignment_a100.hlo index 3b5c92a19388d2..cf05e52efb9dac 100644 --- a/third_party/xla/xla/backends/gpu/transforms/layout_assignment_a100.hlo +++ b/third_party/xla/xla/backends/gpu/transforms/layout_assignment_a100.hlo @@ -1,14 +1,16 @@ -// RUN: hlo-opt %s --platform=gpu --stage=hlo --xla_gpu_target_config_filename=%S/../../../backends/gpu/target_config/specs/a100_pcie_80.txtpb --split-input-file | FileCheck %s +// RUN: hlo-opt %s --platform=gpu --stage=hlo --xla_gpu_target_config_filename=%S/../../../backends/gpu/target_config/specs/a100_pcie_80.txtpb --split-input-file | FileCheck %s +// CHECK: %conv_fprop_fusion_comp +// CHECK: convolution +// CHECK-SAME: window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_o01i->b01f, convolution_kind=fprop // CHECK: %wrapped_transpose_computation // CHECK-NEXT: bf16[3,3,16,32]{3,2,1,0} parameter(0) -// CHECK-NEXT: bf16[32,3,3,16]{3,2,1,0} transpose +// CHECK-NEXT: ROOT {{.*}} = bf16[32,3,3,16]{3,2,1,0} transpose // CHECK-SAME: dimensions={3,0,1,2} -// CHECK: (bf16[1,64,64,32]{3,2,1,0}, u8[0]{0}) custom-call -// CHECK-SAME: window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convForward - +// CHECK: ENTRY %main +// CHECK: %[[WRAPPED_TRANSPOSE:.*]] = bf16[32,3,3,16]{3,2,1,0} fusion({{.*}}), {{.*}}calls=%wrapped_transpose_computation +// CHECK: ROOT {{.*}} = bf16[1,64,64,32]{3,2,1,0} fusion({{.*}}, %[[WRAPPED_TRANSPOSE]]), kind=kCustom, calls=%conv_fprop_fusion_comp HloModule ConvCuDNN - ENTRY main { Arg_0.1 = bf16[1,64,64,16]{3,2,1,0} parameter(0), sharding={replicated} Arg_1.2 = bf16[3,3,16,32]{3,2,1,0} parameter(1), sharding={replicated} diff --git a/third_party/xla/xla/backends/gpu/transforms/layout_assignment_h100.hlo b/third_party/xla/xla/backends/gpu/transforms/layout_assignment_h100.hlo index 1b82bb55c80b2b..07634f1df086a9 100644 --- a/third_party/xla/xla/backends/gpu/transforms/layout_assignment_h100.hlo +++ b/third_party/xla/xla/backends/gpu/transforms/layout_assignment_h100.hlo @@ -1,14 +1,16 @@ -// RUN: hlo-opt %s --platform=gpu --stage=hlo --xla_gpu_target_config_filename=%S/../../../backends/gpu/target_config/specs/h100_sxm.txtpb --split-input-file | FileCheck %s +// RUN: hlo-opt %s --platform=gpu --stage=hlo --xla_gpu_target_config_filename=%S/../../../backends/gpu/target_config/specs/h100_sxm.txtpb --split-input-file | FileCheck %s +// CHECK: %conv_fprop_fusion_comp +// CHECK: convolution +// CHECK-SAME: window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_o01i->b01f, convolution_kind=fprop // CHECK: %wrapped_transpose_computation // CHECK-NEXT: f8e4m3fn[3,3,16,32]{3,2,1,0} parameter(0) -// CHECK-NEXT: f8e4m3fn[32,3,3,16]{3,2,1,0} transpose +// CHECK-NEXT: ROOT {{.*}} = f8e4m3fn[32,3,3,16]{3,2,1,0} transpose // CHECK-SAME: dimensions={3,0,1,2} -// CHECK: (f8e4m3fn[1,64,64,32]{3,2,1,0}, u8[0]{0}) custom-call -// CHECK-SAME: window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convForward - +// CHECK: ENTRY %main +// CHECK: %[[WRAPPED_TRANSPOSE:.*]] = f8e4m3fn[32,3,3,16]{3,2,1,0} fusion({{.*}}), {{.*}}calls=%wrapped_transpose_computation +// CHECK: ROOT {{.*}} = f8e4m3fn[1,64,64,32]{3,2,1,0} fusion({{.*}}, %[[WRAPPED_TRANSPOSE]]), kind=kCustom, calls=%conv_fprop_fusion_comp HloModule ConvCuDNN - ENTRY main { Arg_0.1 = f8e4m3fn[1,64,64,16]{3,2,1,0} parameter(0), sharding={replicated} Arg_1.2 = f8e4m3fn[3,3,16,32]{3,2,1,0} parameter(1), sharding={replicated} diff --git a/third_party/xla/xla/backends/gpu/transforms/layout_assignment_v100.hlo b/third_party/xla/xla/backends/gpu/transforms/layout_assignment_v100.hlo index d5baeb8a42af7d..193e4c27c57fa2 100644 --- a/third_party/xla/xla/backends/gpu/transforms/layout_assignment_v100.hlo +++ b/third_party/xla/xla/backends/gpu/transforms/layout_assignment_v100.hlo @@ -1,14 +1,16 @@ -// RUN: hlo-opt %s --platform=gpu --stage=hlo --xla_gpu_target_config_filename=%S/../../../backends/gpu/target_config/specs/v100.txtpb --split-input-file | FileCheck %s +// RUN: hlo-opt %s --platform=gpu --stage=hlo --xla_gpu_target_config_filename=%S/../../../backends/gpu/target_config/specs/v100.txtpb --split-input-file | FileCheck %s +// CHECK: %conv_fprop_fusion_comp +// CHECK: convolution +// CHECK-SAME: window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_o01i->b01f, convolution_kind=fprop // CHECK: %wrapped_transpose_computation // CHECK-NEXT: f16[3,3,16,32]{3,2,1,0} parameter(0) -// CHECK-NEXT: f16[32,3,3,16]{3,2,1,0} transpose +// CHECK-NEXT: ROOT {{.*}} = f16[32,3,3,16]{3,2,1,0} transpose // CHECK-SAME: dimensions={3,0,1,2} -// CHECK: (f16[1,64,64,32]{3,2,1,0}, u8[0]{0}) custom-call -// CHECK-SAME: window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convForward - +// CHECK: ENTRY %main +// CHECK: %[[WRAPPED_TRANSPOSE:.*]] = f16[32,3,3,16]{3,2,1,0} fusion({{.*}}), {{.*}}calls=%wrapped_transpose_computation +// CHECK: ROOT {{.*}} = f16[1,64,64,32]{3,2,1,0} fusion({{.*}}, %[[WRAPPED_TRANSPOSE]]), kind=kCustom, calls=%conv_fprop_fusion_comp HloModule ConvCuDNN - ENTRY main { Arg_0.1 = f16[1,64,64,16]{3,2,1,0} parameter(0), sharding={replicated} Arg_1.2 = f16[3,3,16,32]{3,2,1,0} parameter(1), sharding={replicated} diff --git a/third_party/xla/xla/backends/gpu/transforms/sort_rewriter_test.cc b/third_party/xla/xla/backends/gpu/transforms/sort_rewriter_test.cc index 66d5226901ea28..6dba87a9612b75 100644 --- a/third_party/xla/xla/backends/gpu/transforms/sort_rewriter_test.cc +++ b/third_party/xla/xla/backends/gpu/transforms/sort_rewriter_test.cc @@ -72,7 +72,7 @@ std::string GetNumpyOrderComparator( rhs_is_zero = pred[] compare(rhs, c_zero), direction=EQ rhs_no_neg_zero = $0[] select(rhs_is_zero, c_zero, rhs) rhs_no_neg_zero_or_nan = $0[] select(rhs_is_nan, c_nan, rhs_no_neg_zero) - ROOT compare = pred[] compare(lhs_no_neg_zero_or_nan, rhs_no_neg_zero_or_nan), direction=$1, type=TOTALORDER + ROOT compare = pred[] compare(lhs_no_neg_zero_or_nan, rhs_no_neg_zero_or_nan), direction=$1, order=TOTAL )"; return absl::StrCat("numpy_order_comparator {\n", params, diff --git a/third_party/xla/xla/backends/gpu/transforms/topk_specializer.cc b/third_party/xla/xla/backends/gpu/transforms/topk_specializer.cc index c6885af8688fb3..f5bd92b7c72206 100644 --- a/third_party/xla/xla/backends/gpu/transforms/topk_specializer.cc +++ b/third_party/xla/xla/backends/gpu/transforms/topk_specializer.cc @@ -17,6 +17,7 @@ limitations under the License. #include +#include #include #include #include @@ -26,6 +27,7 @@ limitations under the License. #include "absl/log/check.h" #include "absl/log/log.h" #include "absl/status/status.h" +#include "absl/status/status_macros.h" #include "absl/strings/string_view.h" #include "xla/hlo/ir/dfs_hlo_visitor_with_default.h" #include "xla/hlo/ir/hlo_casting_utils.h" @@ -34,6 +36,8 @@ limitations under the License. #include "xla/hlo/ir/hlo_instruction_utils.h" #include "xla/hlo/ir/hlo_instructions.h" #include "xla/hlo/ir/hlo_module.h" +#include "xla/hlo/ir/hlo_opcode.h" +#include "xla/literal_util.h" #include "xla/primitive_util.h" #include "xla/service/hlo.pb.h" #include "xla/service/tuple_util.h" @@ -47,6 +51,234 @@ namespace xla::gpu { namespace { +// Broadcast a 32-bit scalar to a target shape. +HloInstruction* BroadcastU32(HloComputation* comp, const Shape& target_shape, + uint32_t value) { + HloInstruction* constant = comp->AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(value))); + return comp->AddInstruction(HloInstruction::CreateBroadcast( + ShapeUtil::ChangeElementType(target_shape, U32), constant, {})); +} + +// Broadcast a 64-bit scalar to a target shape. +HloInstruction* BroadcastU64(HloComputation* comp, const Shape& target_shape, + uint64_t value) { + HloInstruction* constant = comp->AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(value))); + return comp->AddInstruction(HloInstruction::CreateBroadcast( + ShapeUtil::ChangeElementType(target_shape, U64), constant, {})); +} + +// Emits HLO to pack F32 values and Iota indices into U64 priority keys. +HloInstruction* BuildPackF32ToU64(HloInstruction* data, HloComputation* comp) { + const Shape& shape = data->shape(); + int64_t iota_dim = shape.dimensions().size() - 1; + + Shape u32_shape = ShapeUtil::ChangeElementType(shape, U32); + Shape s32_shape = ShapeUtil::ChangeElementType(shape, S32); + Shape u64_shape = ShapeUtil::ChangeElementType(shape, U64); + + // 1. Generate reversed indices: 0xFFFFFFFF - iota + HloInstruction* iota_u32 = + comp->AddInstruction(HloInstruction::CreateIota(u32_shape, iota_dim)); + HloInstruction* broadcast_ff = BroadcastU32(comp, shape, 0xFFFFFFFF); + HloInstruction* iota_neg = comp->AddInstruction(HloInstruction::CreateBinary( + u32_shape, HloOpcode::kSubtract, broadcast_ff, iota_u32)); + + // 2. Pure Bitwise Radix Float Flip (F32 -> U32) + // s32_val = bitcast(data) + HloInstruction* s32_val = comp->AddInstruction( + HloInstruction::CreateBitcastConvert(s32_shape, data)); + + // sign_smeared = s32_val >> 31 (arithmetic shift) + HloInstruction* const_31_s32 = comp->AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(31))); + HloInstruction* broadcast_31_s32 = comp->AddInstruction( + HloInstruction::CreateBroadcast(s32_shape, const_31_s32, {})); + HloInstruction* sign_smeared = comp->AddInstruction( + HloInstruction::CreateBinary(s32_shape, HloOpcode::kShiftRightArithmetic, + s32_val, broadcast_31_s32)); + + // mask = bitcast(sign_smeared) | 0x80000000 + HloInstruction* sign_smeared_u32 = comp->AddInstruction( + HloInstruction::CreateBitcastConvert(u32_shape, sign_smeared)); + HloInstruction* broadcast_8 = BroadcastU32(comp, shape, 0x80000000); + HloInstruction* mask = comp->AddInstruction(HloInstruction::CreateBinary( + u32_shape, HloOpcode::kOr, sign_smeared_u32, broadcast_8)); + + // radix_key = bitcast(data) ^ mask + HloInstruction* u32_val = comp->AddInstruction( + HloInstruction::CreateBitcastConvert(u32_shape, data)); + HloInstruction* radix_key = comp->AddInstruction( + HloInstruction::CreateBinary(u32_shape, HloOpcode::kXor, u32_val, mask)); + + // 3. Pack into U64: (radix_key << 32) | iota_neg + HloInstruction* val_u64 = + comp->AddInstruction(HloInstruction::CreateConvert(u64_shape, radix_key)); + HloInstruction* broadcast_32_u64 = BroadcastU64(comp, shape, 32); + HloInstruction* val_u64_top = + comp->AddInstruction(HloInstruction::CreateBinary( + u64_shape, HloOpcode::kShiftLeft, val_u64, broadcast_32_u64)); + + HloInstruction* iota_neg_u64 = + comp->AddInstruction(HloInstruction::CreateConvert(u64_shape, iota_neg)); + + return comp->AddInstruction(HloInstruction::CreateBinary( + u64_shape, HloOpcode::kOr, val_u64_top, iota_neg_u64)); +} + +// Emits HLO to unpack U64 priority keys back to F32 values. +HloInstruction* BuildUnpackU64ToF32(HloInstruction* u64_values, + HloComputation* comp) { + const Shape& shape = u64_values->shape(); + Shape u32_shape = ShapeUtil::ChangeElementType(shape, U32); + Shape f32_shape = ShapeUtil::ChangeElementType(shape, F32); + + // 1. Shift right by 32 and cast to U32 to extract the radix_key + HloInstruction* broadcast_32_u64 = BroadcastU64(comp, shape, 32); + HloInstruction* rshift = comp->AddInstruction(HloInstruction::CreateBinary( + shape, HloOpcode::kShiftRightLogical, u64_values, broadcast_32_u64)); + HloInstruction* radix_key = + comp->AddInstruction(HloInstruction::CreateConvert(u32_shape, rshift)); + + // 2. Isolate the MSB: msb = radix_key >> 31 (Logical shift) + HloInstruction* const_31_u32 = comp->AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(31))); + HloInstruction* broadcast_31_u32 = comp->AddInstruction( + HloInstruction::CreateBroadcast(u32_shape, const_31_u32, {})); + HloInstruction* msb = comp->AddInstruction(HloInstruction::CreateBinary( + u32_shape, HloOpcode::kShiftRightLogical, radix_key, broadcast_31_u32)); + + // 3. Reconstruct mask: msb_minus_one = msb - 1 + // If MSB was 1: 1 - 1 = 0x00000000. + // If MSB was 0: 0 - 1 = 0xFFFFFFFF (unsigned underflow). + HloInstruction* broadcast_1 = BroadcastU32(comp, shape, 1); + HloInstruction* msb_minus_one = + comp->AddInstruction(HloInstruction::CreateBinary( + u32_shape, HloOpcode::kSubtract, msb, broadcast_1)); + + // 4. unmask = msb_minus_one | 0x80000000 + HloInstruction* broadcast_8 = BroadcastU32(comp, shape, 0x80000000); + HloInstruction* unmask = comp->AddInstruction(HloInstruction::CreateBinary( + u32_shape, HloOpcode::kOr, msb_minus_one, broadcast_8)); + + // 5. Unflip the bits: original_u32 = radix_key ^ unmask + HloInstruction* original_u32 = + comp->AddInstruction(HloInstruction::CreateBinary( + u32_shape, HloOpcode::kXor, radix_key, unmask)); + + // 6. Bitcast back to F32 + return comp->AddInstruction( + HloInstruction::CreateBitcastConvert(f32_shape, original_u32)); +} + +// Checks if we can safely route stable TopK to RAFT using the Uint64 adapter. +bool ShouldRewriteStableTopKToUint64(HloCustomCallInstruction* topk) { + if (!hlo_instruction_utils::IsTopKStable(topk)) { + return false; + } + + Shape data_shape = topk->operand(0)->shape(); + PrimitiveType dtype = data_shape.element_type(); + if (!(dtype == F32 || dtype == BF16)) { + return false; // Only F32 and BF16 are supported for now. + } + + bool has_batch = data_shape.dimensions().size() == 2; + size_t batch = has_batch ? data_shape.dimensions(0) : 1; + size_t n = data_shape.dimensions(has_batch ? 1 : 0); + size_t k = topk->shape().tuple_shapes(0).dimensions(has_batch ? 1 : 0); + + // For small N, sort + slice is faster than RAFT select_k + if (n < 1024) { + return false; + } + + double ratio = static_cast(k) / n; + if (ratio >= 0.85) { + return false; + } + + // Use built-in stable XLA GPU TopK kernel if n/k ranges are supported + if (n >= 1024 && k <= 16) { + return false; + } + + // Upper bounds for using RAFT select_k. + // The heuristic for deciding when to use Raft select_k versus Sort + Slice + // was developed as part of the initial research in b/409009349 + size_t max_k = 128; + if (dtype == F32) { + max_k = 128; + if (batch >= 64 && n >= 16384) { + max_k = 256; + } + } else if (dtype == BF16) { + max_k = 128; + if (batch >= 16 && n >= 65536) { + max_k = 256; + } + if (batch >= 64 && batch <= 128 && n >= 8192 && n <= 32768) { + max_k = 64; + } + } + if (k > max_k) { + return false; + } + return true; +} + +absl::StatusOr RewriteStableTopKToUint64( + HloCustomCallInstruction* topk) { + HloComputation* comp = topk->parent(); + HloInstruction* data = topk->mutable_operand(0); + PrimitiveType original_type = data->shape().element_type(); + + // 1. Pack + // If BF16, upcast to F32 first to match F32 bitmath logic. + if (original_type == BF16) { + data = comp->AddInstruction(HloInstruction::CreateConvert( + ShapeUtil::ChangeElementType(data->shape(), F32), data)); + } + HloInstruction* packed_u64 = BuildPackF32ToU64(data, comp); + + // 2. Create the specialized __gpu$TopK custom call + Shape k_shape = + ShapeUtil::ChangeElementType(topk->shape().tuple_shapes(0), U64); + Shape idx_shape = topk->shape().tuple_shapes(1); + Shape new_cc_shape = ShapeUtil::MakeTupleShape({k_shape, idx_shape}); + + HloInstruction* new_topk = + comp->AddInstruction(HloInstruction::CreateCustomCall( + new_cc_shape, {packed_u64}, topk->to_apply(), "__gpu$TopK", "", + CustomCallApiVersion::API_VERSION_TYPED_FFI)); + + // The packed U64 keys guarantee uniqueness, making ties impossible. + // Therefore, the inner TopK operation no longer requires stability to + // produce a stable overall result. We clear the is_stable flag so the + // backend can freely route this to the fast unstable topk kernel (RAFT lib). + new_topk->set_raw_backend_config_string("{is_stable = false}"); + + // 3. Unpack values and retain indices + HloInstruction* u64_vals = comp->AddInstruction( + HloInstruction::CreateGetTupleElement(k_shape, new_topk, 0)); + HloInstruction* indices = comp->AddInstruction( + HloInstruction::CreateGetTupleElement(idx_shape, new_topk, 1)); + + // Unpack to F32. + HloInstruction* unpacked_data = BuildUnpackU64ToF32(u64_vals, comp); + + // If the original input was BF16, downcast the F32 result back to BF16. + if (original_type == BF16) { + unpacked_data = comp->AddInstruction(HloInstruction::CreateConvert( + ShapeUtil::ChangeElementType(unpacked_data->shape(), BF16), + unpacked_data)); + } + + return comp->AddInstruction( + HloInstruction::CreateTuple({unpacked_data, indices})); +} + absl::StatusOr SmallBufferOptimization( HloCustomCallInstruction* topk, bool is_cuda) { Shape data_shape = topk->operand(0)->shape(); @@ -131,6 +363,12 @@ class SpecializeTopkVisitor : public DfsHloRewriteVisitor { } TF_RET_CHECK(topk->operand_count() == 1); bool is_cuda = compute_capability_.IsCuda(); + // Route stable TopK to RAFT select_k via Uint64 adapter + if (is_cuda && ShouldRewriteStableTopKToUint64(topk)) { + ABSL_ASSIGN_OR_RETURN(HloInstruction * new_topk, + RewriteStableTopKToUint64(topk)); + return ReplaceInstruction(topk, new_topk); + } if (auto small_topk = SmallBufferOptimization(topk, is_cuda); small_topk.ok()) { diff --git a/third_party/xla/xla/backends/gpu/transforms/topk_specializer_test.cc b/third_party/xla/xla/backends/gpu/transforms/topk_specializer_test.cc index f5e82e4a9257af..a97ac72a482f34 100644 --- a/third_party/xla/xla/backends/gpu/transforms/topk_specializer_test.cc +++ b/third_party/xla/xla/backends/gpu/transforms/topk_specializer_test.cc @@ -37,6 +37,7 @@ limitations under the License. #include "xla/hlo/ir/hlo_instructions.h" #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/pass/hlo_pass_interface.h" +#include "xla/hlo/testlib/filecheck.h" #include "xla/service/topk_rewriter.h" #include "xla/shape_util.h" #include "xla/stream_executor/device_description.h" @@ -74,7 +75,7 @@ class TopkTest : public HloPjRtGpuTestBase, public ParameterizedInterface { %broadcast.40631 = pred[] broadcast(pred[] %constant.40630), dimensions={} %p.0.lhs.40626 = $3[] parameter(0) %p.0.rhs.40627 = $3[] parameter(1) - %compare.40632 = pred[] compare($3[] %p.0.lhs.40626, $3[] %p.0.rhs.40627), direction=GT, type=TOTALORDER + %compare.40632 = pred[] compare($3[] %p.0.lhs.40626, $3[] %p.0.rhs.40627), direction=GT, order=TOTAL ROOT %select.40633 = pred[] select(pred[] %broadcast.40631, pred[] %compare.40632, pred[] %broadcast.40631) } @@ -181,7 +182,7 @@ TEST_F(TopkTest, PreservesBackendConfig) { p.1.rhs = s32[] parameter(3) p.0.lhs = f32[] parameter(0) p.0.rhs = f32[] parameter(1) - ROOT compare = pred[] compare(p.0.lhs, p.0.rhs), direction=GT, type=TOTALORDER + ROOT compare = pred[] compare(p.0.lhs, p.0.rhs), direction=GT, order=TOTAL } ENTRY top_k { @@ -213,5 +214,157 @@ TEST_F(TopkTest, PreservesBackendConfig) { EXPECT_EQ(custom_call->raw_backend_config_string(), "{is_stable = false}"); } +TEST_F(TopkTest, RewriteStableTopKF32ToUint64) { + const char* hlo = R"( + HloModule m + + %compare-gt.1 { + p.1.lhs = s32[] parameter(2) + p.1.rhs = s32[] parameter(3) + p.0.lhs = f32[] parameter(0) + p.0.rhs = f32[] parameter(1) + ROOT compare = pred[] compare(p.0.lhs, p.0.rhs), direction=GT, order=TOTAL + } + + ENTRY top_k { + arg = f32[8,1024] parameter(0) + ROOT result = (f32[8,32], s32[8,32]) custom-call(arg), custom_call_target="TopK", called_computations={%compare-gt.1}, backend_config={is_stable = true} + } + )"; + + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(hlo)); + + if (!device_description().gpu_compute_capability().IsCuda()) { + GTEST_SKIP() << "RAFT is CUDA-only."; + } + + ASSERT_OK_AND_ASSIGN( + bool changed, + TopkSpecializer(device_description().gpu_compute_capability()) + .Run(module.get())); + ASSERT_TRUE(changed); + + const char* check_pattern = R"( +// CHECK-LABEL: ENTRY %top_k +// CHECK: %[[ARG:[^ ]+]] = f32[8,1024]{{.*}} parameter(0) + +// 1. Bitcast to U32 and S32 +// CHECK: %[[VAL_U32:[^ ]+]] = u32[8,1024]{{.*}} bitcast-convert(%[[ARG]]) +// CHECK: %[[VAL_S32:[^ ]+]] = s32[8,1024]{{.*}} bitcast-convert(%[[ARG]]) + +// 2. Pure Bitwise Radix Flip (F32 -> U32) +// CHECK: %[[SRA:[^ ]+]] = s32[8,1024]{{.*}} shift-right-arithmetic(%[[VAL_S32]], {{.*}}) +// CHECK: %[[SRA_U32:[^ ]+]] = u32[8,1024]{{.*}} bitcast-convert(%[[SRA]]) +// CHECK: %[[MASK:[^ ]+]] = u32[8,1024]{{.*}} or(%[[SRA_U32]], {{.*}}) +// CHECK: %[[RADIX_KEY:[^ ]+]] = u32[8,1024]{{.*}} xor(%[[VAL_U32]], %[[MASK]]) + +// 3. Pack into U64 +// CHECK: %[[KEY_U64:[^ ]+]] = u64[8,1024]{{.*}} convert(%[[RADIX_KEY]]) +// CHECK: %[[SHIFT_LEFT:[^ ]+]] = u64[8,1024]{{.*}} shift-left(%[[KEY_U64]], {{.*}}) +// CHECK: %[[IOTA:[^ ]+]] = u32[8,1024]{{.*}} iota(), iota_dimension=1 +// CHECK: %[[SUBTRACT:[^ ]+]] = u32[8,1024]{{.*}} subtract({{.*}}, %[[IOTA]]) +// CHECK: %[[PACKED:[^ ]+]] = u64[8,1024]{{.*}} or(%[[SHIFT_LEFT]], {{.*}}) + +// 4. CustomCall (__gpu$TopK) +// CHECK: %[[CUSTOM_CALL:[^ ]+]] = (u64[8,32]{{.*}}, s32[8,32]{{.*}}) custom-call(%[[PACKED]]), custom_call_target="__gpu$TopK", api_version=API_VERSION_TYPED_FFI, {{.*}} backend_config={is_stable = false} + +// 5. Unpack U64 -> U32 +// CHECK: %[[SRL:[^ ]+]] = u64[8,32]{{.*}} shift-right-logical(%[[CUSTOM_CALL]]#0, {{.*}}) +// CHECK: %[[UNPACK_U32:[^ ]+]] = u32[8,32]{{.*}} convert(%[[SRL]]) + +// 6. Pure Bitwise Radix Unflip (U32 -> F32) +// CHECK: %[[MSB:[^ ]+]] = u32[8,32]{{.*}} shift-right-logical(%[[UNPACK_U32]], {{.*}}) +// CHECK: %[[MSB_MINUS_ONE:[^ ]+]] = u32[8,32]{{.*}} subtract(%[[MSB]], {{.*}}) +// CHECK: %[[UNMASK:[^ ]+]] = u32[8,32]{{.*}} or(%[[MSB_MINUS_ONE]], {{.*}}) +// CHECK: %[[UNFLIPPED:[^ ]+]] = u32[8,32]{{.*}} xor(%[[UNPACK_U32]], %[[UNMASK]]) +// CHECK: %[[F32_OUT:[^ ]+]] = f32[8,32]{{.*}} bitcast-convert(%[[UNFLIPPED]]) + +// 7. Tuple output +// CHECK: ROOT {{.*}} tuple({{.*}}%[[F32_OUT]], {{.*}}%[[CUSTOM_CALL]]#1) + )"; + + ASSERT_OK_AND_ASSIGN(bool filecheck_matched, + RunFileCheck(module->ToString(), check_pattern)); + EXPECT_TRUE(filecheck_matched); +} + +TEST_F(TopkTest, RewriteStableTopKBF16ToUint64) { + const char* hlo = R"( + HloModule m + + %compare-gt.1 { + p.1.lhs = s32[] parameter(2) + p.1.rhs = s32[] parameter(3) + p.0.lhs = bf16[] parameter(0) + p.0.rhs = bf16[] parameter(1) + ROOT compare = pred[] compare(p.0.lhs, p.0.rhs), direction=GT, order=TOTAL + } + + ENTRY top_k { + arg = bf16[8,65540] parameter(0) + ROOT result = (bf16[8,32], s32[8,32]) custom-call(arg), custom_call_target="TopK", called_computations={%compare-gt.1}, backend_config={is_stable = true} + } + )"; + + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(hlo)); + + if (!device_description().gpu_compute_capability().IsCuda()) { + GTEST_SKIP() << "RAFT is CUDA-only."; + } + + ASSERT_OK_AND_ASSIGN( + bool changed, + TopkSpecializer(device_description().gpu_compute_capability()) + .Run(module.get())); + ASSERT_TRUE(changed); + + const char* check_pattern = R"( +// CHECK-LABEL: ENTRY %top_k +// CHECK: %[[ARG:[^ ]+]] = bf16[8,65540]{{.*}} parameter(0) + +// 1. Convert BF16 -> F32, then Bitcast to U32 and S32 +// CHECK: %[[F32_VAL:[^ ]+]] = f32[8,65540]{{.*}} convert(%[[ARG]]) +// CHECK: %[[VAL_U32:[^ ]+]] = u32[8,65540]{{.*}} bitcast-convert(%[[F32_VAL]]) +// CHECK: %[[VAL_S32:[^ ]+]] = s32[8,65540]{{.*}} bitcast-convert(%[[F32_VAL]]) + +// 2. Pure Bitwise Radix Flip (F32 -> U32) +// CHECK: %[[SRA:[^ ]+]] = s32[8,65540]{{.*}} shift-right-arithmetic(%[[VAL_S32]], {{.*}}) +// CHECK: %[[SRA_U32:[^ ]+]] = u32[8,65540]{{.*}} bitcast-convert(%[[SRA]]) +// CHECK: %[[MASK:[^ ]+]] = u32[8,65540]{{.*}} or(%[[SRA_U32]], {{.*}}) +// CHECK: %[[RADIX_KEY:[^ ]+]] = u32[8,65540]{{.*}} xor(%[[VAL_U32]], %[[MASK]]) + +// 3. Pack into U64 +// CHECK: %[[KEY_U64:[^ ]+]] = u64[8,65540]{{.*}} convert(%[[RADIX_KEY]]) +// CHECK: %[[SHIFT_LEFT:[^ ]+]] = u64[8,65540]{{.*}} shift-left(%[[KEY_U64]], {{.*}}) +// CHECK: %[[IOTA:[^ ]+]] = u32[8,65540]{{.*}} iota(), iota_dimension=1 +// CHECK: %[[SUBTRACT:[^ ]+]] = u32[8,65540]{{.*}} subtract({{.*}}, %[[IOTA]]) +// CHECK: %[[PACKED:[^ ]+]] = u64[8,65540]{{.*}} or(%[[SHIFT_LEFT]], {{.*}}) + +// 4. CustomCall (__gpu$TopK) +// CHECK: %[[CUSTOM_CALL:[^ ]+]] = (u64[8,32]{{.*}}, s32[8,32]{{.*}}) custom-call(%[[PACKED]]), custom_call_target="__gpu$TopK", api_version=API_VERSION_TYPED_FFI, {{.*}} backend_config={is_stable = false} + +// 5. Unpack U64 -> U32 +// CHECK: %[[SRL:[^ ]+]] = u64[8,32]{{.*}} shift-right-logical(%[[CUSTOM_CALL]]#0, {{.*}}) +// CHECK: %[[UNPACK_U32:[^ ]+]] = u32[8,32]{{.*}} convert(%[[SRL]]) + +// 6. Pure Bitwise Radix Unflip (U32 -> F32) +// CHECK: %[[MSB:[^ ]+]] = u32[8,32]{{.*}} shift-right-logical(%[[UNPACK_U32]], {{.*}}) +// CHECK: %[[MSB_MINUS_ONE:[^ ]+]] = u32[8,32]{{.*}} subtract(%[[MSB]], {{.*}}) +// CHECK: %[[UNMASK:[^ ]+]] = u32[8,32]{{.*}} or(%[[MSB_MINUS_ONE]], {{.*}}) +// CHECK: %[[UNFLIPPED:[^ ]+]] = u32[8,32]{{.*}} xor(%[[UNPACK_U32]], %[[UNMASK]]) +// CHECK: %[[F32_OUT:[^ ]+]] = f32[8,32]{{.*}} bitcast-convert(%[[UNFLIPPED]]) + +// 7. Convert F32 -> BF16 and Tuple output +// CHECK: %[[BF16_OUT:[^ ]+]] = bf16[8,32]{{.*}} convert(%[[F32_OUT]]) +// CHECK: ROOT {{.*}} tuple({{.*}}%[[BF16_OUT]], {{.*}}%[[CUSTOM_CALL]]#1) + )"; + + ASSERT_OK_AND_ASSIGN(bool filecheck_matched, + RunFileCheck(module->ToString(), check_pattern)); + EXPECT_TRUE(filecheck_matched); +} + } // namespace } // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/transforms/topk_splitter_test.cc b/third_party/xla/xla/backends/gpu/transforms/topk_splitter_test.cc index 6af575ed3bbaf0..858909020ca958 100644 --- a/third_party/xla/xla/backends/gpu/transforms/topk_splitter_test.cc +++ b/third_party/xla/xla/backends/gpu/transforms/topk_splitter_test.cc @@ -56,7 +56,7 @@ constexpr absl::string_view kComparator = R"( %broadcast.40631 = pred[] broadcast(pred[] %constant.40630), dimensions={} %p.0.lhs.40626 = f32[] parameter(0) %p.0.rhs.40627 = f32[] parameter(1) - %compare.40632 = pred[] compare(f32[] %p.0.lhs.40626, f32[] %p.0.rhs.40627), direction=GT, type=TOTALORDER + %compare.40632 = pred[] compare(f32[] %p.0.lhs.40626, f32[] %p.0.rhs.40627), direction=GT, order=TOTAL ROOT %select.40633 = pred[] select(pred[] %broadcast.40631, pred[] %compare.40632, pred[] %broadcast.40631) })"; diff --git a/third_party/xla/xla/backends/interpreter/compiler.cc b/third_party/xla/xla/backends/interpreter/compiler.cc index 8e8b94732e1762..66c5b3a9667e74 100644 --- a/third_party/xla/xla/backends/interpreter/compiler.cc +++ b/third_party/xla/xla/backends/interpreter/compiler.cc @@ -91,7 +91,7 @@ absl::StatusOr HandleEvaluatorCustomCall( absl::Status InterpreterCompiler::RunHloOptimization(HloModule* hlo_module) { HloPassPipeline pipeline("Interpreter"); - // The TopkDecomposer generates a compare op with type=TOTALORDER and must + // The TopkDecomposer generates a compare op with order=TOTAL and must // run before the ComparisonExpander which rewrites such comparisons. pipeline.AddPass(); pipeline.AddPass(); diff --git a/third_party/xla/xla/backends/profiler/gpu/BUILD b/third_party/xla/xla/backends/profiler/gpu/BUILD index 919e4b8faa26fd..5d9a42c42e8d38 100644 --- a/third_party/xla/xla/backends/profiler/gpu/BUILD +++ b/third_party/xla/xla/backends/profiler/gpu/BUILD @@ -301,6 +301,7 @@ xla_test( "@local_config_cuda//cuda:cuda_headers", "@local_config_cuda//cuda:cupti_headers", "@tsl//tsl/profiler/lib:scoped_annotation", + "@tsl//tsl/profiler/protobuf:xplane_proto_cc", ], ) diff --git a/third_party/xla/xla/backends/profiler/gpu/cupti_collector.cc b/third_party/xla/xla/backends/profiler/gpu/cupti_collector.cc index b7b3eff6c7b5a4..be8c5c423e54eb 100644 --- a/third_party/xla/xla/backends/profiler/gpu/cupti_collector.cc +++ b/third_party/xla/xla/backends/profiler/gpu/cupti_collector.cc @@ -21,7 +21,6 @@ limitations under the License. #include #include #include -#include #include #include #include diff --git a/third_party/xla/xla/backends/profiler/gpu/cupti_tracer_test.cc b/third_party/xla/xla/backends/profiler/gpu/cupti_tracer_test.cc index 8842a1d852e87e..8e49a762ce58c9 100644 --- a/third_party/xla/xla/backends/profiler/gpu/cupti_tracer_test.cc +++ b/third_party/xla/xla/backends/profiler/gpu/cupti_tracer_test.cc @@ -37,6 +37,7 @@ limitations under the License. #include "xla/tsl/profiler/utils/xplane_schema.h" #include "xla/tsl/profiler/utils/xplane_utils.h" #include "tsl/profiler/lib/scoped_annotation.h" +#include "tsl/profiler/protobuf/xplane.pb.h" namespace xla { namespace profiler { diff --git a/third_party/xla/xla/comparison_util.cc b/third_party/xla/xla/comparison_util.cc index ddf6179c928a9e..9d7ffcd7012ad7 100644 --- a/third_party/xla/xla/comparison_util.cc +++ b/third_party/xla/xla/comparison_util.cc @@ -45,18 +45,6 @@ PrimitiveType DefaultPrimitiveType(Comparison::Type type) { } } -// Returns the expected ordering for each primitive type. -Comparison::Order DefaultPrimitiveOrdering(PrimitiveType type) { - if (primitive_util::IsFloatingPointType(type) || - primitive_util::IsComplexType(type)) { - return Comparison::Order::kPartial; - } - if (primitive_util::IsIntegralType(type) || type == PRED) { - return Comparison::Order::kTotal; - } - LOG(FATAL) << "Unsupported type: " << PrimitiveType_Name(type); -} - // Returns the converse of `direction`. Comparison::Direction Converse(Comparison::Direction direction) { switch (direction) { @@ -223,6 +211,18 @@ Comparison::Order Comparison::DefaultOrdering(Comparison::Type type) { } } +// Returns the expected ordering for each primitive type. +Comparison::Order Comparison::DefaultOrdering(PrimitiveType type) { + if (primitive_util::IsFloatingPointType(type) || + primitive_util::IsComplexType(type)) { + return Comparison::Order::kPartial; + } + if (primitive_util::IsIntegralType(type) || type == PRED) { + return Comparison::Order::kTotal; + } + LOG(FATAL) << "Unsupported type: " << PrimitiveType_Name(type); +} + namespace { Comparison::Type ComparisonTypeFromPrimitiveTypeAndOrder( PrimitiveType type, Comparison::Order order) { @@ -251,7 +251,7 @@ Comparison::Comparison(Direction dir, PrimitiveType type, Order order) Comparison::Comparison(Direction dir, PrimitiveType type) : dir_(dir), primitive_type_(type), - order_(DefaultPrimitiveOrdering(type)), + order_(DefaultOrdering(type)), type_(DefaultComparisonType(type)) {} Comparison::Comparison(Direction dir, Type type) diff --git a/third_party/xla/xla/comparison_util.h b/third_party/xla/xla/comparison_util.h index 86df8e2a3c6fd3..133c6e0cdff06f 100644 --- a/third_party/xla/xla/comparison_util.h +++ b/third_party/xla/xla/comparison_util.h @@ -216,6 +216,9 @@ class Comparison { // Comparison::Type. static Comparison::Order DefaultOrdering(Type type); + // Returns the expected Comparison::Order for each primitive type. + static Comparison::Order DefaultOrdering(PrimitiveType type); + // Returns the Comparison::Type for the given primitive type. This assumes // that each numerical representation follows the standard behavior, e.g., // integers are total order and floats are partial order. diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc index 8d36c0bf83f141..8e350f9b3b5acf 100644 --- a/third_party/xla/xla/debug_options_flags.cc +++ b/third_party/xla/xla/debug_options_flags.cc @@ -486,6 +486,10 @@ DebugOptions DefaultDebugOptionsIgnoringFlags() { opts.set_xla_deduplicate_backend_configs_min_size( std::numeric_limits::max()); + opts.set_xla_force_config(""); + + opts.set_xla_candidate_configs_file(""); + opts.set_xla_gpu_experimental_autotune_cache_mode( DebugOptions::AUTOTUNE_CACHE_MODE_UPDATE); @@ -515,7 +519,7 @@ DebugOptions DefaultDebugOptionsIgnoringFlags() { opts.set_xla_gpu_experimental_matmul_perf_table_path(""); // TODO(b/366475196): Create XLA GPU without cuDNN, cuBLAS. opts.set_xla_gpu_experimental_disable_binary_libraries(false); - opts.set_xla_gpu_experimental_enable_conv_fusion(false); + opts.set_xla_gpu_experimental_enable_conv_fusion(true); opts.set_xla_gpu_dot_merger_threshold_mb(64); opts.set_xla_enable_fast_math(false); opts.set_xla_gpu_experimental_parallel_collective_overlap_limit(1); @@ -3060,6 +3064,18 @@ void MakeDebugOptionsFlags(std::vector* flag_list, "Minimum backend_config size (in bytes) to be eligible for deduplication " "into payloads during serialization. Configs smaller than this threshold " "are kept inline. Default is MAX_INT (feature disabled).")); + flag_list->push_back(tsl::Flag( + "xla_force_config", + string_setter_for(&DebugOptions::set_xla_force_config), + debug_options->xla_force_config(), + "Single serialized config to override config of all instructions, " + "bypassing cache and autotuning.")); + flag_list->push_back(tsl::Flag( + "xla_candidate_configs_file", + string_setter_for(&DebugOptions::set_xla_candidate_configs_file), + debug_options->xla_candidate_configs_file(), + "File containing a list of serialized configs to override supported " + "configs for all instructions.")); flag_list->push_back(tsl::Flag( "xla_gpu_experimental_autotune_backends", SetterForRepeatedEnum( diff --git a/third_party/xla/xla/hlo/analysis/hlo_alias_analysis.cc b/third_party/xla/xla/hlo/analysis/hlo_alias_analysis.cc index 856fe28197eb74..cb70853d9c4632 100644 --- a/third_party/xla/xla/hlo/analysis/hlo_alias_analysis.cc +++ b/third_party/xla/xla/hlo/analysis/hlo_alias_analysis.cc @@ -77,16 +77,10 @@ void ComputeInputOutputAliasedValues(const HloValue& value, std::optional aliased_input = io_alias_config.GetAliasedParameter(pos.index); if (aliased_input) { - // Pipelining and loop crossing can create multi-value value sets. - // Iterate over all aliased values instead of assuming a unique value. - for (const HloValue* aliased_val : - dataflow - .GetValueSet(entry_computation.parameter_instruction( - aliased_input->parameter_number), - aliased_input->parameter_index) - .values()) { - aliased_values.insert(aliased_val); - } + aliased_values.insert( + &dataflow.GetUniqueValueAt(entry_computation.parameter_instruction( + aliased_input->parameter_number), + aliased_input->parameter_index)); } } } @@ -99,15 +93,13 @@ void ComputeWhileAliasedValues(const HloValue& value, // Value is init of a while (use is while). for (const HloUse& use : value.GetUses()) { if (use.instruction->opcode() == HloOpcode::kWhile) { - // Determine all while values that this shares a buffer with. - // A while operand value set may contain multiple aliased values. - for (const HloValue* while_value : - dataflow.GetValueSet(use.instruction, use.operand_index).values()) { - aliased_values.insert(while_value); - VLOG(3) << " value is init value to a while; must share buffer with " - "while value " - << *while_value; - } + // Determine the while value that this shares a buffer with. + const HloValue& while_value = + dataflow.GetUniqueValueAt(use.instruction, use.operand_index); + aliased_values.insert(&while_value); + VLOG(3) << " value is init value to a while; must share buffer with " + "while value " + << while_value; } } // Value is a parameter of a while body/condition. @@ -120,15 +112,12 @@ void ComputeWhileAliasedValues(const HloValue& value, // Call graph must have been flattened. CHECK_EQ(call_graph_node.caller_callsites().size(), 1); - for (const HloValue* while_value : - dataflow - .GetValueSet(callsite.instruction(), value.defining_index()) - .values()) { - VLOG(3) << " value is parameter value of the body or condition of a " - "while; must share buffer with while value " - << *while_value; - aliased_values.insert(while_value); - } + const HloValue& while_value = dataflow.GetUniqueValueAt( + callsite.instruction(), value.defining_index()); + VLOG(3) << " value is parameter value of the body or condition of a " + "while; must share buffer with while value " + << while_value; + aliased_values.insert(&while_value); } } } @@ -149,16 +138,14 @@ void ComputeWhileAliasedValues(const HloValue& value, CHECK_EQ(call_graph_node.caller_callsites().size(), 1) << "Call graph must have been flattened."; - for (const HloValue* while_value : - dataflow.GetValueSet(callsite.instruction(), position.index) - .values()) { - VLOG(3) << " value @ " << position << " is root of " - << callsite.instruction()->name() - << "; body root and while value root must share buffer " - "among them: " - << *while_value; - aliased_values.insert(while_value); - } + const HloValue& while_value = + dataflow.GetUniqueValueAt(callsite.instruction(), position.index); + VLOG(3) << " value @ " << position << " is root of " + << callsite.instruction()->name() + << "; body root and while value root must share buffer " + "among them: " + << while_value; + aliased_values.insert(&while_value); } } } @@ -183,16 +170,13 @@ void ComputeConditionalAliasedValues(const HloValue& value, // Call graph must have been flattened. CHECK_EQ(call_graph_node.caller_callsites().size(), 1); - for (const HloValue* cond_value : - dataflow.GetValueSet(callsite.instruction(), position.index) - .values()) { - VLOG(3) - << " value @ " << position << " is root of " - << callsite.instruction()->name() - << "; branch computation roots must share buffer among them : " - << *cond_value; - aliased_values.insert(cond_value); - } + const HloValue& cond_value = + dataflow.GetUniqueValueAt(callsite.instruction(), position.index); + VLOG(3) << " value @ " << position << " is root of " + << callsite.instruction()->name() + << "; branch computation roots must share buffer among them : " + << cond_value; + aliased_values.insert(&cond_value); } } } @@ -210,15 +194,11 @@ void ComputeInPlaceOperationAliasedValues(const HloValue& value, alias_info->GetInPlaceInputOutputPairs(instruction)) { if (position.index == operand_and_output_index.second) { const HloOperandIndex& operand_index = operand_and_output_index.first; - for (const HloValue* operand_value : - dataflow - .GetValueSet( - instruction->operand(operand_index.operand_number), - operand_index.operand_index) - .values()) { - VLOG(3) << " operand value " << *operand_value << " aliases."; - aliased_values.insert(operand_value); - } + const HloValue& operand_value = dataflow.GetUniqueValueAt( + instruction->operand(operand_index.operand_number), + operand_index.operand_index); + VLOG(3) << " operand value " << operand_value << " aliases."; + aliased_values.insert(&operand_value); } } } @@ -229,13 +209,10 @@ void ComputeInPlaceOperationAliasedValues(const HloValue& value, const HloOperandIndex& operand_index = operand_and_output_index.first; if (use.operand_number == operand_index.operand_number && use.operand_index == operand_index.operand_index) { - for (const HloValue* use_value : - dataflow - .GetValueSet(use.instruction, operand_and_output_index.second) - .values()) { - VLOG(3) << " use value " << *use_value << " aliases."; - aliased_values.insert(use_value); - } + const HloValue& use_value = dataflow.GetUniqueValueAt( + use.instruction, operand_and_output_index.second); + VLOG(3) << " use value " << use_value << " aliases."; + aliased_values.insert(&use_value); } } } diff --git a/third_party/xla/xla/hlo/analysis/hlo_alias_analysis_test.cc b/third_party/xla/xla/hlo/analysis/hlo_alias_analysis_test.cc index 4934f47a7fc879..653ca27f75133e 100644 --- a/third_party/xla/xla/hlo/analysis/hlo_alias_analysis_test.cc +++ b/third_party/xla/xla/hlo/analysis/hlo_alias_analysis_test.cc @@ -648,50 +648,6 @@ ENTRY main { EXPECT_FALSE(AnyValuesInSameBufferInterfere()); } -TEST_F(HloAliasAnalysisTest, WhileOperandWithMultipleValues) { - const char* hlo_string = R"( -HloModule test - -body { - body_param = (f32[], f32[]) parameter(0) - body_element_0 = f32[] get-tuple-element(body_param), index=0 - body_element_1 = f32[] get-tuple-element(body_param), index=1 - add = f32[] add(body_element_0, body_element_1) - ROOT body_tuple = (f32[], f32[]) tuple(body_element_0, add) -} - -condition { - cond_param = (f32[], f32[]) parameter(0) - ROOT cond_constant = pred[] constant(false) -} - -ENTRY main { - c1 = f32[] constant(1.0) - c2 = f32[] constant(2.0) - tuple = (f32[], f32[]) tuple(c1, c2) - ROOT xla_while = (f32[], f32[]) while(tuple), condition=condition, body=body -} - )"; - ASSERT_OK_AND_ASSIGN(module_, ParseAndReturnVerifiedModule(hlo_string)); - - const HloAliasAnalysis& analysis = RunAnalysis(); - HloInstruction* c1 = FindInstruction(module_.get(), "c1"); - HloInstruction* xla_while = FindInstruction(module_.get(), "xla_while"); - HloInstruction* body_param = FindInstruction(module_.get(), "body_param"); - HloInstruction* cond_param = FindInstruction(module_.get(), "cond_param"); - EXPECT_THAT(analysis.GetUniqueBufferAt(xla_while, /*index=*/{0}).values(), - UnorderedElementsAre(&GetValueDefinedAt(c1))); - EXPECT_THAT( - analysis.GetUniqueBufferAt(xla_while, /*index=*/{0}).ComputePositions(), - UnorderedElementsAre( - HloPosition{c1, {}}, - HloPosition{FindInstruction(module_.get(), "tuple"), {0}}, - HloPosition{xla_while, {0}}, HloPosition{body_param, {0}}, - HloPosition{FindInstruction(module_.get(), "body_element_0"), {}}, - HloPosition{cond_param, {0}}, - HloPosition{FindInstruction(module_.get(), "body_tuple"), {0}})); - EXPECT_FALSE(AnyValuesInSameBufferInterfere()); -} TEST_F(HloAliasAnalysisTest, SequentialWhiles) { // Test sequential while instructions. The while body includes a diff --git a/third_party/xla/xla/hlo/builder/xla_builder_test.cc b/third_party/xla/xla/hlo/builder/xla_builder_test.cc index 0fffd4beb5047e..4aa44b97be8114 100644 --- a/third_party/xla/xla/hlo/builder/xla_builder_test.cc +++ b/third_party/xla/xla/hlo/builder/xla_builder_test.cc @@ -2209,14 +2209,15 @@ TEST(XlaBuilderTest, SetAndGetSharding) { hlo_sharding_2); } -TEST(XlaBuilderTest, ComparisonType) { +TEST(XlaBuilderTest, ComparisonOrder) { XlaBuilder b(TestName()); (void)Le(ConstantR0(&b, 1), ConstantR0(&b, 2)); TF_ASSERT_OK_AND_ASSIGN(const auto module, BuildHloModule(b)); const HloInstruction* root = GetRoot(*module); ASSERT_THAT(root, GmockMatch(m::Compare(m::Constant(), m::Constant()))); - EXPECT_EQ(Comparison::Type::kSigned, - DynCast(root)->type()); + const auto* compare = DynCast(root); + EXPECT_EQ(S32, compare->comparison().GetPrimitiveType()); + EXPECT_EQ(ComparisonOrder::kTotal, compare->order()); } TEST(XlaBuilderTest, StableLookUpInstructionByHandle) { diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator_test.cc b/third_party/xla/xla/hlo/evaluator/hlo_evaluator_test.cc index 9803edff14a199..2263849eae41f1 100644 --- a/third_party/xla/xla/hlo/evaluator/hlo_evaluator_test.cc +++ b/third_party/xla/xla/hlo/evaluator/hlo_evaluator_test.cc @@ -3122,6 +3122,231 @@ TEST_P(HloEvaluatorBf16Test, Conv2DGroupedConvolution) { EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); } +TEST_P(HloEvaluatorBf16Test, Conv2DMatrixMultiplyFastPath) { + HloComputation::Builder b(TestName()); + // LHS: [2, 3] + Array2D lhs_array({ + {1.f, 2.f, 3.f}, + {4.f, 5.f, 6.f}, + }); + auto lhs_literal = LiteralUtil::CreateR2FromArray2D(lhs_array); + HloInstruction* lhs_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(lhs_literal))); + + // RHS: [3, 2] + Array2D rhs_array({ + {7.f, 8.f}, + {9.f, 10.f}, + {11.f, 12.f}, + }); + auto rhs_literal = LiteralUtil::CreateR2FromArray2D(rhs_array); + HloInstruction* rhs_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(rhs_literal))); + + ConvolutionDimensionNumbers dnums; + dnums.set_input_batch_dimension(0); + dnums.set_input_feature_dimension(1); + dnums.set_kernel_input_feature_dimension(0); + dnums.set_kernel_output_feature_dimension(1); + dnums.set_output_batch_dimension(0); + dnums.set_output_feature_dimension(1); + + Window window; + + Shape shape = ShapeUtil::MakeShape(F32, {2, 2}); + b.AddInstruction(HloInstruction::CreateConvolve( + shape, {lhs_instruction, rhs_instruction}, /*feature_group_count=*/1, + /*batch_group_count=*/1, window, dnums, DefaultPrecisionConfig(2))); + m_->AddEntryComputation(b.Build()); + + TF_ASSERT_OK_AND_ASSIGN(Literal result, Evaluate()); + + Array2D expected_array({ + {58.f, 64.f}, + {139.f, 154.f}, + }); + auto expected = LiteralUtil::CreateR2FromArray2D(expected_array); + + EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); +} + +TEST_F(HloEvaluatorTest, Conv2DF8E4M3FNToBF16MatrixMultiplyFastPath) { + HloComputation::Builder b(TestName()); + // LHS: [2, 3] in F8E4M3FN + Array2D lhs_array({ + {float8_e4m3fn(1.f), float8_e4m3fn(2.f), float8_e4m3fn(3.f)}, + {float8_e4m3fn(4.f), float8_e4m3fn(5.f), float8_e4m3fn(6.f)}, + }); + auto lhs_literal = LiteralUtil::CreateR2FromArray2D(lhs_array); + HloInstruction* lhs_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(lhs_literal))); + + // RHS: [3, 2] in F8E4M3FN + Array2D rhs_array({ + {float8_e4m3fn(7.f), float8_e4m3fn(8.f)}, + {float8_e4m3fn(9.f), float8_e4m3fn(10.f)}, + {float8_e4m3fn(11.f), float8_e4m3fn(12.f)}, + }); + auto rhs_literal = LiteralUtil::CreateR2FromArray2D(rhs_array); + HloInstruction* rhs_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(rhs_literal))); + + ConvolutionDimensionNumbers dnums; + dnums.set_input_batch_dimension(0); + dnums.set_input_feature_dimension(1); + dnums.set_kernel_input_feature_dimension(0); + dnums.set_kernel_output_feature_dimension(1); + dnums.set_output_batch_dimension(0); + dnums.set_output_feature_dimension(1); + + Window window; + + Shape shape = ShapeUtil::MakeShape(BF16, {2, 2}); + b.AddInstruction(HloInstruction::CreateConvolve( + shape, {lhs_instruction, rhs_instruction}, /*feature_group_count=*/1, + /*batch_group_count=*/1, window, dnums, DefaultPrecisionConfig(2))); + m_->AddEntryComputation(b.Build()); + + TF_ASSERT_OK_AND_ASSIGN(Literal result, Evaluate()); + + Array2D expected_array({ + {bfloat16(58.f), bfloat16(64.f)}, + {bfloat16(139.f), bfloat16(154.f)}, + }); + auto expected = LiteralUtil::CreateR2FromArray2D(expected_array); + + EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); +} + +TEST_P(HloEvaluatorBf16Test, Conv2DMatrixMultiplyZeroContractingDim) { + HloComputation::Builder b(TestName()); + // LHS: [2, 0] + Array2D lhs_array(2, 0); + auto lhs_literal = LiteralUtil::CreateR2FromArray2D(lhs_array); + HloInstruction* lhs_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(lhs_literal))); + + // RHS: [0, 2] + Array2D rhs_array(0, 2); + auto rhs_literal = LiteralUtil::CreateR2FromArray2D(rhs_array); + HloInstruction* rhs_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(rhs_literal))); + + ConvolutionDimensionNumbers dnums; + dnums.set_input_batch_dimension(0); + dnums.set_input_feature_dimension(1); + dnums.set_kernel_input_feature_dimension(0); + dnums.set_kernel_output_feature_dimension(1); + dnums.set_output_batch_dimension(0); + dnums.set_output_feature_dimension(1); + + Window window; + + Shape shape = ShapeUtil::MakeShape(F32, {2, 2}); + b.AddInstruction(HloInstruction::CreateConvolve( + shape, {lhs_instruction, rhs_instruction}, /*feature_group_count=*/1, + /*batch_group_count=*/1, window, dnums, DefaultPrecisionConfig(2))); + m_->AddEntryComputation(b.Build()); + + TF_ASSERT_OK_AND_ASSIGN(Literal result, Evaluate()); + + Array2D expected_array({ + {0.f, 0.f}, + {0.f, 0.f}, + }); + auto expected = LiteralUtil::CreateR2FromArray2D(expected_array); + + EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); +} + +TEST_P(HloEvaluatorBf16Test, Conv2DMatrixMultiplyColumnMajorFallback) { + HloComputation::Builder b(TestName()); + // LHS: [2, 3] with column-major layout {0, 1} + Array2D lhs_array({ + {1.f, 2.f, 3.f}, + {4.f, 5.f, 6.f}, + }); + auto lhs_literal = + LiteralUtil::CreateR2FromArray2D(lhs_array).Relayout( + LayoutUtil::MakeLayout({0, 1})); + HloInstruction* lhs_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(lhs_literal))); + + // RHS: [3, 2] + Array2D rhs_array({ + {7.f, 8.f}, + {9.f, 10.f}, + {11.f, 12.f}, + }); + auto rhs_literal = LiteralUtil::CreateR2FromArray2D(rhs_array); + HloInstruction* rhs_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(rhs_literal))); + + ConvolutionDimensionNumbers dnums; + dnums.set_input_batch_dimension(0); + dnums.set_input_feature_dimension(1); + dnums.set_kernel_input_feature_dimension(0); + dnums.set_kernel_output_feature_dimension(1); + dnums.set_output_batch_dimension(0); + dnums.set_output_feature_dimension(1); + + Window window; + + Shape shape = ShapeUtil::MakeShape(F32, {2, 2}); + b.AddInstruction(HloInstruction::CreateConvolve( + shape, {lhs_instruction, rhs_instruction}, /*feature_group_count=*/1, + /*batch_group_count=*/1, window, dnums, DefaultPrecisionConfig(2))); + m_->AddEntryComputation(b.Build()); + + TF_ASSERT_OK_AND_ASSIGN(Literal result, Evaluate()); + + Array2D expected_array({ + {58.f, 64.f}, + {139.f, 154.f}, + }); + auto expected = LiteralUtil::CreateR2FromArray2D(expected_array); + + EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); +} + +TEST_P(HloEvaluatorBf16Test, Conv2DMatrixMultiplyZeroBatchDim) { + HloComputation::Builder b(TestName()); + // LHS: [0, 3] + Array2D lhs_array(0, 3); + auto lhs_literal = LiteralUtil::CreateR2FromArray2D(lhs_array); + HloInstruction* lhs_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(lhs_literal))); + + // RHS: [3, 2] + Array2D rhs_array(3, 2); + auto rhs_literal = LiteralUtil::CreateR2FromArray2D(rhs_array); + HloInstruction* rhs_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(rhs_literal))); + + ConvolutionDimensionNumbers dnums; + dnums.set_input_batch_dimension(0); + dnums.set_input_feature_dimension(1); + dnums.set_kernel_input_feature_dimension(0); + dnums.set_kernel_output_feature_dimension(1); + dnums.set_output_batch_dimension(0); + dnums.set_output_feature_dimension(1); + + Window window; + + Shape shape = ShapeUtil::MakeShape(F32, {0, 2}); + b.AddInstruction(HloInstruction::CreateConvolve( + shape, {lhs_instruction, rhs_instruction}, /*feature_group_count=*/1, + /*batch_group_count=*/1, window, dnums, DefaultPrecisionConfig(2))); + m_->AddEntryComputation(b.Build()); + + TF_ASSERT_OK_AND_ASSIGN(Literal result, Evaluate()); + + Array2D expected_array(0, 2); + auto expected = LiteralUtil::CreateR2FromArray2D(expected_array); + + EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); +} + // Initialization of data sets for FFT tests: void HloEvaluatorTest::InitializeFftData() { diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator_typed_visitor.h b/third_party/xla/xla/hlo/evaluator/hlo_evaluator_typed_visitor.h index ea59214d5e0288..4047aad199a1a2 100644 --- a/third_party/xla/xla/hlo/evaluator/hlo_evaluator_typed_visitor.h +++ b/third_party/xla/xla/hlo/evaluator/hlo_evaluator_typed_visitor.h @@ -922,6 +922,50 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault { const int64_t feature_group_count = conv->feature_group_count(); const int64_t batch_group_count = conv->batch_group_count(); + if constexpr (std::is_same_v) { + auto is_row_major_r2 = [](const Shape& s) { + return s.dimensions_size() == 2 && + (!s.has_layout() || + LayoutUtil::IsMonotonicWithDim0Major(s.layout())); + }; + + if (parent_->trace_mac_handler_ == nullptr && feature_group_count == 1 && + batch_group_count == 1 && num_spatial_dims == 0 && + dnums.input_batch_dimension() == 0 && + dnums.input_feature_dimension() == 1 && + dnums.kernel_input_feature_dimension() == 0 && + dnums.kernel_output_feature_dimension() == 1 && + dnums.output_batch_dimension() == 0 && + dnums.output_feature_dimension() == 1 && is_row_major_r2(lhs_shape) && + is_row_major_r2(rhs_shape) && is_row_major_r2(result_shape)) { + const int64_t m = lhs_shape.dimensions(0); + const int64_t k = lhs_shape.dimensions(1); + const int64_t n = rhs_shape.dimensions(1); + + if (m > 0 && k > 0 && n > 0) { + Literal lhs_f32 = lhs_literal.Convert(F32).value(); + Literal rhs_f32 = rhs_literal.Convert(F32).value(); + + Array2D lhs_array(m, k); + lhs_array.SetValues(lhs_f32.data()); + Array2D rhs_array(k, n); + rhs_array.SetValues(rhs_f32.data()); + + std::unique_ptr> result_array = + HloEvaluator::MatmulArray2D(lhs_array, rhs_array); + + Literal result_f32(ShapeUtil::MakeShape(F32, {m, n})); + result_f32.PopulateR2FromArray2D(*result_array); + + parent_->SetEvaluatedLiteralFor( + conv, std::move(result_f32) + .Convert(result_shape.element_type()) + .value()); + return absl::OkStatus(); + } + } + } + auto func = [&window_shape, &dnums, &lhs_shape, &rhs_shape, &window, &lhs_dim_multipliers, &rhs_dim_multipliers, lhs_literal_data, rhs_literal_data, feature_group_count, batch_group_count, diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction_test.cc b/third_party/xla/xla/hlo/ir/hlo_instruction_test.cc index 44691380b98efb..4da2aca17d6c97 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instruction_test.cc +++ b/third_party/xla/xla/hlo/ir/hlo_instruction_test.cc @@ -622,23 +622,22 @@ TEST_F(HloInstructionTest, PrintCompareOpWorksIfDead) { ENTRY main { p0 = f32[] parameter(0) p1 = f32[] parameter(1) - ROOT result = pred[] compare(p0, p1), direction=GT, type=TOTALORDER + ROOT result = pred[] compare(p0, p1), direction=GT, order=TOTAL } )"; ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kModuleStr)); HloInstruction* root = module->entry_computation()->root_instruction(); - EXPECT_EQ( - root->ToString(), - "%result = pred[] compare(%p0, %p1), direction=GT, type=TOTALORDER"); + EXPECT_EQ(root->ToString(), + "%result = pred[] compare(%p0, %p1), direction=GT, order=TOTAL"); module->entry_computation()->set_root_instruction( root->mutable_operand(0), /*accept_different_shape=*/true); root->DetachFromOperandsAndUsers(); EXPECT_EQ( root->ToString(), - "%result = pred[] compare(null , null ), direction=GT, type=TOTALORDER"); + "%result = pred[] compare(null , null ), direction=GT, order=TOTAL"); TF_ASSERT_OK(module->entry_computation()->RemoveInstruction(root)); EXPECT_EQ(root->ToString(), - "%result = pred[] compare(), direction=GT, type=TOTALORDER"); + "%result = pred[] compare(), direction=GT, order=TOTAL"); *module->mutable_entry_computation_layout() = module->compute_computation_layout(); } @@ -650,7 +649,7 @@ TEST_F(HloInstructionTest, CanonicalPrintingSupportsInt64) { ENTRY main { p0 = f32[] parameter(0) p1 = f32[] parameter(1) - ROOT result = pred[] compare(p0, p1), direction=GT, type=TOTALORDER + ROOT result = pred[] compare(p0, p1), direction=GT, order=TOTAL } )")); @@ -683,7 +682,7 @@ TEST_F(HloInstructionTest, CanonicalPrintingSupportsInt64) { EXPECT_EQ(param2_to_string, "tmp_1 = f32[] parameter(1)"); EXPECT_EQ(param3_to_string, "tmp_2 = pred[] compare(f32[] tmp_0, f32[] tmp_1), direction=GT, " - "type=TOTALORDER"); + "order=TOTAL"); } TEST_F(HloInstructionTest, CanonicalPrintingSupportsCustomCall) { diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction_utils_test.cc b/third_party/xla/xla/hlo/ir/hlo_instruction_utils_test.cc index ace67d7f44f153..b7e30fdcc5e405 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instruction_utils_test.cc +++ b/third_party/xla/xla/hlo/ir/hlo_instruction_utils_test.cc @@ -671,7 +671,7 @@ TEST_F(HloInstructionUtilsTest, IsTopKStable) { p.1.rhs = s32[] parameter(3) p.0.lhs = f32[] parameter(0) p.0.rhs = f32[] parameter(1) - ROOT compare = pred[] compare(p.0.lhs, p.0.rhs), direction=GT, type=TOTALORDER + ROOT compare = pred[] compare(p.0.lhs, p.0.rhs), direction=GT, order=TOTAL } ENTRY main { diff --git a/third_party/xla/xla/hlo/ir/hlo_instructions.cc b/third_party/xla/xla/hlo/ir/hlo_instructions.cc index 4a04449f0504f5..9f5c773561d98f 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instructions.cc +++ b/third_party/xla/xla/hlo/ir/hlo_instructions.cc @@ -972,10 +972,10 @@ void HloCompareInstruction::PrintExtraAttributesImpl( // We might want to print a HloInstruction which has been cleand up and has no // operands anymore. This should not result in a crash. if (operand_count() == 0 || operand(0) == nullptr || - compare_.GetType() != Comparison::DefaultComparisonType( - operand(0)->shape().element_type())) { + order() != + Comparison::DefaultOrdering(operand(0)->shape().element_type())) { printer.Next([this](Printer* printer) { - AppendCat(printer, "type=", ComparisonTypeToString(compare_.GetType())); + AppendCat(printer, "order=", ComparisonOrderToShortString(order())); }); } } diff --git a/third_party/xla/xla/hlo/ir/hlo_schedule.cc b/third_party/xla/xla/hlo/ir/hlo_schedule.cc index 5ffb9647bbc524..653e81fdb922ef 100644 --- a/third_party/xla/xla/hlo/ir/hlo_schedule.cc +++ b/third_party/xla/xla/hlo/ir/hlo_schedule.cc @@ -154,134 +154,160 @@ const HloInstructionSequence& HloSchedule::sequence( return sequences_.at(computation->unique_id()); } -absl::Status HloSchedule::UpdateComputationSchedule( - const HloComputation* computation) { - // Map from unique ID to HloInstruction pointer for instructions in the - // computation. - absl::flat_hash_map id_to_instruction; - for (HloInstruction* instruction : computation->instructions()) { - InsertOrDie(&id_to_instruction, instruction->unique_id(), instruction); - } - - // Invalidate the schedule of instructions that conflict with control - // dependencies. - auto sched_sequence = sequence(computation).instructions(); - absl::flat_hash_set invalid_instructions; - absl::flat_hash_set seen_instructions; - for (HloInstruction* inst : sched_sequence) { - for (HloInstruction* pred : inst->control_predecessors()) { - // Found a pair of instructions whose schedule order is inconsistent with - // their control dependencies. - if (pred == inst) { - ABSL_RETURN_IF_ERROR(pred->RemoveControlDependencyTo(inst)); - } - if (pred->parent() == computation && !seen_instructions.contains(pred)) { - invalid_instructions.insert(inst); - invalid_instructions.insert(pred); - } - } - seen_instructions.insert(inst); - } - for (HloInstruction* inst : invalid_instructions) { - sequences_.at(computation->unique_id()).remove_instruction(inst); - } +namespace { - // Set of all HloInstructions in the schedule. - absl::flat_hash_set ids_in_schedule; - for (int64_t id : sequence(computation).ids()) { - InsertOrDie(&ids_in_schedule, id); - } +// Tracks in-degrees and successors for dependency-driven schedule +// reconstruction. +struct DependencyTracker { + // Number of unscheduled operands and control predecessors for each + // instruction in the computation. + absl::flat_hash_map in_degree; - // Map from HloInstruction X to newly added instructions (instruction is in - // computation, but not in schedule) which depend on X. If an instruction is - // not in the map, then it has no users or control successors which are newly - // added instructions. + // Map from HloInstruction X to instructions in the computation which depend + // on X. absl::flat_hash_map> - new_instruction_successors; - - // For each newly added instruction, this is the count of the instruction's - // operands and control predecessors that have not yet been scheduled. When - // this value reaches zero, then the instruction may be placed in the - // schedule. - absl::flat_hash_map unscheduled_predecessor_count; - - // Create a worklist of newly added instructions which are ready to be added - // to the schedule. Initialize worklist with those that have zero operands. - std::queue worklist; + successors; + + // Worklist of instructions whose dependencies are fully satisfied and are + // ready to be placed in the schedule. + std::queue ready_worklist; +}; + +// Builds in-degree counts and successor lists for all instructions in the +// computation, and seeds the ready worklist with newly added 0-predecessor +// instructions. +absl::StatusOr BuildDependencyTracker( + const HloComputation* computation, + const absl::flat_hash_set& ids_in_schedule) { + DependencyTracker tracker; + absl::flat_hash_set predecessors; for (HloInstruction* instruction : computation->instructions()) { - if (!ids_in_schedule.contains(instruction->unique_id())) { - // `instruction` is a newly added instruction which is not in the - // schedule. - if (instruction->operands().empty() && - instruction->control_predecessors().empty()) { - // `instruction` has no operands or control dependencies. It may be - // added to the schedule immediately (once the worklist is processed). - worklist.push(instruction); - } else { - absl::flat_hash_set predecessors; - auto add_predecessor = [&](const HloInstruction* predecessor) { - std::vector& successors = - new_instruction_successors[predecessor]; - if (!absl::c_linear_search(successors, instruction)) { - // Only add an instruction once. - successors.push_back(instruction); - } - predecessors.insert(predecessor); - }; - for (const HloInstruction* operand : instruction->operands()) { - add_predecessor(operand); - } - for (const HloInstruction* control_predecessor : - instruction->control_predecessors()) { - add_predecessor(control_predecessor); - } - unscheduled_predecessor_count[instruction] = predecessors.size(); + predecessors.clear(); + for (const HloInstruction* operand : instruction->operands()) { + if (predecessors.insert(operand).second) { + tracker.successors[operand].push_back(instruction); } } + for (HloInstruction* control_pred : instruction->control_predecessors()) { + if (control_pred == instruction) { + ABSL_RETURN_IF_ERROR(control_pred->RemoveControlDependencyTo(instruction)); + } else if (predecessors.insert(control_pred).second) { + tracker.successors[control_pred].push_back(instruction); + } + } + + tracker.in_degree[instruction] = predecessors.size(); + + // Newly added instructions with zero predecessors are ready immediately. + if (predecessors.empty() && + !ids_in_schedule.contains(instruction->unique_id())) { + tracker.ready_worklist.push(instruction); + } } + return tracker; +} - // Update the schedule with the newly added instructions, and remove any - // instructions no longer in the graph. +// Reconstructs the schedule by replaying existing_sequence and interleaving +// ready instructions as their dependencies are satisfied. +HloInstructionSequence BuildUpdatedSequence( + const HloInstructionSequence& existing_sequence, + const absl::flat_hash_map& id_to_instruction, + const absl::flat_hash_set& ids_in_schedule, + DependencyTracker& tracker) { HloInstructionSequence new_sequence; + absl::flat_hash_set scheduled; + absl::flat_hash_set deferred_instructions; - // Lambda which schedules all instructions on the worklist. - auto schedule_worklist = [&]() { - while (!worklist.empty()) { - HloInstruction* instruction = worklist.front(); - worklist.pop(); - new_sequence.push_back(instruction); - std::vector* new_successors = - tsl::gtl::FindOrNull(new_instruction_successors, instruction); - if (new_successors != nullptr) { - // This just-scheduled instruction has users which are newly added to - // the module. Update the number of unscheduled operands and push the - // newly added instruction to the worklist if it is ready to - // schedule. - for (HloInstruction* new_successor : *new_successors) { - unscheduled_predecessor_count.at(new_successor)--; - CHECK_GE(unscheduled_predecessor_count.at(new_successor), 0); - if (unscheduled_predecessor_count.at(new_successor) == 0) { - worklist.push(new_successor); + auto schedule_instruction = [&](HloInstruction* instruction) { + if (!scheduled.insert(instruction).second) { + return; + } + new_sequence.push_back(instruction); + + auto it = tracker.successors.find(instruction); + if (it != tracker.successors.end()) { + for (HloInstruction* successor : it->second) { + int& in_degree = tracker.in_degree.at(successor); + in_degree--; + CHECK_GE(in_degree, 0); + if (in_degree == 0) { + // Push newly added instructions and previously deferred instructions + // to the ready worklist. Instructions still awaiting their turn in + // existing_sequence are not pushed early. + if (!ids_in_schedule.contains(successor->unique_id()) || + deferred_instructions.contains(successor)) { + tracker.ready_worklist.push(successor); } } } } }; - schedule_worklist(); - for (int64_t id : sequences_.at(computation->unique_id()).ids()) { + auto drain_ready_worklist = [&]() { + while (!tracker.ready_worklist.empty()) { + HloInstruction* ready_instruction = tracker.ready_worklist.front(); + tracker.ready_worklist.pop(); + schedule_instruction(ready_instruction); + } + }; + + // 1. Schedule any 0-predecessor newly added instructions first. + drain_ready_worklist(); + + // 2. Iterate through the existing sequence in order. + for (int64_t id : existing_sequence.ids()) { auto it = id_to_instruction.find(id); if (it == id_to_instruction.end()) { - // This instruction in the schedule is no longer in the module. Do not add - // it to the new schedule. + // Instruction was removed from computation; skip. + continue; + } + HloInstruction* instruction = it->second; + if (scheduled.contains(instruction)) { + // Already scheduled via the ready worklist. continue; } - worklist.push(it->second); - schedule_worklist(); + if (tracker.in_degree.at(instruction) == 0) { + schedule_instruction(instruction); + drain_ready_worklist(); + } else { + // Instruction has unsatisfied dependencies and cannot be scheduled yet; + // defer it until its in-degree reaches 0. + deferred_instructions.insert(instruction); + } } - set_sequence(computation, std::move(new_sequence)); + // 3. Drain any remaining ready instructions. + drain_ready_worklist(); + + return new_sequence; +} + +} // namespace + +absl::Status HloSchedule::UpdateComputationSchedule( + const HloComputation* computation) { + // Map from unique ID to HloInstruction pointer for instructions in the + // computation. + absl::flat_hash_map id_to_instruction; + for (HloInstruction* instruction : computation->instructions()) { + InsertOrDie(&id_to_instruction, instruction->unique_id(), instruction); + } + + // Set of all HloInstructions in the schedule. + HloInstructionSequence& current_sequence = + sequences_.at(computation->unique_id()); + absl::flat_hash_set ids_in_schedule(current_sequence.ids().begin(), + current_sequence.ids().end()); + + // Build in-degree dependency tracking for all instructions. + ABSL_ASSIGN_OR_RETURN(DependencyTracker tracker, + BuildDependencyTracker(computation, ids_in_schedule)); + + // Reconstruct schedule in topological order. + set_sequence(computation, + BuildUpdatedSequence(current_sequence, id_to_instruction, + ids_in_schedule, tracker)); return absl::OkStatus(); } diff --git a/third_party/xla/xla/hlo/parser/hlo_parser.cc b/third_party/xla/xla/hlo/parser/hlo_parser.cc index 21abddde1b2211..3e4398bdea7818 100644 --- a/third_party/xla/xla/hlo/parser/hlo_parser.cc +++ b/third_party/xla/xla/hlo/parser/hlo_parser.cc @@ -9213,7 +9213,7 @@ void HloParserImpl::UpdateAsyncWrappedComputation( if (i < async_wrapped_computation->num_parameters()) { Shape* param_shape = async_wrapped_computation->parameter_instruction(i)->mutable_shape(); - if (!ShapeUtil::Compatible( + if (!ShapeUtil::Equal( *param_shape, called_computation->parameter_instruction(i)->shape())) { *param_shape = called_computation->parameter_instruction(i)->shape(); @@ -9230,7 +9230,7 @@ void HloParserImpl::UpdateAsyncWrappedComputation( Shape* root_shape = async_wrapped_computation->root_instruction()->mutable_shape(); const Shape& result_shape = called_computation->root_instruction()->shape(); - if (!ShapeUtil::Compatible(*root_shape, result_shape)) { + if (!ShapeUtil::Equal(*root_shape, result_shape)) { *root_shape = result_shape; } } diff --git a/third_party/xla/xla/hlo/parser/hlo_parser_test.cc b/third_party/xla/xla/hlo/parser/hlo_parser_test.cc index 0de052fc0a1420..7b734e8f010136 100644 --- a/third_party/xla/xla/hlo/parser/hlo_parser_test.cc +++ b/third_party/xla/xla/hlo/parser/hlo_parser_test.cc @@ -311,7 +311,7 @@ R"(HloModule SelectR1F32WithCmpR1F32sFromParamsSmall_module, entry_computation_l ENTRY %SelectR1F32WithCmpR1F32sFromParamsSmall.v4 (v1: f32[4], v2: f32[4]) -> f32[4] { %v1 = f32[4]{0} parameter(0), sharding={maximal device=1} %v2 = f32[4]{0} parameter(1), sharding={maximal device=1} - %greater-than = pred[4]{0} compare(f32[4]{0} %v1, f32[4]{0} %v2), direction=GT, type=TOTALORDER, sharding={replicated} + %greater-than = pred[4]{0} compare(f32[4]{0} %v1, f32[4]{0} %v2), direction=GT, order=TOTAL, sharding={replicated} ROOT %select = f32[4]{0} select(pred[4]{0} %greater-than, f32[4]{0} %v1, f32[4]{0} %v2), sharding={replicated} } @@ -798,7 +798,7 @@ R"(HloModule R4F32OverlapSmall_module, entry_computation_layout={()->f32[4,5,1,1 %ge_F32.v3 (lhs: f32[], rhs: f32[]) -> pred[] { %lhs = f32[] parameter(0) %rhs = f32[] parameter(1) - ROOT %greater-than-or-equal-to = pred[] compare(f32[] %lhs, f32[] %rhs), direction=GE, type=TOTALORDER + ROOT %greater-than-or-equal-to = pred[] compare(f32[] %lhs, f32[] %rhs), direction=GE, order=TOTAL } %add_F32.v3 (lhs.1: f32[], rhs.1: f32[]) -> f32[] { @@ -7928,6 +7928,64 @@ ENTRY main { "f32[64]"); } +TEST_F(HloParserTest, + DesugarParsingTest_CallStart_LayoutSyncFromCalledComputation) { + const char* const hlo = R"( +HloModule main + +comp { + ROOT root = f32[16,8]{1,0} parameter(0) +} + +ENTRY main { + arg.0 = f32[16,8]{0,1} parameter(0) + call-start = ((f32[16,8]{0,1}), f32[16,8]{0,1}, s32[]) call-start(arg.0), async_execution_thread="thread", to_apply=comp + call-update = ((f32[16,8]{0,1}), f32[16,8]{0,1}, s32[]) call-update(call-start) + ROOT call-done = f32[16,8]{0,1} call-done(call-update) +} +)"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnUnverifiedModule(hlo)); + HloInstruction* async_done = module->entry_computation()->root_instruction(); + HloComputation* async_wrapped = async_done->async_wrapped_computation(); + ASSERT_NE(async_wrapped, nullptr); + // Parameters and root of the async-wrapped computation must synchronize their + // layouts with the called computation `comp` ({1,0}), even if call-start + // initially specified a different layout ({0,1}). + EXPECT_EQ(async_wrapped->parameter_instruction(0)->shape().ToString( + /*print_layout=*/true), + "f32[16,8]{1,0}"); + EXPECT_EQ(async_wrapped->root_instruction()->shape().ToString( + /*print_layout=*/true), + "f32[16,8]{1,0}"); +} + +TEST_F(HloParserTest, + DesugarParsingTest_FusionStart_LayoutSyncFromFusedComputation) { + const char* const hlo = R"( +HloModule main + +ENTRY main { + arg.0 = f32[16,8]{0,1} parameter(0) + fusion-start = ((f32[16,8]{0,1}), f32[16,8]{0,1}, s32[]) fusion-start(arg.0), kind=kLoop, calls={ + p0 = f32[16,8]{1,0} parameter(0) + ROOT root = f32[16,8]{1,0} negate(p0) + } + fusion-update = ((f32[16,8]{0,1}), f32[16,8]{0,1}, s32[]) fusion-update(fusion-start) + ROOT fusion-done = f32[16,8]{0,1} fusion-done(fusion-update) +} +)"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnUnverifiedModule(hlo)); + HloInstruction* async_done = module->entry_computation()->root_instruction(); + HloComputation* async_wrapped = async_done->async_wrapped_computation(); + ASSERT_NE(async_wrapped, nullptr); + EXPECT_EQ(async_wrapped->parameter_instruction(0)->shape().ToString( + /*print_layout=*/true), + "f32[16,8]{1,0}"); + EXPECT_EQ(async_wrapped->root_instruction()->shape().ToString( + /*print_layout=*/true), + "f32[16,8]{1,0}"); +} + TEST_F(HloParserTest, DeeplyNestedOperandsExceedsRecursionLimit) { constexpr int kTestRecursionDepth = 10; std::string deeply_nested = "f32[] "; diff --git a/third_party/xla/xla/hlo/transforms/expanders/permutation_sort_expander_test.cc b/third_party/xla/xla/hlo/transforms/expanders/permutation_sort_expander_test.cc index 92b7df572c6d71..1f77dba53a63a4 100644 --- a/third_party/xla/xla/hlo/transforms/expanders/permutation_sort_expander_test.cc +++ b/third_party/xla/xla/hlo/transforms/expanders/permutation_sort_expander_test.cc @@ -40,7 +40,7 @@ TEST_F(PermutationSortExpanderTest, ReplacePermutationSortWithScatter) { p.0.rhs = f32[] parameter(1) p.1.lhs = s32[] parameter(2) p.1.rhs = s32[] parameter(3) - ROOT lt = pred[] compare(p.0.lhs, p.0.rhs), direction=LT, type=TOTALORDER + ROOT lt = pred[] compare(p.0.lhs, p.0.rhs), direction=LT, order=TOTAL } lt_s32 { @@ -83,7 +83,7 @@ TEST_F(PermutationSortExpanderTest, DontReplaceIfWrongComparisonDirection) { p.0.rhs = f32[] parameter(1) p.1.lhs = s32[] parameter(2) p.1.rhs = s32[] parameter(3) - ROOT lt = pred[] compare(p.0.lhs, p.0.rhs), direction=LT, type=TOTALORDER + ROOT lt = pred[] compare(p.0.lhs, p.0.rhs), direction=LT, order=TOTAL } lt_s32 { @@ -117,7 +117,7 @@ TEST_F(PermutationSortExpanderTest, DontReplaceIfComparingWrongParameters) { p.0.rhs = f32[] parameter(1) p.1.lhs = s32[] parameter(2) p.1.rhs = s32[] parameter(3) - ROOT lt = pred[] compare(p.0.lhs, p.0.rhs), direction=LT, type=TOTALORDER + ROOT lt = pred[] compare(p.0.lhs, p.0.rhs), direction=LT, order=TOTAL } lt_s32 { @@ -153,7 +153,7 @@ TEST_F(PermutationSortExpanderTest, DontReplacePermutationSortIfNonIntegral) { p.0.rhs = f32[] parameter(1) p.1.lhs = f32[] parameter(2) p.1.rhs = f32[] parameter(3) - ROOT lt = pred[] compare(p.0.lhs, p.0.rhs), direction=LT, type=TOTALORDER + ROOT lt = pred[] compare(p.0.lhs, p.0.rhs), direction=LT, order=TOTAL } ENTRY sort_computation { @@ -181,7 +181,7 @@ TEST_F(PermutationSortExpanderTest, DontReplacePermutationSortWrongDimensions) { p.0.rhs = f32[] parameter(1) p.1.lhs = s32[] parameter(2) p.1.rhs = s32[] parameter(3) - ROOT lt = pred[] compare(p.0.lhs, p.0.rhs), direction=LT, type=TOTALORDER + ROOT lt = pred[] compare(p.0.lhs, p.0.rhs), direction=LT, order=TOTAL } lt_s32 { @@ -216,7 +216,7 @@ TEST_F(PermutationSortExpanderTest, k_rhs = f32[] parameter(1) v_lhs = s64[] parameter(2) v_rhs = s64[] parameter(3) - ROOT lt = pred[] compare(k_lhs, k_rhs), direction=LT, type=TOTALORDER + ROOT lt = pred[] compare(k_lhs, k_rhs), direction=LT, order=TOTAL } lt_s64 { diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/BUILD b/third_party/xla/xla/hlo/transforms/simplifiers/BUILD index feff54ac3b9fad..617859e472cf2e 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/BUILD +++ b/third_party/xla/xla/hlo/transforms/simplifiers/BUILD @@ -1231,12 +1231,14 @@ xla_cc_test( "//xla:xla_data_proto_cc", "//xla/hlo/ir:hlo", "//xla/hlo/parser:hlo_parser", + "//xla/hlo/testlib:filecheck", "//xla/hlo/testlib:hlo_hardware_independent_test_base", "//xla/hlo/testlib:pattern_matcher_gmock", "//xla/hlo/testlib:test", "//xla/hlo/utils:hlo_matchers", "//xla/service:pattern_matcher", "//xla/tsl/platform:statusor", + "@com_google_absl//absl/status:status_matchers", "@com_google_absl//absl/strings:string_view", "@com_google_absl//absl/types:span", "@com_google_googletest//:gtest_main", diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc index 2a3ef2c1cb5aef..61d2a12d53d450 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc @@ -5701,8 +5701,7 @@ absl::Status AlgebraicSimplifierVisitor::HandleCompare( HloInstruction* rhs; CHECK(Match(compare, m::Compare(m::Op(&lhs), m::Op(&rhs)))); - if (Cast(compare)->type() == - Comparison::Type::kUnsigned) { + if (primitive_util::IsUnsignedIntegralType(lhs->shape().element_type())) { // X u< 0 -> false if (compare->comparison_direction() == ComparisonDirection::kLt && IsAll(rhs, 0)) { diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.cc b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.cc index 546a90bfd2e598..b1dc2bc4498dbd 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.cc @@ -470,8 +470,22 @@ absl::StatusOr HloConstantFolding::RunImpl( // Visit computations in reverse post-order, so that we can propagate constant // arguments from callers to callees. + absl::flat_hash_set live_computations; for (auto it = computations.rbegin(); it != computations.rend(); ++it) { HloComputation* computation = *it; + // If all callers were folded, skip instruction folding entirely. + if (!computation->IsEntryComputation()) { + bool has_live_caller = + absl::c_any_of(computation->caller_instructions(), + [&](const HloInstruction* caller) { + return live_computations.contains(caller->parent()); + }); + if (!has_live_caller) { + continue; + } + } + live_computations.insert(computation); + // If the computation is only used by call instructions, check whether for // any of the parameters of the computation, the argument passed by the // call-sites is always the same constant. In that case, we can sink the diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding_test.cc index 12f90085cdc2c1..bcde8b57c71a98 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding_test.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding_test.cc @@ -21,12 +21,14 @@ limitations under the License. #include #include +#include "absl/status/status_matchers.h" #include "absl/strings/string_view.h" #include "absl/types/span.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/hlo/parser/hlo_parser.h" +#include "xla/hlo/testlib/filecheck.h" #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" #include "xla/hlo/testlib/pattern_matcher_gmock.h" #include "xla/hlo/testlib/test.h" @@ -46,6 +48,7 @@ limitations under the License. namespace xla { namespace { +using ::absl_testing::IsOkAndHolds; namespace op = xla::testing::opcode_matchers; namespace m = xla::match; using HloConstantFoldingTest = HloHardwareIndependentTestBase; @@ -1400,5 +1403,33 @@ TEST_F(HloConstantFoldingTest, LateOptionsDontFoldGteWithControlDependency) { EXPECT_FALSE(result); } +TEST_F(HloConstantFoldingTest, SkipComputationsWithFoldedCallers) { + absl::string_view hlo_string = R"( + HloModule test + + // CHECK-LABEL: %foo + // CHECK: %p0 = s32[] parameter(0) + // CHECK: %c0 = s32[] constant(10) + // CHECK: %c1 = s32[] constant(20) + // CHECK: ROOT %add = s32[] add(%c0, %c1) + foo { + p0 = s32[] parameter(0) + c0 = s32[] constant(10) + c1 = s32[] constant(20) + ROOT add = s32[] add(c0, c1) + } + + // CHECK-LABEL: ENTRY %entry + // CHECK: ROOT %constant{{.*}} = s32[] constant(30) + ENTRY entry { + c = s32[] constant(5) + ROOT call = s32[] call(c), to_apply=foo + })"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + HloConstantFolding constant_folding; + EXPECT_THAT(constant_folding.Run(module.get()), IsOkAndHolds(true)); + EXPECT_THAT(RunFileCheck(module->ToString(), hlo_string), IsOkAndHolds(true)); +} + } // namespace } // namespace xla diff --git a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/tests/export.mlir b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/tests/export.mlir index fde846e97737e8..00c9ac46865642 100644 --- a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/tests/export.mlir +++ b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/tests/export.mlir @@ -2410,7 +2410,7 @@ func.func @main(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<10x12x12x64xf32>) } // CHECK: %[[SELECT_COMPUTATION:.*]] ([[ARG0:.*]]: f32[], [[ARG1:.*]]: f32[]) -> pred[] { -// CHECK: ROOT %[[RESULT:.*]] = pred[] compare(%[[ARG0]], %[[ARG1]]), direction=GE, type=TOTALORDER +// CHECK: ROOT %[[RESULT:.*]] = pred[] compare(%[[ARG0]], %[[ARG1]]), direction=GE, order=TOTAL // CHECK: %[[SCATTER_COMPUTATION:.*]] ([[ARG0:.*]]: f32[], [[ARG1:.*]]: f32[]) -> f32[] { // CHECK: ROOT %[[RESULT:.*]] = f32[] add(%[[ARG0]], %[[ARG1]]) diff --git a/third_party/xla/xla/hlo/translate/tests/stablehlo.mlir b/third_party/xla/xla/hlo/translate/tests/stablehlo.mlir index 05136c2152959e..392483610a2342 100644 --- a/third_party/xla/xla/hlo/translate/tests/stablehlo.mlir +++ b/third_party/xla/xla/hlo/translate/tests/stablehlo.mlir @@ -1897,7 +1897,7 @@ module { // CHECK: %[[$region_0_4:[^ ]+]] // CHECK-NEXT: %[[Arg_0_5:[^ ]+]] = f32[] parameter(0) // CHECK-NEXT: %[[Arg_1_6:[^ ]+]] = f32[] parameter(1) -// CHECK-NEXT: ROOT %[[compare_7:[^ ]+]] = pred[] compare(%[[Arg_0_5]], %[[Arg_1_6]]), direction=GE, type=TOTALORDER, metadata= +// CHECK-NEXT: ROOT %[[compare_7:[^ ]+]] = pred[] compare(%[[Arg_0_5]], %[[Arg_1_6]]), direction=GE, order=TOTAL, metadata= // CHECK: %[[$region_1_8:[^ ]+]] // CHECK-NEXT: %[[Arg_0_9:[^ ]+]] = f32[] parameter(0) diff --git a/third_party/xla/xla/hlo/utils/hlo_live_range.cc b/third_party/xla/xla/hlo/utils/hlo_live_range.cc index 1db60423a53b3e..f3027bb4cff609 100644 --- a/third_party/xla/xla/hlo/utils/hlo_live_range.cc +++ b/third_party/xla/xla/hlo/utils/hlo_live_range.cc @@ -537,4 +537,72 @@ std::string HloLiveRange::ToString() const { return output; } +int64_t ViewExtendedTransitiveUseTime( + const HloInstruction* view, int64_t view_color, + const absl::flat_hash_map& + instruction_schedule) { + CHECK(!view->shape().IsTuple() && view->shape().has_layout() && + view->shape().layout().memory_space() == view_color) + << "not a view: " << view->ToString(); + auto is_view_colored = [view_color](const HloInstruction* instruction) { + return instruction->shape().has_layout() && + instruction->shape().layout().memory_space() == view_color; + }; + int64_t use_time = -1; + absl::flat_hash_set visited = {view}; + std::vector worklist = {view}; + while (!worklist.empty()) { + const HloInstruction* current = worklist.back(); + worklist.pop_back(); + auto time_it = instruction_schedule.find(current); + if (time_it != instruction_schedule.end()) { + use_time = std::max(use_time, time_it->second); + } + for (const HloInstruction* user : current->users()) { + if (is_view_colored(user)) { + if (visited.insert(user).second) { + worklist.push_back(user); + } + } else { + auto user_time_it = instruction_schedule.find(user); + if (user_time_it != instruction_schedule.end()) { + use_time = std::max(use_time, user_time_it->second); + } + } + } + } + return use_time; +} + +void ExtendViewBaseLiveRanges(HloLiveRange* hlo_live_range, + const HloDataflowAnalysis& dataflow_analysis, + int64_t view_color) { + const absl::flat_hash_map& + instruction_schedule = hlo_live_range->instruction_schedule(); + absl::flat_hash_map& + buffer_live_ranges = hlo_live_range->buffer_live_ranges(); + // dataflow_analysis.values() is id ordered, so the walk is deterministic. + for (const HloValue* value : dataflow_analysis.values()) { + auto live_range_it = buffer_live_ranges.find(value); + if (live_range_it == buffer_live_ranges.end()) { + continue; + } + HloLiveRange::LiveRangeBounds& live_range = live_range_it->second; + for (const HloUse& use : value->GetUses()) { + const HloInstruction* user = use.instruction; + // Only the viewed value itself (operand 0 of the view) needs the + // extension; a view's start index operands are consumed at the view's + // own time. + if (use.operand_number != 0 || user->shape().IsTuple() || + !user->shape().has_layout() || + user->shape().layout().memory_space() != view_color) { + continue; + } + live_range.end = + std::max(live_range.end, ViewExtendedTransitiveUseTime( + user, view_color, instruction_schedule)); + } + } +} + } // namespace xla diff --git a/third_party/xla/xla/hlo/utils/hlo_live_range.h b/third_party/xla/xla/hlo/utils/hlo_live_range.h index 92758777416b47..be33e1a9bca6da 100644 --- a/third_party/xla/xla/hlo/utils/hlo_live_range.h +++ b/third_party/xla/xla/hlo/utils/hlo_live_range.h @@ -278,6 +278,32 @@ class HloLiveRange { absl::flat_hash_set execution_threads_; }; +// Returns the latest schedule time at which `view` (a value colored +// `view_color`, e.g. memory_space_assignment::Options::dus_view_color or +// BufferAssigner::Options::dus_view_color) still has its underlying storage +// read through it: the max schedule time over the transitive closure of the +// view's readers, following users that are themselves view colored. A view +// is an address into another buffer with no storage of its own, so that +// buffer must stay reserved until this time. +// +// REQUIRES: view->shape().IsTuple() == false. +int64_t ViewExtendedTransitiveUseTime( + const HloInstruction* view, int64_t view_color, + const absl::flat_hash_map& + instruction_schedule); + +// Extends, in `hlo_live_range`, the live range end of every value used as +// the base (operand 0) of a view to the view's last transitive reader (see +// ViewExtendedTransitiveUseTime). The view is an address into the base's +// buffer with no storage of its own, so every consumer of liveness that can +// recycle or overlap storage (allocation reuse, heap simulation) must see +// the base held live until its last reader through the view. Values are +// visited in `dataflow_analysis.values()` order (the analysis `hlo_live_range` +// was built from). +void ExtendViewBaseLiveRanges(HloLiveRange* hlo_live_range, + const HloDataflowAnalysis& dataflow_analysis, + int64_t view_color); + } // namespace xla #endif // XLA_HLO_UTILS_HLO_LIVE_RANGE_H_ diff --git a/third_party/xla/xla/hlo/utils/sort_utils_test.cc b/third_party/xla/xla/hlo/utils/sort_utils_test.cc index 08d36b9a5666f0..195ab46ca299d0 100644 --- a/third_party/xla/xla/hlo/utils/sort_utils_test.cc +++ b/third_party/xla/xla/hlo/utils/sort_utils_test.cc @@ -127,7 +127,7 @@ compare { rhs_is_zero = pred[] compare(rhs, c_zero), direction=EQ rhs_no_neg_zero = f32[] select(rhs_is_zero, c_zero, rhs) rhs_canonical = f32[] select(rhs_is_nan, c_nan, rhs_no_neg_zero) - ROOT cmp = pred[] compare(lhs_canonical, rhs_canonical), direction=LT, type=TOTALORDER + ROOT cmp = pred[] compare(lhs_canonical, rhs_canonical), direction=LT, order=TOTAL } ENTRY main { @@ -201,7 +201,7 @@ HloModule test_module compare { p0 = f32[] parameter(0) p1 = f32[] parameter(1) - ROOT cmp = pred[] compare(p0, p1), direction=LT, type=TOTALORDER + ROOT cmp = pred[] compare(p0, p1), direction=LT, order=TOTAL } ENTRY main { @@ -269,7 +269,7 @@ compare { lhs_no_neg_zero = f32[] select(lhs_is_zero, c_zero, lhs) rhs_is_zero = pred[] compare(rhs, c_zero), direction=EQ rhs_no_neg_zero = f32[] select(rhs_is_zero, c_zero, rhs) - ROOT cmp = pred[] compare(lhs_no_neg_zero, rhs_no_neg_zero), direction=LT, type=TOTALORDER + ROOT cmp = pred[] compare(lhs_no_neg_zero, rhs_no_neg_zero), direction=LT, order=TOTAL } ENTRY main { @@ -307,7 +307,7 @@ compare { rhs_no_neg_zero = f32[] select(rhs_is_zero, c_zero, rhs) rhs_canonical = f32[] select(rhs_is_nan, c_not_nan, rhs_no_neg_zero) - ROOT cmp = pred[] compare(lhs_canonical, rhs_canonical), direction=LT, type=TOTALORDER + ROOT cmp = pred[] compare(lhs_canonical, rhs_canonical), direction=LT, order=TOTAL } ENTRY main { @@ -412,7 +412,7 @@ compare { p1_zero_sel = bf16[] select(p1_eq_zero, c_zero, p1) p1_canon = bf16[] select(p1_ne, c_nan, p1_zero_sel) - ROOT cmp = pred[] compare(p0_canon, p1_canon), direction=LT, type=TOTALORDER + ROOT cmp = pred[] compare(p0_canon, p1_canon), direction=LT, order=TOTAL } ENTRY main { diff --git a/third_party/xla/xla/literal.h b/third_party/xla/xla/literal.h index 5a965de7394736..9324772cce53d4 100644 --- a/third_party/xla/xla/literal.h +++ b/third_party/xla/xla/literal.h @@ -597,14 +597,14 @@ class LiteralBase { void WriteElement(NativeT element) { constexpr PrimitiveType primitive_type = primitive_util::NativeToPrimitiveType(); - static_assert(primitive_util::StorageBitWidth(primitive_type) % 8 == 0); + static_assert(primitive_util::BitWidth(primitive_type) % 8 == 0); if constexpr (primitive_util::IsComplexType(primitive_type)) { WriteElement(element.real()); WriteElement(element.imag()); } else { constexpr PrimitiveType unsigned_type = primitive_util::UnsignedIntegralTypeForBitWidth( - primitive_util::StorageBitWidth(primitive_type)); + primitive_util::BitWidth(primitive_type)); using UnsignedT = primitive_util::NativeTypeOf; UnsignedT unsigned_element = absl::bit_cast(element); if constexpr (sizeof(UnsignedT) == 1) { @@ -625,7 +625,7 @@ class LiteralBase { constexpr PrimitiveType primitive_type = primitive_util::NativeToPrimitiveType(); constexpr int bits_per_element = primitive_util::BitWidth(primitive_type); - if constexpr (primitive_util::IsSubByteNonPredType(primitive_type)) { + if constexpr (bits_per_element < 8) { static_assert(!primitive_util::IsComplexType(primitive_type)); static_assert(8 % bits_per_element == 0); @@ -688,7 +688,7 @@ class LiteralBase { ABSL_MUST_USE_RESULT bool ReadElement(NativeT& element) { constexpr PrimitiveType primitive_type = primitive_util::NativeToPrimitiveType(); - static_assert(primitive_util::StorageBitWidth(primitive_type) % 8 == 0); + static_assert(primitive_util::BitWidth(primitive_type) % 8 == 0); if constexpr (primitive_util::IsComplexType(primitive_type)) { using ComponentT = primitive_util::NativeTypeOf; if constexpr (sizeof(UnsignedT) == 1) { if (at_end()) { @@ -736,7 +736,7 @@ class LiteralBase { constexpr PrimitiveType primitive_type = primitive_util::NativeToPrimitiveType(); constexpr int bits_per_element = primitive_util::BitWidth(primitive_type); - if constexpr (primitive_util::IsSubByteNonPredType(primitive_type)) { + if constexpr (bits_per_element < 8) { static_assert(!primitive_util::IsComplexType(primitive_type)); static_assert(8 % bits_per_element == 0); @@ -1758,7 +1758,7 @@ bool LiteralBase::Piece::DeserializeData( // - If a piece is dynamic, we first write the sizes of the dynamic dimensions. // // - The elements of the piece are then written. Elements smaller than a single -// byte (e.g. S4, U4) are packed into bytes. Otherwise, they are written in +// byte (PRED, S4, U4) are packed into bytes. Otherwise, they are written in // little-endian byte order. template absl::Status LiteralBase::SerializeWithShapeProto(const ShapeProto& shape_proto, diff --git a/third_party/xla/xla/mlir_hlo/transforms/passes.h b/third_party/xla/xla/mlir_hlo/transforms/passes.h index ce5a0740baf80b..fb1e1cb5c0f422 100644 --- a/third_party/xla/xla/mlir_hlo/transforms/passes.h +++ b/third_party/xla/xla/mlir_hlo/transforms/passes.h @@ -51,18 +51,6 @@ std::unique_ptr> createFinalBufferizePass( uint64_t alignment, BufferizeDialectsCallback dc = {}, BufferizePatternsCallback pc = {}); -// Creates a TileLoopsPass with tiles sizes provided through `tile_sizes` -// and unroll factors provided through `unroll_factors`. -inline std::unique_ptr createTileLoopsPass( - ArrayRef tileSizes, ArrayRef unrollFactors) { - TileLoopsPassOptions options; - options.tile_sizes_ = - SmallVector(tileSizes.begin(), tileSizes.end()); - options.unroll_factors_ = - SmallVector(unrollFactors.begin(), unrollFactors.end()); - return createTileLoopsPass(options); -} - namespace hlo { #define GEN_PASS_REGISTRATION #include "transforms/passes.h.inc" diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc index 9c4a1ad856af1a..2132b408ac400e 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc @@ -85,6 +85,9 @@ static FailureOr convertFloatValue( } std::optional getRefCoreType(TypedValue value) { + if (!value) { + return std::nullopt; + } auto space = dyn_cast_if_present( value.getType().getMemorySpace()); if (!space) { @@ -1649,6 +1652,19 @@ mlir::tpu::CoreType SemaphoreSignalOp::getTargetCoreType() { return getRefCoreType(getSemaphore()).value_or(GetCoreTypeOfParentOp(**this)); } +namespace { + +bool isRemote(Value device_id, Value core_id) { + return device_id != nullptr || core_id != nullptr; +} + +template +bool isRemote(OpTy op) { + return isRemote(op.getDeviceId(), op.getCoreId()); +} + +} // namespace + LogicalResult SemaphoreSignalOp::verify() { MemRefType sem_type = getSemaphore().getType(); if (sem_type.getRank() != 0) { @@ -1658,7 +1674,7 @@ LogicalResult SemaphoreSignalOp::verify() { CoreType issuing_core_type = GetCoreTypeOfParentOp(**this); CoreType target_core_type = getTargetCoreType(); - if (getCoreId() == nullptr && getDeviceId() == nullptr) { + if (!isRemote(*this)) { if (target_core_type != issuing_core_type) { return emitOpError( absl::StrFormat("Target core type (%s) must match source core type " @@ -1685,113 +1701,202 @@ LogicalResult SemaphoreWaitOp::verify() { return success(); } -void EnqueueDMAOp::build(OpBuilder& builder, OperationState& state, - Value source, Value source_semaphore, Value target, - Value target_semaphore, Value device_id, Value core_id, - uint32_t priority, bool strict_ordering) { - build(builder, state, source, source_semaphore, target, target_semaphore, - device_id, core_id, /*subcore_id=*/nullptr, priority, strict_ordering); -} -mlir::tpu::CoreType EnqueueDMAOp::getTargetCoreType() { - return getRefCoreType(getTargetSemaphore()) - .value_or(GetCoreTypeOfParentOp(**this)); +namespace { + +bool isSparseCoreStreamLocalMemory(tpu::MemorySpaceAttr mem_space, + CoreType issuing_core) { + if (mem_space.getCoreType().value_or(issuing_core) != issuing_core) { + return false; + } + if (issuing_core == CoreType::kScVectorSubcore) { + return mem_space.getValue() == MemorySpace::kVmem; + } + if (issuing_core == CoreType::kScScalarSubcore) { + return mem_space.getValue() == MemorySpace::kSmem; + } + return false; } -LogicalResult EnqueueDMAOp::verify() { - Value target_sem = getTargetSemaphore(); - if (!target_sem) { - // TODO: b/501204503 - Support optional source and destination semaphores. - return emitOpError("EnqueueDMA target semaphore must be provided."); +LogicalResult verifySparseCoreDmaSemaphores( + Operation* op, tpu::MemorySpaceAttr source_mem_space, + tpu::MemorySpaceAttr target_mem_space, Value source_semaphore, + Value target_semaphore, bool is_remote, CoreType issuing_core) { + if (is_remote) { + return success(); + } + bool src_is_local = + isSparseCoreStreamLocalMemory(source_mem_space, issuing_core); + bool tgt_is_local = + isSparseCoreStreamLocalMemory(target_mem_space, issuing_core); + if (!src_is_local && tgt_is_local) { + if (source_semaphore != nullptr) { + return op->emitOpError( + "Source semaphores are unsupported for transfers from a " + "non-local memory to local memory on SparseCore"); + } } - MemRefType target_sem_type = getMemRefType(target_sem); - if (target_sem_type.getRank() != 0) { - return emitOpError("DMA target semaphore must be rank 0"); + return success(); +} + +LogicalResult verifyCommonDmaParams( + Operation* op, MemRefType source_ty, MemRefType target_ty, + TypedValue source_sem, TypedValue target_sem, + Value device_id, Value core_id, Value subcore_id, int priority, + bool strict_ordering, CoreType issuing_core) { + Type sem_elem_type = nullptr; + + if (target_sem) { + MemRefType target_sem_type = getMemRefType(target_sem); + if (target_sem_type.getRank() != 0) { + return op->emitOpError("DMA target semaphore must be rank 0"); + } + sem_elem_type = target_sem_type.getElementType(); } - Type target_sem_elem_type = target_sem_type.getElementType(); - Value source_sem = getSourceSemaphore(); if (source_sem) { MemRefType source_sem_type = getMemRefType(source_sem); if (source_sem_type.getRank() != 0) { - return emitOpError("DMA source semaphore reference must be rank 0"); + return op->emitOpError("DMA source semaphore reference must be rank 0"); } - if (source_sem_type.getElementType() != target_sem_elem_type) { - return emitOpError( + if (sem_elem_type && source_sem_type.getElementType() != sem_elem_type) { + return op->emitOpError( "DMA source and target semaphore must have the same type"); } + if (!sem_elem_type) { + sem_elem_type = source_sem_type.getElementType(); + } } - MemRefType source_ty = getMemRefType(getSource()); - MemRefType target_ty = getMemRefType(getTarget()); + if (source_ty.getElementType() != target_ty.getElementType()) { - return emitOpError("DMA source and target element type mismatch"); + return op->emitOpError("DMA source and target element type mismatch"); } if (source_ty.getShape() != target_ty.getShape()) { - return emitOpError("DMA source and target shape mismatch."); + return op->emitOpError("DMA source and target shape mismatch."); } - if (getDeviceId() || getCoreId()) { - if (!getSourceSemaphore()) { - return emitOpError( - "DMA source semaphore must be specified when device_id or core_id is " - "specified"); + bool is_remote = isRemote(device_id, core_id); + bool is_sc = issuing_core == CoreType::kScVectorSubcore || + issuing_core == CoreType::kScScalarSubcore; + if (is_sc) { + if (source_ty.getMemorySpace() == nullptr) { + return op->emitOpError( + "Source memory space must be provided for SparseCore DMA"); } - } - bool is_remote = getDeviceId() || getCoreId(); - if (getSourceSemaphore()) { - // TODO: b/501204503 - Support optional source and destination semaphores. - if (!is_remote) { - return emitOpError( - "DMA destination device_id or core_id must be specified when source " - "semaphore is specified"); + if (target_ty.getMemorySpace() == nullptr) { + return op->emitOpError( + "Target memory space must be provided for SparseCore DMA"); } } - int priority = getPriority(); + if (priority < 0 || priority > 1) { - return emitOpError( + return op->emitOpError( "Not implemented: only support priority 0 or 1, but got ") << priority; } if (priority != 0 && is_remote) { - return emitOpError( + return op->emitOpError( "Not implemented: non-zero priority is not supported for remote DMA"); } + if (source_sem && + getRefCoreType(source_sem).value_or(issuing_core) != issuing_core) { + return op->emitOpError( + "Source semaphore and source ref core type mismatched"); + } // If the target core_type is different from the issuing core_type, // the specific core_id must be provided. The device_id is irrelevant here. - CoreType issuing_core = GetCoreTypeOfParentOp(**this); - CoreType target_core = getTargetCoreType(); - if (getSourceSemaphore() && - getRefCoreType(getSourceSemaphore()).value_or(issuing_core) != - issuing_core) { - return emitOpError("Source semaphore and source ref core type mismatched"); - } - if (target_core != issuing_core && getCoreId() == nullptr) { - return emitOpError( + CoreType target_core = getRefCoreType(target_sem).value_or(issuing_core); + if (target_core != issuing_core && core_id == nullptr) { + return op->emitOpError( absl::StrFormat("Core id must be specified when target core type (%v) " "is different from source core type (%v)", target_core, issuing_core)); } - if (getStrictOrdering() && issuing_core != CoreType::kScScalarSubcore && + if (strict_ordering && issuing_core != CoreType::kScScalarSubcore && issuing_core != CoreType::kScVectorSubcore) { - return emitOpError( + return op->emitOpError( "Strict ordering is only supported on the SC scalar and vector " "subcores"); } - if (isa(target_sem_elem_type)) { + if (sem_elem_type && isa(sem_elem_type)) { if (HasMemorySpace(source_ty, MemorySpace::kSmem, CoreType::kTc) || HasMemorySpace(target_ty, MemorySpace::kSmem, CoreType::kTc)) { - return emitOpError( + return op->emitOpError( "Non-DMA semaphores are not supported for DMAs involving SMEM"); } } // Subcore ID applies only to SC vector subcore DMAs. - if (target_core != CoreType::kScVectorSubcore && getSubcoreId() != nullptr) { - return emitOpError( + if (target_core != CoreType::kScVectorSubcore && subcore_id != nullptr) { + return op->emitOpError( "Subcore id should not be set unless target core type is SC vector " "subcore"); } return success(); } +} // namespace + +void EnqueueDMAOp::build(OpBuilder& builder, OperationState& state, + Value source, Value source_semaphore, Value target, + Value target_semaphore, Value device_id, Value core_id, + uint32_t priority, bool strict_ordering) { + build(builder, state, source, source_semaphore, target, target_semaphore, + device_id, core_id, /*subcore_id=*/nullptr, priority, strict_ordering); +} +mlir::tpu::CoreType EnqueueDMAOp::getTargetCoreType() { + return getRefCoreType(getTargetSemaphore()) + .value_or(GetCoreTypeOfParentOp(**this)); +} + +LogicalResult EnqueueDMAOp::verify() { + MemRefType source_ty = getMemRefType(getSource()); + MemRefType target_ty = getMemRefType(getTarget()); + CoreType issuing_core = GetCoreTypeOfParentOp(**this); + bool is_remote = isRemote(*this); + bool is_sc = issuing_core == CoreType::kScVectorSubcore || + issuing_core == CoreType::kScScalarSubcore; + + if (failed(verifyCommonDmaParams( + getOperation(), source_ty, target_ty, getSourceSemaphore(), + getTargetSemaphore(), getDeviceId(), getCoreId(), getSubcoreId(), + getPriority(), getStrictOrdering(), issuing_core))) { + return failure(); + } + + if (is_remote) { + if (!getSourceSemaphore()) { + return emitOpError( + "DMA source semaphore must be specified when device_id or core_id is " + "specified"); + } + } + + if (is_sc) { + auto source_mem_space = + dyn_cast_or_null(source_ty.getMemorySpace()); + auto target_mem_space = + dyn_cast_or_null(target_ty.getMemorySpace()); + if (source_mem_space != nullptr && target_mem_space != nullptr) { + if (failed(verifySparseCoreDmaSemaphores( + getOperation(), source_mem_space, target_mem_space, + getSourceSemaphore(), getTargetSemaphore(), is_remote, + issuing_core))) { + return failure(); + } + } + } else { + if (getTargetSemaphore() == nullptr) { + return emitOpError("DMA target semaphore must be specified"); + } + if (!is_remote && getSourceSemaphore() != nullptr) { + return emitOpError( + "DMA destination device_id or core_id must be specified when source " + "semaphore is specified"); + } + } + + return success(); +} + FailureOr EnqueueIndirectDMAOp::isGather() { const auto source_ms = dyn_cast_if_present( getMemRefType(getSource()).getMemorySpace()); @@ -1872,29 +1977,63 @@ LogicalResult WaitDMA2Op::verify() { } LogicalResult WaitDMAOp::verify() { + MemRefType source_ty = getMemRefType(getSource()); + MemRefType target_ty = getMemRefType(getTarget()); + CoreType issuing_core = GetCoreTypeOfParentOp(**this); + bool is_remote = isRemote(*this); + bool is_sc = issuing_core == CoreType::kScVectorSubcore || + issuing_core == CoreType::kScScalarSubcore; + + if (failed(verifyCommonDmaParams( + getOperation(), source_ty, target_ty, getSourceSemaphore(), + getTargetSemaphore(), getDeviceId(), getCoreId(), getSubcoreId(), + getPriority(), getStrictOrdering(), issuing_core))) { + return failure(); + } + bool wait_destination = getWaitTarget(); TypedValue sem = wait_destination ? getTargetSemaphore() : getSourceSemaphore(); - if (!sem) { - // TODO: b/501204503 - Support optional source and destination semaphores - // with global reserved semaphore allocation tracking. - return emitOpError("The awaited semaphore must be provided"); - } - if (getMemRefType(sem).getRank() != 0) { - return emitOpError("DMA wait semaphore must be rank 0"); - } - - CoreType issuing_core = GetCoreTypeOfParentOp(**this); - if (getRefCoreType(sem).value_or(issuing_core) != issuing_core) { - return emitOpError("Can only await semaphores attached to the local core"); + if (is_sc) { + if (!is_remote) { + auto source_mem_space = + dyn_cast_or_null(source_ty.getMemorySpace()); + auto target_mem_space = + dyn_cast_or_null(target_ty.getMemorySpace()); + if (source_mem_space != nullptr && target_mem_space != nullptr) { + bool src_is_local = + isSparseCoreStreamLocalMemory(source_mem_space, issuing_core); + bool tgt_is_local = + isSparseCoreStreamLocalMemory(target_mem_space, issuing_core); + if (!src_is_local && tgt_is_local && !wait_destination) { + return emitOpError( + "Awaiting source read completion is unsupported for transfers " + "from a non-local memory to local memory on SparseCore"); + } + } + } + } else { + if (wait_destination) { + if (getTargetSemaphore() == nullptr) { + return emitOpError("DMA target semaphore must be specified"); + } + if (!is_remote && getSourceSemaphore() != nullptr) { + return emitOpError( + "DMA destination device_id or core_id must be specified when " + "source semaphore is specified"); + } + } else { + if (getSourceSemaphore() == nullptr) { + return emitOpError("The awaited semaphore must be provided"); + } + } } - - // Subcore ID applies only to SC vector subcore. - if (issuing_core != CoreType::kScVectorSubcore && getSubcoreId() != nullptr) { - return emitOpError( - "Subcore id should not be set unless issuing core type is SC vector " - "subcore"); + if (sem) { + if (getRefCoreType(sem).value_or(issuing_core) != issuing_core) { + return emitOpError( + "Can only await semaphores attached to the local core"); + } } return success(); diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops_verification_test.cc b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops_verification_test.cc index b2e04ea4e28374..c76082db583498 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops_verification_test.cc +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops_verification_test.cc @@ -60,7 +60,7 @@ class TpuOpsVerificationTest : public ::testing::Test { context_.loadAllAvailableDialects(); context_.printOpOnDiagnostic(true); } - ~TpuOpsVerificationTest() { + ~TpuOpsVerificationTest() override { for (int i = ops_.size() - 1; i >= 0; --i) { ops_[i]->erase(); } @@ -730,36 +730,37 @@ TEST_F(TpuOpsVectorSubcoreVerificationTest, ScanVerificationInvalidDimension) { } TEST_F(TpuOpsVectorSubcoreVerificationTest, DmaElementTypeMismatch) { - auto dma = Create( - /*source=*/AllocaI32({1024, 256, 128}, MemorySpace::kHbm), - /*source_semaphore=*/AllocaSemaphore(), - /*target=*/ - Create(GetMemRefType({1024, 256, 128}, - builder().getI64Type(), - MemorySpace::kHbm)) - .getMemref(), - /*target_semaphore=*/AllocaSemaphore(), - /*device_id=*/nullptr, - /*core_id=*/nullptr); + auto src = AllocaI32({1024, 256, 128}, MemorySpace::kHbm); + auto dst = Create(GetMemRefType({1024, 256, 128}, + builder().getI64Type(), + MemorySpace::kHbm)) + .getMemref(); + auto sem = AllocaSemaphore(); ASSERT_THAT( - VerifyOp(dma), + VerifyOp(Create(src, sem, dst, sem, nullptr, nullptr)), + StatusIs(_, HasSubstr("DMA source and target element type mismatch"))); + + ASSERT_THAT( + VerifyOp(Create(src, nullptr, dst, sem, nullptr, nullptr, + nullptr, 0, false, /*wait_target=*/true)), StatusIs(_, HasSubstr("DMA source and target element type mismatch"))); } -TEST_F(TpuOpsVectorSubcoreVerificationTest, DmaDynamicRankMismatch) { - auto dma = Create( - /*source=*/AllocaI32({ShapedType::kDynamic, 256, 128}, MemorySpace::kHbm), - /*source_semaphore=*/AllocaSemaphore(), - /*target=*/ - AllocaI32({ShapedType::kDynamic, ShapedType::kDynamic, 128}, - MemorySpace::kHbm), - /*target_semaphore=*/AllocaSemaphore(), - /*device_id=*/nullptr, - /*core_id=*/nullptr); +TEST_F(TpuOpsVectorSubcoreVerificationTest, DmaShapeMismatch) { + auto src = AllocaI32({ShapedType::kDynamic, 256, 128}, MemorySpace::kHbm); + auto dst = AllocaI32({ShapedType::kDynamic, ShapedType::kDynamic, 128}, + MemorySpace::kHbm); + auto sem = AllocaSemaphore(); - ASSERT_THAT(VerifyOp(dma), - StatusIs(_, HasSubstr("DMA source and target shape mismatch."))); + ASSERT_THAT( + VerifyOp(Create(src, sem, dst, sem, nullptr, nullptr)), + StatusIs(_, HasSubstr("DMA source and target shape mismatch."))); + + ASSERT_THAT( + VerifyOp(Create(src, nullptr, dst, sem, nullptr, nullptr, + nullptr, 0, false, /*wait_target=*/true)), + StatusIs(_, HasSubstr("DMA source and target shape mismatch."))); } TEST_F(TpuOpsVectorSubcoreVerificationTest, DmaStrictOrderingSupported) { @@ -776,6 +777,82 @@ TEST_F(TpuOpsVectorSubcoreVerificationTest, DmaStrictOrderingSupported) { ASSERT_OK(VerifyOp(dma)); } +TEST_F(TpuOpsVectorSubcoreVerificationTest, DmaVectorSubcoreValidStreams) { + auto vmem = AllocaI32({1024, 256, 128}, MemorySpace::kVmem); + auto hbm = AllocaI32({1024, 256, 128}, MemorySpace::kHbm); + auto sem = AllocaSemaphore(); + + // Push stream (VMEM -> HBM): supports target sem, source sem, or no sem. + ASSERT_OK(VerifyOp( + Create(vmem, nullptr, hbm, sem, nullptr, nullptr))); + ASSERT_OK(VerifyOp( + Create(vmem, sem, hbm, nullptr, nullptr, nullptr))); + ASSERT_OK(VerifyOp( + Create(vmem, nullptr, hbm, nullptr, nullptr, nullptr))); + ASSERT_OK(VerifyOp(Create(vmem, nullptr, hbm, sem, nullptr, + nullptr, nullptr, 0, false, + /*wait_target=*/true))); + ASSERT_OK(VerifyOp(Create(vmem, sem, hbm, nullptr, nullptr, + nullptr, nullptr, 0, false, + /*wait_target=*/false))); + + // Pull stream (HBM -> VMEM): supports no sem. + ASSERT_OK(VerifyOp( + Create(hbm, nullptr, vmem, nullptr, nullptr, nullptr))); +} + +TEST_F(TpuOpsVectorSubcoreVerificationTest, + DmaPullStreamSourceSemaphoreNotSupported) { + auto vmem = AllocaI32({1024, 256, 128}, MemorySpace::kVmem); + auto hbm = AllocaI32({1024, 256, 128}, MemorySpace::kHbm); + auto smem = AllocaI32({1024, 256, 128}, MemorySpace::kSmem); + auto sem = AllocaSemaphore(); + + const char* err_msg = + "Source semaphores are unsupported for transfers from a " + "non-local memory to local memory on SparseCore"; + + // HBM -> VMEM pull stream with source semaphore. + ASSERT_THAT( + VerifyOp(Create(hbm, sem, vmem, nullptr, nullptr, nullptr)), + StatusIs(_, HasSubstr(err_msg))); + + // SMEM -> VMEM pull stream with source semaphore. + ASSERT_THAT(VerifyOp(Create(smem, sem, vmem, nullptr, nullptr, + nullptr)), + StatusIs(_, HasSubstr(err_msg))); + + // Awaiting source on pull stream without semaphore. + ASSERT_THAT( + VerifyOp(Create(hbm, nullptr, vmem, nullptr, nullptr, nullptr, + nullptr, 0, false, /*wait_target=*/false)), + StatusIs(_, HasSubstr("Awaiting source read completion is unsupported"))); +} + +TEST_F(TpuOpsVectorSubcoreVerificationTest, DmaMissingMemorySpaceOnSparseCore) { + auto vmem = AllocaI32({1024, 256, 128}, MemorySpace::kVmem); + auto unannotated = AllocaI32({1024, 256, 128}, std::nullopt); + auto sem = AllocaSemaphore(); + + const char* src_err_msg = + "Source memory space must be provided for SparseCore DMA"; + const char* tgt_err_msg = + "Target memory space must be provided for SparseCore DMA"; + + ASSERT_THAT(VerifyOp(Create(unannotated, nullptr, vmem, sem, + nullptr, nullptr)), + StatusIs(_, HasSubstr(src_err_msg))); + + ASSERT_THAT(VerifyOp(Create(vmem, sem, unannotated, nullptr, + nullptr, nullptr)), + StatusIs(_, HasSubstr(tgt_err_msg))); + + ASSERT_THAT(VerifyOp(Create(unannotated, nullptr, vmem, sem, + nullptr, nullptr, nullptr, 0, false, + /*wait_target=*/true)), + StatusIs(_, HasSubstr(src_err_msg))); +} + TEST_F(TpuOpsVerificationTest, DmaStrictOrderingNotSupportedOnTc) { auto func_op = Create("tc_kernel", builder().getFunctionType({}, {})); @@ -796,6 +873,89 @@ TEST_F(TpuOpsVerificationTest, DmaStrictOrderingNotSupportedOnTc) { ASSERT_THAT(VerifyOp(dma), StatusIs(_, HasSubstr("Strict ordering is only supported on the " "SC scalar and vector subcores"))); + + auto wait = Create( + /*source=*/AllocaI32({1024, 256, 128}, MemorySpace::kHbm), + /*source_semaphore=*/nullptr, + /*target=*/AllocaI32({1024, 256, 128}, MemorySpace::kVmem), + /*target_semaphore=*/AllocaSemaphore(), + /*device_id=*/nullptr, + /*core_id=*/nullptr, + /*subcore_id=*/nullptr, + /*priority=*/0, + /*strict_ordering=*/true, + /*wait_target=*/true); + + ASSERT_THAT(VerifyOp(wait), + StatusIs(_, HasSubstr("Strict ordering is only supported on the " + "SC scalar and vector subcores"))); +} + +TEST_F(TpuOpsVerificationTest, DmaTargetSemaphoreRequiredOnTc) { + auto func_op = + Create("tc_kernel", builder().getFunctionType({}, {})); + func_op->setAttr(TPUDialect::GetCoreTypeKey(), + CoreTypeAttr::get(builder().getContext(), CoreType::kTc)); + builder().setInsertionPointToStart(func_op.addEntryBlock()); + + auto src = AllocaI32({1024, 256, 128}, MemorySpace::kHbm); + auto dst = AllocaI32({1024, 256, 128}, MemorySpace::kVmem); + + auto dma = Create(src, /*source_semaphore=*/nullptr, dst, + /*target_semaphore=*/nullptr, + /*device_id=*/nullptr, /*core_id=*/nullptr); + ASSERT_THAT(VerifyOp(dma), + StatusIs(_, HasSubstr("DMA target semaphore must be specified"))); + + auto wait = Create( + src, /*source_semaphore=*/nullptr, dst, /*target_semaphore=*/nullptr, + /*device_id=*/nullptr, /*core_id=*/nullptr, /*subcore_id=*/nullptr, + /*priority=*/0, /*strict_ordering=*/false, /*wait_target=*/true); + ASSERT_THAT(VerifyOp(wait), + StatusIs(_, HasSubstr("DMA target semaphore must be specified"))); +} + +TEST_F(TpuOpsVerificationTest, DmaLocalSourceSemaphoreRejectedOnTc) { + auto func_op = + Create("tc_kernel", builder().getFunctionType({}, {})); + func_op->setAttr(TPUDialect::GetCoreTypeKey(), + CoreTypeAttr::get(builder().getContext(), CoreType::kTc)); + builder().setInsertionPointToStart(func_op.addEntryBlock()); + + auto src = AllocaI32({1024, 256, 128}, MemorySpace::kHbm); + auto dst = AllocaI32({1024, 256, 128}, MemorySpace::kVmem); + auto sem = AllocaSemaphore(); + + auto dma = Create(src, /*source_semaphore=*/sem, dst, + /*target_semaphore=*/sem, + /*device_id=*/nullptr, /*core_id=*/nullptr); + ASSERT_THAT( + VerifyOp(dma), + StatusIs(_, HasSubstr("DMA destination device_id or core_id must be " + "specified when source semaphore is specified"))); + + auto wait = Create( + src, /*source_semaphore=*/sem, dst, /*target_semaphore=*/sem, + /*device_id=*/nullptr, /*core_id=*/nullptr, /*subcore_id=*/nullptr, + /*priority=*/0, /*strict_ordering=*/false, /*wait_target=*/true); + ASSERT_THAT( + VerifyOp(wait), + StatusIs(_, HasSubstr("DMA destination device_id or core_id must be " + "specified when source semaphore is specified"))); +} + +TEST_F(TpuOpsVerificationTest, DmaRemoteTargetWaitAllowsNullSourceSemaphore) { + auto src = AllocaI32({1024, 256, 128}, MemorySpace::kHbm); + auto dst = AllocaI32({1024, 256, 128}, MemorySpace::kVmem); + auto sem = AllocaSemaphore(); + auto device_id = Create(0, 32).getResult(); + auto core_id = Create(0, 32).getResult(); + + auto wait = Create( + src, /*source_semaphore=*/nullptr, dst, /*target_semaphore=*/sem, + /*device_id=*/device_id, /*core_id=*/core_id, /*subcore_id=*/nullptr, + /*priority=*/0, /*strict_ordering=*/false, /*wait_target=*/true); + ASSERT_OK(VerifyOp(wait)); } TEST_F(TpuOpsVectorSubcoreVerificationTest, diff --git a/third_party/xla/xla/pjrt/BUILD b/third_party/xla/xla/pjrt/BUILD index 4becb8e978a00b..579a430fe77dbd 100644 --- a/third_party/xla/xla/pjrt/BUILD +++ b/third_party/xla/xla/pjrt/BUILD @@ -218,6 +218,7 @@ cc_library( "//xla:executable_run_options", "//xla:future", "//xla:literal", + "//xla:literal_util", "//xla:shape_util", "//xla:util", "//xla:xla_data_proto_cc", @@ -225,6 +226,7 @@ cc_library( "//xla/error:error_codes", "//xla/hlo/ir:hlo", "//xla/pjrt/c:pjrt_c_api_device_event_hdrs", + "//xla/pjrt/plugin/xla_cpu:cpu_topology", "//xla/runtime:device_id", "//xla/service:dump", "//xla/tsl/concurrency:async_value", diff --git a/third_party/xla/xla/pjrt/common_pjrt_client.cc b/third_party/xla/xla/pjrt/common_pjrt_client.cc index 13df6c962b4810..35bcc8f6ab5ff8 100644 --- a/third_party/xla/xla/pjrt/common_pjrt_client.cc +++ b/third_party/xla/xla/pjrt/common_pjrt_client.cc @@ -61,6 +61,7 @@ limitations under the License. #include "xla/layout.h" #include "xla/layout_util.h" #include "xla/literal.h" +#include "xla/literal_util.h" #include "xla/pjrt/abstract_tracked_device_buffer.h" #include "xla/pjrt/async_work_runner.h" #include "xla/pjrt/c/pjrt_c_api_device_event.h" @@ -73,6 +74,7 @@ limitations under the License. #include "xla/pjrt/pjrt_client.h" #include "xla/pjrt/pjrt_compiler.h" #include "xla/pjrt/pjrt_executable.h" +#include "xla/pjrt/plugin/xla_cpu/cpu_topology.h" #include "xla/pjrt/raw_buffer.h" #include "xla/pjrt/raw_pjrt_client.h" #include "xla/pjrt/staging_buffer.h" @@ -454,10 +456,15 @@ CommonPjRtClient::LoadInternal(std::shared_ptr executable, addressable_devices.reserve(num_replicas * num_partitions); for (int replica = 0; replica < num_replicas; ++replica) { for (int partition = 0; partition < num_partitions; ++partition) { - int64_t device_id = (*device_assignment)(replica, partition); - GlobalDeviceId global_device_id(device_id); + GlobalDeviceId device_id((*device_assignment)(replica, partition)); + // TODO(parkers): remove this dependence on CPU's UnpackCpuProcessIndex. + if (IsCpuId(platform_id()) && + UnpackCpuProcessIndex(device_id) != process_index()) { + VLOG(3) << "Non-local device: " << device_id; + continue; + } - ABSL_ASSIGN_OR_RETURN(PjRtDevice * device, LookupDevice(global_device_id)); + ABSL_ASSIGN_OR_RETURN(PjRtDevice * device, LookupDevice(device_id)); if (device->process_index() != process_index()) { VLOG(3) << "Non-local device: " << device_id; continue; @@ -486,7 +493,14 @@ CommonPjRtClient::LoadInternal(std::shared_ptr executable, *hlo_module, kAfterOptimizationsDumpName, ex_options.has_debug_options() ? &ex_options.debug_options() : nullptr); } - xla::Shape result_shape = hlo_module->result_shape(); + xla::Shape result_shape; + if (IsGpuId(platform_id())) { + // TODO(parkers): GPU isn't properly setting entry_computation_layout? + result_shape = hlo_module->result_shape(); + } else { + result_shape = + hlo_module->entry_computation_layout().result_layout().shape(); + } absl::Span result_shapes = result_shape.IsTuple() ? absl::MakeSpan(result_shape.tuple_shapes()) : absl::MakeSpan(&result_shape, 1); @@ -533,7 +547,12 @@ CommonPjRtClient::LoadInternal(std::shared_ptr executable, std::move(addressable_devices), nullptr, compile_options.parameter_is_tupled_arguments, std::move(input_hlo_snapshot_bits))); - + // TODO(parkers): Clear cpu memory spaces because CPU previously would + // accept inputs on any memory space. If this is fixed at some point, remove + // this. + if (IsCpuId(platform_id())) { + dispatch_info.parameter_memory_space_kind_ids.clear(); + } auto load_state = raw_client()->MakeLoadState(); ABSL_RETURN_IF_ERROR(load_state->Preload(executable.get())); auto loaded_executable = std::make_unique( @@ -2724,11 +2743,113 @@ CommonPjRtLoadedExecutable::ExecutePortable( return std::move(result.buffers); } +static void MaybeDumpHloSnapshot( + const HloModule& module, RunId run_id, + const std::vector& arguments, + const std::vector>& results, + absl::string_view file_name_prefix = "") { + if (!DumpingEnabledForHloModule(module)) { + return; + } + if (!module.config().debug_options().xla_dump_hlo_snapshots()) { + return; + } + xla::HloSnapshot hlo_snapshot; + *hlo_snapshot.mutable_hlo()->mutable_hlo_module() = module.ToProto(); + + for (auto* argument : arguments) { + auto literal_or = argument->ToLiteral().Await(); + if (!literal_or.ok()) { + LOG(ERROR) << "Failed to get literal for argument: " + << literal_or.status(); + return; + } + *hlo_snapshot.add_arguments() = (*literal_or)->ToProto(); + } + + // If there are multiple results, wrap them in a tuple. + if (results.size() == 1) { + auto literal_or = results[0]->ToLiteral().Await(); + if (!literal_or.ok()) { + LOG(ERROR) << "Failed to get literal for result: " << literal_or.status(); + return; + } + *hlo_snapshot.mutable_result() = (*literal_or)->ToProto(); + } else { + std::vector result_literals; + result_literals.reserve(results.size()); + for (auto& result : results) { + auto literal_or = result->ToLiteral().Await(); + if (!literal_or.ok()) { + LOG(ERROR) << "Failed to get literal for result: " + << literal_or.status(); + return; + } + result_literals.push_back(std::move(**literal_or)); + } + *hlo_snapshot.mutable_result() = + LiteralUtil::MakeTupleOwned(std::move(result_literals)).ToProto(); + } + + DumpToFileInDir( + module, "", + absl::StrCat(file_name_prefix, "snapshot.", run_id.ToInt(), ".pb"), + hlo_snapshot.SerializeAsString()); +} + absl::StatusOr>>> CommonPjRtLoadedExecutable::Execute( absl::Span> argument_handles, const ExecuteOptions& options, std::optional>>& returned_futures) const { + if (addressable_devices_.size() == 1 && argument_handles.size() == 1 && + IsCpuId(client()->platform_id())) { + std::vector>> wrapped_results(1); + RunId run_id = options.launch_id != 0 ? RunId(options.launch_id) + : RunId::CreateUniqueId(); + // Fast-path if there is only one device — run the computation on the + // current thread. + const int replica = addressable_device_logical_ids_[0].replica; + const int partition = addressable_device_logical_ids_[0].partition; + + ABSL_ASSIGN_OR_RETURN(auto hlo_module, GetExecutable()->GetHloModule()); + + // Dump once before running, in case there's a crash. + MaybeDumpHloSnapshot(*hlo_module, run_id, argument_handles[0], {}); + if (auto unoptimized_hlo_module = + GetExecutable()->GetUnoptimizedHloModule()) { + HloUnoptimizedSnapshot hlo_snapshot; + *hlo_snapshot.mutable_hlo_module() = *std::move(unoptimized_hlo_module); + for (const auto& argument_handle : argument_handles) { + HloInputs hlo_inputs; + for (const auto& buffer : argument_handle) { + ABSL_ASSIGN_OR_RETURN(auto literal, buffer->ToLiteral().Await()); + *hlo_inputs.add_arguments() = literal->ToProto(); + } + *hlo_snapshot.add_partitions() = std::move(hlo_inputs); + } + + DumpHloUnoptimizedSnapshotIfEnabled(hlo_snapshot, + hlo_module->config().debug_options()); + } + auto statusor = ExecuteHelperOnSingleDevice(argument_handles[0], run_id, + replica, partition, options, + returned_futures.has_value()); + + if (!statusor.ok()) { + return std::move(statusor).status(); + } + + wrapped_results[0] = std::move(statusor->buffers); + if (returned_futures.has_value()) { + returned_futures->push_back(std::move(*statusor->future)); + } + + MaybeDumpHloSnapshot(*hlo_module, run_id, argument_handles[0], + wrapped_results[0]); + return wrapped_results; + } + if (input_hlo_snapshot_bits()) { HloUnoptimizedSnapshot hlo_snapshot; *hlo_snapshot.mutable_hlo_module() = input_hlo_snapshot_bits()->hlo_module; diff --git a/third_party/xla/xla/pjrt/cpu/BUILD b/third_party/xla/xla/pjrt/cpu/BUILD index e8170242698e7f..9b18d45c866e6f 100644 --- a/third_party/xla/xla/pjrt/cpu/BUILD +++ b/third_party/xla/xla/pjrt/cpu/BUILD @@ -164,15 +164,12 @@ cc_library( hdrs = ["cpu_device.h"], visibility = internal_visibility(["//xla/pjrt/cpu:legacy_cpu_topology_users"]), deps = [ - ":cpu_async_execution_tracker", - ":execution_stream_event_map", "//xla:literal", + "//xla/pjrt:common_pjrt_client", "//xla/pjrt:host_memory_spaces", "//xla/pjrt:pjrt_client", "//xla/pjrt:pjrt_common", - "//xla/pjrt:semaphore", "//xla/pjrt/plugin/xla_cpu:cpu_device_description", - "//xla/service/cpu:cpu_xfeed", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:inlined_vector", @@ -208,6 +205,7 @@ cc_library( ":cpu_async_execution_tracker", ":cpu_device", ":cpu_event", + ":execution_stream_event_map", "//xla:array", "//xla:debug_options_flags", "//xla:executable_run_options", @@ -274,6 +272,7 @@ cc_library( "//xla/service/cpu:cpu_compiler", "//xla/service/cpu:cpu_executable", "//xla/service/cpu:cpu_executable_run_options", + "//xla/service/cpu:cpu_xfeed", "//xla/service/cpu:executable_proto_cc", "//xla/service/llvm_ir:llvm_command_line_options", "//xla/stream_executor:device_address", diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.cc b/third_party/xla/xla/pjrt/cpu/cpu_client.cc index 0aae628700402a..2e375712b2cc80 100644 --- a/third_party/xla/xla/pjrt/cpu/cpu_client.cc +++ b/third_party/xla/xla/pjrt/cpu/cpu_client.cc @@ -108,6 +108,7 @@ limitations under the License. #include "xla/service/cpu/cpu_compiler.h" #include "xla/service/cpu/cpu_executable.h" #include "xla/service/cpu/cpu_executable_run_options.h" +#include "xla/service/cpu/cpu_xfeed.h" #include "xla/service/cpu/executable.pb.h" #include "xla/service/device_assignment.h" #include "xla/service/dump.h" @@ -367,20 +368,20 @@ absl::StatusOr> GetPjRtCpuClient( } } + auto raw_client = std::make_unique( + std::move(allocator), std::move(options.collectives), num_threads, + options.asynchronous, options.max_transpose_threads, + std::move(options.customize_hlo_module_config), cpu_device_count, + options.max_inflight_computations_per_device); + std::vector> devices; devices.reserve(topology->cpu_topology().number_of_devices()); for (const auto& topology_device : topology->cpu_topology().devices()) { auto device = std::make_unique( - topology_device.process_id, topology_device.local_device_id, - options.max_inflight_computations_per_device); + topology_device.process_id, topology_device.local_device_id); devices.push_back(std::move(device)); } - auto raw_client = std::make_unique( - std::move(allocator), std::move(options.collectives), num_threads, - options.asynchronous, options.max_transpose_threads, - std::move(options.customize_hlo_module_config)); - return std::unique_ptr( new PjRtCpuClient(options.process_id, std::move(devices), std::move(raw_client), std::move(topology))); @@ -401,12 +402,58 @@ static tsl::ThreadOptions GetThreadOptions() { return thread_options; } +absl::StatusOr PjRtCpuRawClient::PoisonExecution( + LocalDeviceId local_device_id, int32_t launch_id, absl::Status error) { + if (!(local_device_id >= 0 && + local_device_id < local_device_states_.size())) { + return absl::InvalidArgumentError(absl::StrCat( + "PjRtCpuRawClient: ", local_device_id, " is out of range: [0, ", + local_device_states_.size(), ") for local devices")); + } + return local_device_states_[local_device_id.value()] + ->async_execution_tracker() + ->SetError(launch_id, std::move(error)); +} + +absl::Status PjRtCpuRawClient::TransferToInfeed(LocalDeviceId local_device_id, + const LiteralSlice& literal) { + return TransferLiteralToInfeedOnCpu(local_device_id.value(), literal); +} + +absl::Status PjRtCpuRawClient::TransferFromOutfeed( + LocalDeviceId local_device_id, MutableBorrowingLiteral literal) { + return TransferLiteralFromOutfeedOnCpu(local_device_id.value(), literal); +} + +PjRtCpuRawClient::LocalDeviceState* PjRtCpuRawClient::GetLocalDeviceState( + LocalDeviceId local_device_id) { + CHECK((local_device_id.value() < local_device_states_.size()) && + (local_device_id.value() >= 0)) + << "Local device id " << local_device_id << " not in range: [0, " + << local_device_states_.size() << ")"; + return local_device_states_[local_device_id.value()].get(); +} + +static std::vector> +MakeLocalDeviceStates(int cpu_device_count, int max_inflight_computations) { + std::vector> results; + results.reserve(cpu_device_count); + for (size_t i = 0; i < cpu_device_count; ++i) { + results.push_back(std::make_unique( + max_inflight_computations)); + } + return results; +} + PjRtCpuRawClient::PjRtCpuRawClient( std::shared_ptr allocator, std::shared_ptr collectives, size_t num_threads, bool asynchronous, int max_transpose_threads, - std::function customize_hlo_module_config) - : allocator_(std::move(allocator)), + std::function customize_hlo_module_config, + int cpu_device_count, int max_inflight_computations) + : local_device_states_( + MakeLocalDeviceStates(cpu_device_count, max_inflight_computations)), + allocator_(std::move(allocator)), collectives_(std::move(collectives)), asynchronous_(asynchronous), max_transpose_threads_(max_transpose_threads), @@ -647,67 +694,6 @@ PjRtCpuExecutable::Deserialize(riegeli::Any reader, return cpu_executable; } -absl::StatusOr> -PjRtCpuClient::LoadInternal( - std::shared_ptr cpu_executable, - std::shared_ptr device_assignment) { - int num_replicas = cpu_executable->num_replicas(); - int num_partitions = cpu_executable->num_partitions(); - std::vector - addressable_device_logical_ids; - std::vector addressable_devices; - if (device_assignment != nullptr) { - addressable_device_logical_ids.reserve(num_replicas * num_partitions); - addressable_devices.reserve(num_replicas * num_partitions); - for (int replica = 0; replica < num_replicas; ++replica) { - for (int partition = 0; partition < num_partitions; ++partition) { - GlobalDeviceId device_id((*device_assignment)(replica, partition)); - if (UnpackCpuProcessIndex(device_id) != process_index()) { - VLOG(3) << "Non-local device: " << device_id; - continue; - } - ABSL_ASSIGN_OR_RETURN(PjRtDevice * device, LookupDevice(device_id)); - PjRtLoadedExecutable::LogicalDeviceIds logica_device_ids; - logica_device_ids.replica = replica; - logica_device_ids.partition = partition; - addressable_device_logical_ids.push_back(std::move(logica_device_ids)); - addressable_devices.push_back(device); - } - } - } - const auto& result_shape = cpu_executable->cpu_executable()->result_shape(); - if (result_shape.IsTuple()) { - for (auto& leaf_shape : result_shape.tuple_shapes()) { - if (leaf_shape.IsTuple()) { - return absl::InternalError(absl::StrCat( - "Nested tuples are not supported with PjRtCpuClient. got: ", - result_shape.ToString())); - } - } - } - auto load_state = tsl::MakeRef(this); - ABSL_ASSIGN_OR_RETURN( - auto parameters_that_may_be_donated, - ComputeParametersThatMayBeDonated( - *cpu_executable->cpu_executable_->shared_module(), - cpu_executable->compile_options_.parameter_is_tupled_arguments)); - return std::make_unique( - this, tsl::MakeAvailableAsyncValueRef(cpu_executable), - CommonPjRtLoadedExecutable::DispatchInfo{ - cpu_executable->parameter_device_shapes_, - std::make_shared( - cpu_executable->cpu_executable_->result_shape()), - std::vector(), - cpu_executable->output_memory_space_kind_ids_, - std::move(addressable_devices), - std::move(addressable_device_logical_ids), - std::move(device_assignment), - std::move(parameters_that_may_be_donated), - cpu_executable->input_buffer_sizes_in_bytes_, - }, - std::move(load_state)); -} - tsl::AsyncValueRef PjRtCpuRawClient::ToAsyncExecutable( std::shared_ptr executable) const { return tsl::MakeAvailableAsyncValueRef( @@ -801,31 +787,6 @@ absl::StatusOr> PjRtCpuRawClient::Compile( std::move(options), topology, process_index); } -absl::StatusOr> PjRtCpuClient::Load( - std::shared_ptr executable, - const LoadOptions& load_options) { - auto* cpu_exec_ptr = dynamic_cast(executable.get()); - if (cpu_exec_ptr == nullptr) { - return absl::InvalidArgumentError( - "PjRtCpuClient::Load: executable is not a PjRtCpuExecutable."); - } - auto cpu_executable = - std::static_pointer_cast(std::move(executable)); - CompileOptions options = cpu_executable->compile_options(); - int unused_num_replicas; - int unused_num_partitions; - std::shared_ptr device_assignment; - ABSL_RETURN_IF_ERROR(ParseDeviceAssignmentCompileOptions( - options.compile_portable_executable, &options.executable_build_options, - [this](int num_replicas, int num_partitions) { - return topology().GetDefaultDeviceAssignment(process_index(), - num_replicas, std::nullopt, - num_partitions, nullptr); - }, - &unused_num_replicas, &unused_num_partitions, &device_assignment)); - return LoadInternal(std::move(cpu_executable), std::move(device_assignment)); -} - absl::StatusOr> PjRtCpuRawClient::CompileAheadOfTime(const XlaComputation& computation, CompileOptions options, @@ -1388,6 +1349,10 @@ CreateBufferTable(const BufferAssignment& assignment, return std::move(buffer_table); } +tsl::RCReference PjRtCpuRawClient::MakeLoadState() { + return tsl::MakeRef(); +} + absl::StatusOr> CpuExecutableLoadState::LoadRawExecutable( tsl::AsyncValueRef executable, @@ -1395,7 +1360,10 @@ CpuExecutableLoadState::LoadRawExecutable( DeviceAndAssignment device_and_assign, int attempt) { auto result = std::make_unique(run_id); result->executable_ = absl::down_cast(&executable.get()); - result->raw_client_ = client_->raw_client(); + auto* client = + absl::down_cast(device_and_assign.device->client()); + result->raw_client_ = + absl::down_cast(client->raw_client()); int num_addressable_devices = 0; if (device_and_assign.device_assignment != nullptr) { for (int r = 0; r < device_and_assign.device_assignment->replica_count(); @@ -1403,7 +1371,7 @@ CpuExecutableLoadState::LoadRawExecutable( for (int p = 0; p < device_and_assign.device_assignment->computation_count(); ++p) { GlobalDeviceId device_id((*device_and_assign.device_assignment)(r, p)); - if (UnpackCpuProcessIndex(device_id) == client_->process_index()) { + if (UnpackCpuProcessIndex(device_id) == client->process_index()) { ++num_addressable_devices; } } @@ -1411,7 +1379,8 @@ CpuExecutableLoadState::LoadRawExecutable( } result->num_addressable_devices_ = num_addressable_devices; result->device_assignment_ = std::move(device_and_assign.device_assignment); - result->device_ = absl::down_cast(device_and_assign.device); + result->local_device_id_ = device_and_assign.device->local_device_id(); + result->global_device_id_ = device_and_assign.device->global_device_id(); return result; } @@ -1450,6 +1419,7 @@ PjRtRawLoadedExecutable::RawExecuteResult CpuPjRtRawLoadedExecutable::Execute( absl::Span output_leaf_buffers, PjRtDeviceEventRefVector extra_deps, PjRtDeviceEventRefVector control_deps, bool is_predetermined_error, bool fill_future) && { + auto* local_device_state = raw_client_->GetLocalDeviceState(local_device_id_); PjRtRawLoadedExecutable::RawExecuteResult result; // `returned_future_can_be_set_event` indicates when `returned_future` can be // set using `execute_event`. This is necessary to delay setting the @@ -1528,7 +1498,8 @@ PjRtRawLoadedExecutable::RawExecuteResult CpuPjRtRawLoadedExecutable::Execute( // allows the inputs for the next executable to be fetched even if the // launch is delayed. auto compute_reservation = std::make_unique( - device_->max_inflight_computations_semaphore().ScopedAcquire(1)); + local_device_state->max_inflight_computations_semaphore().ScopedAcquire( + 1)); ExecutableRunOptions run_options; run_options.set_run_id(run_id_); @@ -1548,10 +1519,10 @@ PjRtRawLoadedExecutable::RawExecuteResult CpuPjRtRawLoadedExecutable::Execute( cpu_execute_context->process_index().has_value()) { run_options.set_device_ordinal( PackCpuDeviceId(*cpu_execute_context->process_index(), - UnpackCpuLocalDeviceId(device_->global_device_id())) + UnpackCpuLocalDeviceId(global_device_id_)) .value()); } else { - run_options.set_device_ordinal(device_->global_device_id().value()); + run_options.set_device_ordinal(global_device_id_.value()); } if (cpu_execute_context != nullptr && cpu_execute_context->collectives() != nullptr) { @@ -1575,8 +1546,9 @@ PjRtRawLoadedExecutable::RawExecuteResult CpuPjRtRawLoadedExecutable::Execute( } else { // This is a non-parallel computation. Add the last enqueue event as a // dependency with any error cleared. - auto last_enqueue_event = device_->stream_event_map()->GetLastEnqueueEvent( - options.execution_stream_id); + auto last_enqueue_event = + local_device_state->stream_event_map()->GetLastEnqueueEvent( + options.execution_stream_id); if (!last_enqueue_event.IsAvailable()) { auto last_enqueue_done_event = tsl::MakeUnconstructedAsyncValueRef(); @@ -1704,7 +1676,7 @@ PjRtRawLoadedExecutable::RawExecuteResult CpuPjRtRawLoadedExecutable::Execute( if (!is_a_collective_launch) { // This is a non-parallel computation. Set the execute event as the new // last enqueue event. - auto* stream_event_map = device_->stream_event_map(); + auto* stream_event_map = local_device_state->stream_event_map(); stream_event_map->SetLastEnqueueEvent(options.execution_stream_id, execute_event.CopyRef()); execute_event.AndThen([stream_event_map, @@ -1714,7 +1686,7 @@ PjRtRawLoadedExecutable::RawExecuteResult CpuPjRtRawLoadedExecutable::Execute( }); } CpuScopedAsyncExecution scoped_async_execution = - device_->async_execution_tracker()->NewAsyncExecution( + local_device_state->async_execution_tracker()->NewAsyncExecution( run_id_.ToInt(), std::move(ready_on_exit).Release()); PjRtDeviceEventSpan events_ref(input_deps); xla::ExecuteWhenReady( @@ -1789,139 +1761,4 @@ PjRtRawLoadedExecutable::RawExecuteResult CpuPjRtRawLoadedExecutable::Execute( return result; } -static void MaybeDumpHloSnapshot( - const HloModule& module, RunId run_id, - const std::vector& arguments, - const std::vector>& results, - absl::string_view file_name_prefix = "") { - if (!DumpingEnabledForHloModule(module)) { - return; - } - if (!module.config().debug_options().xla_dump_hlo_snapshots()) { - return; - } - xla::HloSnapshot hlo_snapshot; - *hlo_snapshot.mutable_hlo()->mutable_hlo_module() = module.ToProto(); - - for (auto* argument : arguments) { - auto literal_or = argument->ToLiteral().Await(); - if (!literal_or.ok()) { - LOG(ERROR) << "Failed to get literal for argument: " - << literal_or.status(); - return; - } - *hlo_snapshot.add_arguments() = (*literal_or)->ToProto(); - } - - // If there are multiple results, wrap them in a tuple. - if (results.size() == 1) { - auto literal_or = results[0]->ToLiteral().Await(); - if (!literal_or.ok()) { - LOG(ERROR) << "Failed to get literal for result: " << literal_or.status(); - return; - } - *hlo_snapshot.mutable_result() = (*literal_or)->ToProto(); - } else { - std::vector result_literals; - result_literals.reserve(results.size()); - for (auto& result : results) { - auto literal_or = result->ToLiteral().Await(); - if (!literal_or.ok()) { - LOG(ERROR) << "Failed to get literal for result: " - << literal_or.status(); - return; - } - result_literals.push_back(std::move(**literal_or)); - } - *hlo_snapshot.mutable_result() = - LiteralUtil::MakeTupleOwned(std::move(result_literals)).ToProto(); - } - - DumpToFileInDir( - module, "", - absl::StrCat(file_name_prefix, "snapshot.", run_id.ToInt(), ".pb"), - hlo_snapshot.SerializeAsString()); -} - -absl::StatusOr>>> -PjRtCpuLoadedExecutable::Execute( - absl::Span> argument_handles, - const ExecuteOptions& options, - std::optional>>& returned_futures) const { - if (device_assignment_ == nullptr) { - return InvalidArgument("Execute expects a non-null device_assignment"); - } - if (addressable_devices_.size() == 1 && argument_handles.size() == 1) { - std::vector>> wrapped_results(1); - RunId run_id = options.launch_id != 0 ? RunId(options.launch_id) - : RunId::CreateUniqueId(); - // Fast-path if there is only one device — run the computation on the - // current thread. - const int replica = addressable_device_logical_ids_[0].replica; - const int partition = addressable_device_logical_ids_[0].partition; - - // Dump once before running, in case there's a crash. - MaybeDumpHloSnapshot(GetExecutable()->cpu_executable_->module(), run_id, - argument_handles[0], {}); - if (GetExecutable()->unoptimized_hlo_module_ != nullptr) { - HloUnoptimizedSnapshot hlo_snapshot; - *hlo_snapshot.mutable_hlo_module() = - GetExecutable()->unoptimized_hlo_module_->ToProto(); - for (const auto& argument_handle : argument_handles) { - HloInputs hlo_inputs; - for (const auto& buffer : argument_handle) { - ABSL_ASSIGN_OR_RETURN(auto literal, buffer->ToLiteral().Await()); - *hlo_inputs.add_arguments() = literal->ToProto(); - } - *hlo_snapshot.add_partitions() = std::move(hlo_inputs); - } - - DumpHloUnoptimizedSnapshotIfEnabled( - hlo_snapshot, - GetExecutable()->cpu_executable_->module().config().debug_options()); - } - auto statusor = ExecuteHelperOnSingleDevice(argument_handles[0], run_id, - replica, partition, options, - returned_futures.has_value()); - - if (!statusor.ok()) { - return std::move(statusor).status(); - } - - wrapped_results[0] = std::move(statusor->buffers); - if (returned_futures.has_value()) { - returned_futures->push_back(std::move(*statusor->future)); - } - - MaybeDumpHloSnapshot(GetExecutable()->cpu_executable_->module(), run_id, - argument_handles[0], wrapped_results[0]); - return wrapped_results; - } - return CommonPjRtLoadedExecutable::Execute(argument_handles, options, - returned_futures); -} - -absl::StatusOr>> -PjRtCpuLoadedExecutable::ExecuteSharded( - absl::Span argument_handles, PjRtDevice* device, - const ExecuteOptions& options, std::optional>& returned_future, - bool fill_future) const { - if (device_assignment_ == nullptr) { - return InvalidArgument("ExecuteShard expects a non-null device_assignment"); - } - return CommonPjRtLoadedExecutable::ExecuteSharded( - argument_handles, device, options, returned_future, fill_future); -} - -absl::StatusOr>> -PjRtCpuLoadedExecutable::ExecutePortable( - absl::Span argument_handles, PjRtDevice* device, - const ExecuteOptions& options, std::optional>& returned_future, - bool fill_future) const { - if (device_assignment_ != nullptr) { - return InvalidArgument("ExecutePortable gets a non-portable executable"); - } - return CommonPjRtLoadedExecutable::ExecutePortable( - argument_handles, device, options, returned_future, fill_future); -} } // namespace xla diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.h b/third_party/xla/xla/pjrt/cpu/cpu_client.h index 6c72e89b619bea..96adb5a65abb9c 100644 --- a/third_party/xla/xla/pjrt/cpu/cpu_client.h +++ b/third_party/xla/xla/pjrt/cpu/cpu_client.h @@ -48,9 +48,11 @@ limitations under the License. #include "xla/pjrt/async_work_runner.h" #include "xla/pjrt/common_pjrt_client.h" #include "xla/pjrt/compiled_memory_stats.h" +#include "xla/pjrt/cpu/cpu_async_execution_tracker.h" #include "xla/pjrt/cpu/cpu_device.h" #include "xla/pjrt/cpu/cpu_device_memory.h" #include "xla/pjrt/cpu/cpu_event.h" +#include "xla/pjrt/cpu/execution_stream_event_map.h" #include "xla/pjrt/device_event.h" #include "xla/pjrt/dynamic_shapes.h" #include "xla/pjrt/maybe_owning_mlir_module.h" @@ -62,6 +64,7 @@ limitations under the License. #include "xla/pjrt/plugin/xla_cpu/cpu_topology_description.h" #include "xla/pjrt/raw_buffer.h" #include "xla/pjrt/raw_pjrt_client.h" +#include "xla/pjrt/semaphore.h" #include "xla/pjrt/thread_pool_async_work_runner.h" #include "xla/runtime/device_id.h" #include "xla/service/buffer_assignment.h" @@ -103,8 +106,8 @@ class PjRtCpuRawClient : public PjRtRawClient { std::shared_ptr allocator, std::shared_ptr collectives, size_t num_threads, bool asynchronous, int max_transpose_threads, - std::function customize_hlo_module_config = - nullptr); + std::function customize_hlo_module_config, + int cpu_device_count, int max_inflight_computations); ~PjRtCpuRawClient() override; @@ -182,6 +185,53 @@ class PjRtCpuRawClient : public PjRtRawClient { tsl::AsyncValueRef ToAsyncExecutable( std::shared_ptr executable) const override; + tsl::RCReference MakeLoadState() override; + + absl::StatusOr PoisonExecution(LocalDeviceId local_device_id, + int32_t launch_id, + absl::Status error) override; + + absl::Status TransferToInfeed(LocalDeviceId local_device_id, + const LiteralSlice& literal) override; + + absl::Status TransferFromOutfeed(LocalDeviceId local_device_id, + MutableBorrowingLiteral literal) override; + + class LocalDeviceState { + public: + explicit LocalDeviceState(int max_inflight_computations = 32) + : max_inflight_computations_semaphore_( + /*capacity=*/max_inflight_computations), + async_execution_tracker_( + std::make_unique()), + stream_event_map_(std::make_unique()) {} + + // Returns a semaphore for admission control on inflight computations. + Semaphore& max_inflight_computations_semaphore() { + return max_inflight_computations_semaphore_; + } + + CpuAsyncExecutionTracker* async_execution_tracker() { + return async_execution_tracker_.get(); + } + + ExecutionStreamEventMap* stream_event_map() const { + return stream_event_map_.get(); + } + + private: + // TODO(zhangqiaorjc): Optimize semaphore related overhead. + // Semaphore used to limit how many programs can be enqueued by the host + // ahead of the device. + Semaphore max_inflight_computations_semaphore_; + + std::unique_ptr async_execution_tracker_; + + std::unique_ptr stream_event_map_; + }; + + LocalDeviceState* GetLocalDeviceState(LocalDeviceId local_device_id); + private: friend class PjRtCpuClient; friend class CpuExecutableLoadState; @@ -195,6 +245,8 @@ class PjRtCpuRawClient : public PjRtRawClient { int process_index, const AotCompilationOptions* absl_nullable aot_options = nullptr); + std::vector> local_device_states_; + // A memory allocator used to allocate host memory for PjRtBuffers, and // temporary allocations passed to XLA:CPU executable. std::shared_ptr allocator_; @@ -246,20 +298,6 @@ class PjRtCpuClient final : public CommonPjRtClientImpl { public: ~PjRtCpuClient() override; - PjRtCpuRawClient* raw_client() const override { - return absl::down_cast( - CommonPjRtClientImpl::raw_client()); - } - - absl::StatusOr> Load( - std::shared_ptr executable, - const LoadOptions& load_options) override; - - const xla::CpuTopologyDescription& topology() const { - return *absl::down_cast( - &CommonPjRtClientImpl::topology()); - } - private: friend class PjRtCpuLoadedExecutable; friend class CpuPjRtRawLoadedExecutable; @@ -271,10 +309,6 @@ class PjRtCpuClient final : public CommonPjRtClientImpl { std::vector> devices, std::unique_ptr raw_client, std::unique_ptr topology); - - absl::StatusOr> LoadInternal( - std::shared_ptr cpu_executable, - std::shared_ptr device_assignment); }; class PjRtCpuLoadedExecutable; @@ -300,14 +334,15 @@ class CpuPjRtRawLoadedExecutable : public PjRtRawLoadedExecutable { const PjRtCpuExecutable* executable_; std::shared_ptr device_assignment_; size_t num_addressable_devices_; - PjRtCpuDevice* device_; + LocalDeviceId local_device_id_; + GlobalDeviceId global_device_id_; PjRtCpuRawClient* raw_client_; RunId run_id_; }; class CpuExecutableLoadState : public PjRtExecutableLoadState { public: - explicit CpuExecutableLoadState(PjRtCpuClient* client) : client_(client) {} + explicit CpuExecutableLoadState() = default; ~CpuExecutableLoadState() override = default; @@ -320,10 +355,7 @@ class CpuExecutableLoadState : public PjRtExecutableLoadState { xla::RunId run_id, DeviceAndAssignment device_and_assign, int attempt) override; - PjRtCpuClient* client() const { return client_; } - private: - PjRtCpuClient* client_; std::atomic is_deleted_{false}; }; @@ -380,6 +412,13 @@ class PjRtCpuExecutable final : public PjRtExecutable { const CompileOptions& compile_options() const { return compile_options_; } + std::optional GetUnoptimizedHloModule() const override { + if (!unoptimized_hlo_module_) { + return std::nullopt; + } + return unoptimized_hlo_module_->ToProto(); + } + static absl::StatusOr> Deserialize( riegeli::Any reader, const xla::CpuTopologyDescription& topology, @@ -425,47 +464,6 @@ class PjRtCpuExecutable final : public PjRtExecutable { const CpuTopologyDescription* topology_; }; -class PjRtCpuLoadedExecutable final : public CommonPjRtLoadedExecutable { - public: - using CommonPjRtLoadedExecutable::CommonPjRtLoadedExecutable; - - ~PjRtCpuLoadedExecutable() override = default; - - PjRtCpuExecutable* GetExecutable() const override { - return absl::down_cast( - CommonPjRtLoadedExecutable::GetExecutable()); - } - - PjRtCpuClient* client() const override { - return absl::down_cast( - CommonPjRtLoadedExecutable::client()); - } - - using PjRtLoadedExecutable::Execute; - absl::StatusOr>>> Execute( - absl::Span> argument_handles, - const ExecuteOptions& options, - std::optional>>& returned_futures) const override; - - using PjRtLoadedExecutable::ExecuteSharded; - absl::StatusOr>> ExecuteSharded( - absl::Span argument_handles, PjRtDevice* device, - const ExecuteOptions& options, std::optional>& returned_future, - bool fill_future) const override; - - using PjRtLoadedExecutable::ExecutePortable; - absl::StatusOr>> ExecutePortable( - absl::Span argument_handles, PjRtDevice* device, - const ExecuteOptions& options, std::optional>& returned_future, - bool fill_future) const override; - - const HloInputOutputAliasConfig& input_output_alias_config() const override { - return GetExecutable() - ->cpu_executable_->module() - .input_output_alias_config(); - } -}; - absl::StatusOr> ABSL_DEPRECATED( "Use public XLA:CPU GetXlaPjRtCpuClient instead") GetPjRtCpuClient(CpuClientOptions options); diff --git a/third_party/xla/xla/pjrt/cpu/cpu_device.cc b/third_party/xla/xla/pjrt/cpu/cpu_device.cc index ffe6d935b88746..fda33f34deea15 100644 --- a/third_party/xla/xla/pjrt/cpu/cpu_device.cc +++ b/third_party/xla/xla/pjrt/cpu/cpu_device.cc @@ -28,29 +28,25 @@ limitations under the License. #include "absl/strings/string_view.h" #include "absl/types/span.h" #include "xla/literal.h" -#include "xla/pjrt/cpu/cpu_async_execution_tracker.h" -#include "xla/pjrt/cpu/execution_stream_event_map.h" #include "xla/pjrt/host_memory_spaces.h" #include "xla/pjrt/pjrt_client.h" -#include "xla/service/cpu/cpu_xfeed.h" namespace xla { -PjRtCpuDevice::PjRtCpuDevice(int process_id, int local_device_id, - int max_inflight_computations) - : description_(process_id, local_device_id), - max_inflight_computations_semaphore_( - /*capacity=*/max_inflight_computations), - async_execution_tracker_(std::make_unique()), - stream_event_map_(std::make_unique()) {} +PjRtCpuDevice::PjRtCpuDevice(int process_id, int local_device_id) + : description_(process_id, local_device_id) {} absl::Status PjRtCpuDevice::TransferToInfeed(const LiteralSlice& literal) { - return TransferLiteralToInfeedOnCpu(local_hardware_id().value(), literal); + return absl::down_cast(client_) + ->raw_client() + ->TransferToInfeed(local_device_id(), literal); } absl::Status PjRtCpuDevice::TransferFromOutfeed( MutableBorrowingLiteral literal) { - return TransferLiteralFromOutfeedOnCpu(local_hardware_id().value(), literal); + return absl::down_cast(client_) + ->raw_client() + ->TransferFromOutfeed(local_device_id(), literal); } void PjRtCpuDevice::AttachMemorySpace(PjRtMemorySpace* memory_space) { @@ -102,7 +98,9 @@ absl::StatusOr PjRtCpuDevice::memory_space_by_kind_id( absl::StatusOr PjRtCpuDevice::PoisonExecution(int32_t launch_id, absl::Status error) { - return async_execution_tracker_->SetError(launch_id, std::move(error)); + return absl::down_cast(client_) + ->raw_client() + ->PoisonExecution(local_device_id(), launch_id, std::move(error)); } } // namespace xla diff --git a/third_party/xla/xla/pjrt/cpu/cpu_device.h b/third_party/xla/xla/pjrt/cpu/cpu_device.h index 13cbef9cfaef63..24b2b7067b439e 100644 --- a/third_party/xla/xla/pjrt/cpu/cpu_device.h +++ b/third_party/xla/xla/pjrt/cpu/cpu_device.h @@ -27,19 +27,16 @@ limitations under the License. #include "absl/strings/string_view.h" #include "absl/types/span.h" #include "xla/literal.h" -#include "xla/pjrt/cpu/cpu_async_execution_tracker.h" -#include "xla/pjrt/cpu/execution_stream_event_map.h" +#include "xla/pjrt/common_pjrt_client.h" #include "xla/pjrt/pjrt_client.h" #include "xla/pjrt/pjrt_common.h" #include "xla/pjrt/plugin/xla_cpu/cpu_device_description.h" -#include "xla/pjrt/semaphore.h" namespace xla { class PjRtCpuDevice final : public PjRtDevice { public: - explicit PjRtCpuDevice(int process_id, int local_device_id, - int max_inflight_computations = 32); + explicit PjRtCpuDevice(int process_id, int local_device_id); const CpuDeviceDescription& description() const override { return description_; @@ -79,11 +76,6 @@ class PjRtCpuDevice final : public PjRtDevice { absl::StatusOr memory_space_by_kind_id(int id) const; - // Returns a semaphore for admission control on inflight computations. - Semaphore& max_inflight_computations_semaphore() { - return max_inflight_computations_semaphore_; - } - std::unique_ptr CreateAsyncTrackingEvent( absl::string_view description) const override { return nullptr; @@ -92,28 +84,11 @@ class PjRtCpuDevice final : public PjRtDevice { absl::StatusOr PoisonExecution(int32_t launch_id, absl::Status error) override; - CpuAsyncExecutionTracker* async_execution_tracker() { - return async_execution_tracker_.get(); - } - - ExecutionStreamEventMap* stream_event_map() const { - return stream_event_map_.get(); - } - private: PjRtClient* client_ = nullptr; CpuDeviceDescription description_; absl::InlinedVector memory_spaces_; absl::flat_hash_map memory_spaces_by_id_; - - // TODO(zhangqiaorjc): Optimize semaphore related overhead. - // Semaphore used to limit how many programs can be enqueued by the host - // ahead of the device. - Semaphore max_inflight_computations_semaphore_; - - std::unique_ptr async_execution_tracker_; - - std::unique_ptr stream_event_map_; }; } // namespace xla diff --git a/third_party/xla/xla/pjrt/interpreter/interpreter_executable.cc b/third_party/xla/xla/pjrt/interpreter/interpreter_executable.cc index 8cae4d04e3d132..d9d8d197bb9aa4 100644 --- a/third_party/xla/xla/pjrt/interpreter/interpreter_executable.cc +++ b/third_party/xla/xla/pjrt/interpreter/interpreter_executable.cc @@ -207,7 +207,7 @@ absl::StatusOr> RunInterpreterHloPasses( std::unique_ptr hlo_module) { HloPassPipeline pipeline("Interpreter"); - // The TopkDecomposer generates a compare op with type=TOTALORDER and must + // The TopkDecomposer generates a compare op with order=TOTAL and must // run before the ComparisonExpander which rewrites such comparisons. pipeline.AddPass(); pipeline.AddPass(); diff --git a/third_party/xla/xla/pjrt/raw_pjrt_client.h b/third_party/xla/xla/pjrt/raw_pjrt_client.h index 0bce5ae8f8a7ef..e0804e2f9953f6 100644 --- a/third_party/xla/xla/pjrt/raw_pjrt_client.h +++ b/third_party/xla/xla/pjrt/raw_pjrt_client.h @@ -208,6 +208,22 @@ class PjRtRawClient { LOG(FATAL) << "Implement MakeLoadState()"; } + virtual absl::StatusOr PoisonExecution(LocalDeviceId local_device_id, + int32_t launch_id, + absl::Status error) { + return absl::UnimplementedError("PoisonExecution is not supported"); + } + + virtual absl::Status TransferToInfeed(LocalDeviceId local_device_id, + const LiteralSlice& literal) { + return absl::UnimplementedError("TransferToInfeed is not supported"); + } + + virtual absl::Status TransferFromOutfeed(LocalDeviceId local_device_id, + MutableBorrowingLiteral literal) { + return absl::UnimplementedError("TransferToOutfeed is not supported"); + } + virtual void ScheduleRemoteSend(PjRtMemorySpace* memory_space, PjRtRawBufferRef raw_buffer, PjRtDeviceEventRefVector definition_events, diff --git a/third_party/xla/xla/pjrt/se/buffer_sequencing_event.cc b/third_party/xla/xla/pjrt/se/buffer_sequencing_event.cc index 0bc24ee1bfbaa0..fb5af6bba7fceb 100644 --- a/third_party/xla/xla/pjrt/se/buffer_sequencing_event.cc +++ b/third_party/xla/xla/pjrt/se/buffer_sequencing_event.cc @@ -40,10 +40,7 @@ namespace xla { void BufferSequencingEvent::SetSequencingEvent(EventPool::Handle event, se::Stream* stream) { - EventState state; - state.event = std::move(event); - state.definition_stream = stream; - event_.emplace(std::move(state)); + event_.emplace(EventState{std::move(event), stream}); } void BufferSequencingEvent::SetDefinedStatus(absl::Status status) { diff --git a/third_party/xla/xla/pjrt/se/event_pool.cc b/third_party/xla/xla/pjrt/se/event_pool.cc index 67ea04b501a6f0..948989c41bf272 100644 --- a/third_party/xla/xla/pjrt/se/event_pool.cc +++ b/third_party/xla/xla/pjrt/se/event_pool.cc @@ -29,6 +29,8 @@ limitations under the License. namespace xla { EventPool::Handle::~Handle() { + // event_ == nullptr is ONLY possible when Handle is in an invalid + // 'moved from' state. if (pool_ && event_) { absl::MutexLock lock(pool_->mu_free_events_); pool_->free_events_.push(std::move(event_)); @@ -40,22 +42,23 @@ EventPool::EventPool(bool allow_reuse) absl::StatusOr EventPool::AllocateEvent( AsyncWorkRunner* async_work_runner, se::StreamExecutor* executor) { - Handle event; + std::unique_ptr event; + EventPool* pool = nullptr; if (allow_reuse_) { - event.pool_ = this; + pool = this; absl::MutexLock lock(mu_free_events_); if (!free_events_.empty()) { - event.event_ = std::move(free_events_.top()); + event = std::move(free_events_.top()); free_events_.pop(); } } - if (!event.event_) { - ABSL_ASSIGN_OR_RETURN(event.event_, NeverRunOnFiber(async_work_runner, [&]() { + if (!event) { + ABSL_ASSIGN_OR_RETURN(event, NeverRunOnFiber(async_work_runner, [&]() { return executor->CreateEvent(); })); } - return event; + return EventPool::Handle(pool, std::move(event)); } void EventPool::ThenRecordEvent(se::Stream* stream, EventPool::Handle& handle) { diff --git a/third_party/xla/xla/pjrt/se/event_pool.h b/third_party/xla/xla/pjrt/se/event_pool.h index 7a746d65525eb8..cb9e7c431674d0 100644 --- a/third_party/xla/xla/pjrt/se/event_pool.h +++ b/third_party/xla/xla/pjrt/se/event_pool.h @@ -35,7 +35,6 @@ class EventPool { public: class Handle { public: - Handle() = default; ~Handle(); Handle(const Handle&) = delete; @@ -60,6 +59,8 @@ class EventPool { private: friend class EventPool; + Handle(EventPool* pool, std::unique_ptr event) + : pool_(pool), event_(std::move(event)) {} EventPool* pool_ = nullptr; std::unique_ptr event_; diff --git a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client_test.cc b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client_test.cc index 7449168d3d26a6..9d08e8d509e83c 100644 --- a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client_test.cc +++ b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client_test.cc @@ -512,7 +512,7 @@ TEST(PjRtStreamExecutorClientTest, MakeAllocationReadyEventAsync) { data.data(), S32, {1024}, /*byte_strides=*/std::nullopt, PjRtClient::HostBufferSemantics::kImmutableZeroCopy, nullptr, memory_space, /*device_layout=*/nullptr)); - + TF_ASSERT_OK(buffer->GetReadyFuture().Await()); Shape shape = buffer->on_device_shape(); TF_ASSERT_OK_AND_ASSIGN(auto result, client->CreateAliasBuffer(shape, memory_space)); @@ -671,6 +671,9 @@ TEST(PjRtStreamExecutorClientTest, CrossHostSendBuffersCleanupAfterFailure) { /*memory_space=*/memory_space, /*device_layout=*/nullptr)); + TF_ASSERT_OK(buffer0->GetReadyFuture().Await()); + TF_ASSERT_OK(buffer1->GetReadyFuture().Await()); + // Delete buffer1 so that AcquireScopedRawBuffer fails on it mid-loop in // CrossHostSendBuffers. buffer1->Delete(); diff --git a/third_party/xla/xla/pjrt/transpose.cc b/third_party/xla/xla/pjrt/transpose.cc index c0dcb91b3f611a..41afad3c6099ed 100644 --- a/third_party/xla/xla/pjrt/transpose.cc +++ b/third_party/xla/xla/pjrt/transpose.cc @@ -86,6 +86,7 @@ limitations under the License. #include #include "absl/algorithm/container.h" +#include "absl/base/attributes.h" #include "absl/base/optimization.h" #include "absl/container/inlined_vector.h" #include "absl/functional/function_ref.h" @@ -188,6 +189,7 @@ struct TransposePlan::Node { template +ABSL_ATTRIBUTE_FUNC_ALIGN(64) void MacroKernel(const char* __restrict a, int64_t lda, int outer_bs_a, char* __restrict b, int64_t ldb, int outer_bs_b, void* __restrict scratch, int bits_per_element) { @@ -253,6 +255,7 @@ void MacroKernel(const char* __restrict a, int64_t lda, int outer_bs_a, // following by iterating over the linked Node data structure. template +ABSL_ATTRIBUTE_FUNC_ALIGN(64) void Transpose(const char* __restrict a, int outer_bs_a, char* __restrict b, int outer_bs_b, TransposePlan::Node const* __restrict node, void* __restrict scratch, int bits_per_element) { @@ -462,6 +465,7 @@ void TransposeConstStride1(const char* __restrict a, char* __restrict b, } template +ABSL_ATTRIBUTE_FUNC_ALIGN(64) void TransposePlan::ExecuteTyped(const char* a, char* b, absl::Span nodes, int bits_per_element) const { @@ -516,6 +520,7 @@ struct uint128 { }; static_assert(sizeof(uint128) == 16, "uint128 should be 16 bytes in size"); +ABSL_ATTRIBUTE_FUNC_ALIGN(64) void TransposePlan::ExecuteChunk(int chunk_id, const void* a, void* b, bool input_is_global, bool output_is_global) const { @@ -619,6 +624,7 @@ void TransposePlan::ExecuteInternal( } } +ABSL_ATTRIBUTE_FUNC_ALIGN(64) void TransposePlan::Execute( const void* a, void* b, std::optional)>> schedule_work) diff --git a/third_party/xla/xla/python/ifrt/BUILD b/third_party/xla/xla/python/ifrt/BUILD index 23bc8a98a5d733..43c2f57771d894 100644 --- a/third_party/xla/xla/python/ifrt/BUILD +++ b/third_party/xla/xla/python/ifrt/BUILD @@ -1225,8 +1225,6 @@ xla_cc_test( ":sharding_serdes", "//xla:shape_util", "//xla/pjrt:pjrt_layout", - "//xla/tsl/lib/core:status_test_util", - "//xla/tsl/platform:statusor", "//xla/tsl/platform:test", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/container:flat_hash_map", diff --git a/third_party/xla/xla/python/ifrt/array_impl_test_lib.cc b/third_party/xla/xla/python/ifrt/array_impl_test_lib.cc index fb30db8b988f90..4cacd83c869b6a 100644 --- a/third_party/xla/xla/python/ifrt/array_impl_test_lib.cc +++ b/third_party/xla/xla/python/ifrt/array_impl_test_lib.cc @@ -697,12 +697,12 @@ TEST(ArrayImplTest, CopyArraysToHostBufferShardsSingleDevice) { copy_specs.push_back({ /*array=*/arrays[0], /*buffers=*/ - {{{0}, {out0.data(), dtype, shape, /*byte_strides=*/std::nullopt}}}, + {{out0.data(), dtype, shape, /*byte_strides=*/std::nullopt}}, }); copy_specs.push_back({ /*array=*/arrays[1], /*buffers=*/ - {{{0}, {out1.data(), dtype, shape, /*byte_strides=*/std::nullopt}}}, + {{out1.data(), dtype, shape, /*byte_strides=*/std::nullopt}}, }); ASSERT_OK_AND_ASSIGN( @@ -735,9 +735,12 @@ TEST(ArrayImplTest, CopyArraysToHostBufferShardsMultiDevice) { client->addressable_devices().subspan(0, 2); ASSERT_OK_AND_ASSIGN(DeviceListRef device_list, client->MakeDeviceList(devices)); - ShardingRef sharding = - ConcreteEvenSharding::Create(device_list, MemoryKind(), shape, - shard_shape, /*is_fully_replicated=*/false); + std::vector shard_shapes = {shard_shape, shard_shape}; + std::vector index_domains = { + IndexDomain(Index({0, 0}), shard_shape), + IndexDomain(Index({1, 0}), shard_shape)}; + ShardingRef sharding = ConcreteSharding::Create( + device_list, MemoryKind(), shape, shard_shapes, index_domains); std::vector make_specs; make_specs.push_back({ @@ -765,12 +768,8 @@ TEST(ArrayImplTest, CopyArraysToHostBufferShardsMultiDevice) { copy_specs.push_back({ /*array=*/arrays[0], /*buffers=*/ - {{{0}, - {out_shard0.data(), dtype, shard_shape, - /*byte_strides=*/std::nullopt}}, - {{1}, - {out_shard1.data(), dtype, shard_shape, - /*byte_strides=*/std::nullopt}}}, + {{out_shard0.data(), dtype, shard_shape, /*byte_strides=*/std::nullopt}, + {out_shard1.data(), dtype, shard_shape, /*byte_strides=*/std::nullopt}}, }); ASSERT_OK_AND_ASSIGN( @@ -826,7 +825,7 @@ TEST(ArrayImplTest, CopyArraysToHostBufferShardsReplicated) { copy_specs.push_back({ /*array=*/arrays[0], /*buffers=*/ - {{{0, 1}, {out.data(), dtype, shape, /*byte_strides=*/std::nullopt}}}, + {{out.data(), dtype, shape, /*byte_strides=*/std::nullopt}}, }); ASSERT_OK_AND_ASSIGN( diff --git a/third_party/xla/xla/python/ifrt/client.h b/third_party/xla/xla/python/ifrt/client.h index 7bd41fbcdcc1e0..3cce987aeccb1c 100644 --- a/third_party/xla/xla/python/ifrt/client.h +++ b/third_party/xla/xla/python/ifrt/client.h @@ -214,17 +214,11 @@ class Client : public RTTIExtends { // Represents the specification of copying an array to host buffer shards. // - // `buffers` is a list of pairs of addressable shard indices and a destination - // mutable host buffer. - // - // For replicated or partially-replicated arrays, multiple shard indices that - // hold the same shard data can be passed in `ShardIndices`. The runtime - // implementation can choose which shard index (or combination of shard - // indices) to copy the data to the host buffer from. + // `buffers` is a list of destination host buffers that have one-to-one + // correspondence to the unique index domains in + // `Sharding::UniqueIndexDomains()`. struct CopyArraysToHostBufferShardsSpec { - using ShardIndices = absl::InlinedVector; - using Buffers = - absl::InlinedVector, 1>; + using Buffers = absl::InlinedVector; ArrayRef array; Buffers buffers; }; diff --git a/third_party/xla/xla/python/ifrt/client_impl_util.cc b/third_party/xla/xla/python/ifrt/client_impl_util.cc index 4820717f094ef6..03bec7b1b0edfc 100644 --- a/third_party/xla/xla/python/ifrt/client_impl_util.cc +++ b/third_party/xla/xla/python/ifrt/client_impl_util.cc @@ -236,10 +236,6 @@ absl::StatusOr>> ClientCopyArraysToHostBufferShards( std::vector> result; result.reserve(specs.size()); for (Client::CopyArraysToHostBufferShardsSpec& spec : specs) { - if (spec.buffers.empty()) { // Nothing to copy. - result.push_back(absl::OkStatus()); - continue; - } if (spec.array == nullptr) { return absl::InvalidArgumentError( "CopyArraysToHostBufferShards called with a null array."); @@ -249,6 +245,21 @@ absl::StatusOr>> ClientCopyArraysToHostBufferShards( "CopyArraysToHostBufferShards called with an array with some " "non-addressable devices."); } + using UniqueIndexDomains = + absl::InlinedVector; + ABSL_ASSIGN_OR_RETURN( + UniqueIndexDomains unique_index_domains, + spec.array->sharding().UniqueIndexDomains(spec.array->shape())); + if (spec.buffers.size() != unique_index_domains.size()) { + return absl::InvalidArgumentError(absl::StrCat( + "The number of buffers (", spec.buffers.size(), + ") does not match the number of unique index domains (", + unique_index_domains.size(), ") in CopyArraysToHostBufferShards.")); + } + if (spec.buffers.empty()) { // Nothing to copy. + result.push_back(absl::OkStatus()); + continue; + } // Split the array into single-device arrays. ABSL_ASSIGN_OR_RETURN(std::vector single_device_arrays, spec.array->DisassembleIntoSingleDeviceArrays( @@ -256,18 +267,18 @@ absl::StatusOr>> ClientCopyArraysToHostBufferShards( absl::InlinedVector, 4> buffer_futures; buffer_futures.reserve(spec.buffers.size()); - for (auto& buffer : spec.buffers) { - Client::CopyArraysToHostBufferShardsSpec::ShardIndices& shard_indices = - buffer.first; - Client::MutableHostBuffer& host_buffer = buffer.second; + for (int i = 0; i < spec.buffers.size(); ++i) { + Client::MutableHostBuffer& host_buffer = spec.buffers[i]; + absl::Span shard_indices = + unique_index_domains[i].shard_indices; if (shard_indices.empty()) { - return absl::InvalidArgumentError( - "No source shard indices specified for a host buffer in " + return absl::InternalError( + "No source shard indices found for a unique index domain in " "CopyArraysToHostBufferShards."); } - // If multiple array source shards are specified, pick the first one to + // If multiple array source shards are available, pick the first one to // copy from. - const int64_t shard_idx = shard_indices.front(); + const int shard_idx = shard_indices.front(); if (shard_idx < 0 || shard_idx >= single_device_arrays.size()) { return absl::OutOfRangeError( absl::StrCat("Shard index ", shard_idx, " out of range [0, ", diff --git a/third_party/xla/xla/python/ifrt/ir/BUILD b/third_party/xla/xla/python/ifrt/ir/BUILD index 76e3e968b407c0..63f62e8bb1e89c 100644 --- a/third_party/xla/xla/python/ifrt/ir/BUILD +++ b/third_party/xla/xla/python/ifrt/ir/BUILD @@ -574,6 +574,7 @@ cc_library( ":program_interpreter", ":utils", "//xla:xla_data_proto_cc", + "//xla/pjrt:host_memory_spaces", "//xla/pjrt:pjrt_layout", "//xla/python/ifrt", "//xla/python/ifrt/ir/transforms:debug", @@ -722,6 +723,7 @@ cc_library( ":version", "//xla/pjrt:pjrt_compiler", "//xla/pjrt:pjrt_executable", + "//xla/pjrt:pjrt_layout", "//xla/python/ifrt", "//xla/python/ifrt:serdes", "//xla/python/ifrt:serdes_proto_cc", @@ -769,6 +771,7 @@ cc_library( "//xla/tsl/concurrency:future", "//xla/tsl/concurrency:ref_count", "//xla/tsl/platform:errors", + "@com_google_absl//absl/container:btree", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/functional:any_invocable", diff --git a/third_party/xla/xla/python/ifrt/ir/compiled_ifrt_ir_program.cc b/third_party/xla/xla/python/ifrt/ir/compiled_ifrt_ir_program.cc index 0a3bb477f5ae07..7bcebf6211ec88 100644 --- a/third_party/xla/xla/python/ifrt/ir/compiled_ifrt_ir_program.cc +++ b/third_party/xla/xla/python/ifrt/ir/compiled_ifrt_ir_program.cc @@ -41,6 +41,7 @@ limitations under the License. #include "mlir/Pass/PassManager.h" #include "mlir/Support/LLVM.h" #include "mlir/Support/LogicalResult.h" +#include "xla/pjrt/host_memory_spaces.h" #include "xla/pjrt/pjrt_layout.h" #include "xla/python/ifrt/array_spec.h" #include "xla/python/ifrt/client.h" @@ -50,6 +51,7 @@ limitations under the License. #include "xla/python/ifrt/executable.h" #include "xla/python/ifrt/ir/atom_program_compiler.h" #include "xla/python/ifrt/ir/constants.h" +#include "xla/python/ifrt/ir/ifrt_dialect.h" #include "xla/python/ifrt/ir/ifrt_ir_program.h" #include "xla/python/ifrt/ir/ifrt_ops.h" #include "xla/python/ifrt/ir/program_interpreter.h" @@ -89,7 +91,7 @@ class FutureExecutor : public tsl::Executor { absl::StatusOr> BuildDefaultLayout( const ArraySpec& arg_spec, Client* client) { - ABSL_ASSIGN_OR_RETURN(auto shard_shape, + ABSL_ASSIGN_OR_RETURN(xla::ifrt::Shape shard_shape, arg_spec.sharding->GetShardShape(arg_spec.shape)); return client->GetDefaultPjRtLayout( arg_spec.dtype, shard_shape.dims(), @@ -109,18 +111,42 @@ GetParameterLayoutFromLoadedExecutable( auto atom_program_name = loaded_exec_op.getSymName().str(); auto exec_it = atom_program_executables.find(atom_program_name); if (exec_it != atom_program_executables.end()) { - ABSL_ASSIGN_OR_RETURN(auto exec_layouts, exec_it->second->GetParameterLayouts()); + ABSL_ASSIGN_OR_RETURN( + std::vector> exec_layouts, + exec_it->second->GetParameterLayouts()); return std::move(exec_layouts[param_operand_number]); } return absl::FailedPreconditionError( absl::StrFormat("Could not find SPMD executable %s", atom_program_name)); } +bool CanPropagateLayoutAcrossCopyArrays(IfrtArrayType input_type, + IfrtArrayType output_type, + const DeviceListRef& device_list) { + llvm::ArrayRef src_device_ids = input_type.getDevices(); + llvm::ArrayRef dst_device_ids = output_type.getDevices(); + CHECK(!src_device_ids.empty()); + CHECK(!dst_device_ids.empty()); + CHECK_LT(src_device_ids.front(), device_list->devices().size()); + CHECK_LT(dst_device_ids.front(), device_list->devices().size()); + Device* src_device = device_list->devices()[src_device_ids.front()]; + Device* dst_device = device_list->devices()[dst_device_ids.front()]; + if (src_device->PlatformName() != dst_device->PlatformName()) { + return false; + } + if (input_type.MemoryKind().value() == xla::UnpinnedHostMemorySpace::kKind || + output_type.MemoryKind().value() == xla::UnpinnedHostMemorySpace::kKind) { + return false; + } + return true; +} + absl::StatusOr> GetLayoutForValue( mlir::Value value, Client* client, const AtomExecutableMap& atom_program_executables, absl::Span in_specs, - mlir::SymbolTableCollection& symbol_table) { + mlir::SymbolTableCollection& symbol_table, + const DeviceListRef& device_list) { if (auto block_arg = llvm::dyn_cast(value)) { if (in_specs[block_arg.getArgNumber()].layout != nullptr) { return in_specs[block_arg.getArgNumber()].layout; @@ -138,15 +164,23 @@ absl::StatusOr> GetLayoutForValue( return absl::FailedPreconditionError(absl::StrFormat( "Could not find SPMD executable %s", atom_program_name)); } - ABSL_ASSIGN_OR_RETURN(auto exec_layouts, exec_it->second->GetOutputLayouts()); + ABSL_ASSIGN_OR_RETURN( + std::vector> exec_layouts, + exec_it->second->GetOutputLayouts()); return exec_layouts[op_result.getResultNumber()]; } if (auto copy_arrays = llvm::dyn_cast(op_result.getOwner())) { - return GetLayoutForValue( - copy_arrays.getInputs()[op_result.getResultNumber()], client, - atom_program_executables, in_specs, symbol_table); + mlir::Value input = copy_arrays.getInputs()[op_result.getResultNumber()]; + IfrtArrayType input_type = GetArrayType(input); + IfrtArrayType output_type = GetArrayType(op_result); + if (!CanPropagateLayoutAcrossCopyArrays(input_type, output_type, + device_list)) { + return nullptr; + } + return GetLayoutForValue(input, client, atom_program_executables, in_specs, + symbol_table, device_list); } return absl::FailedPreconditionError(absl::StrFormat( @@ -158,7 +192,8 @@ absl::StatusOr> GetLayoutForValue( absl::Status PopulateLayouts(mlir::ModuleOp mlir_module, Client* client, const AtomExecutableMap& atom_program_executables, absl::Span in_specs, - absl::Span out_specs) { + absl::Span out_specs, + const DeviceListRef& device_list) { tsl::profiler::TraceMe traceme("PopulateLayouts"); auto main_func = GetMainFunction(mlir_module); @@ -223,10 +258,10 @@ absl::Status PopulateLayouts(mlir::ModuleOp mlir_module, Client* client, for (mlir::OpOperand& return_operand : main_func.front().getTerminator()->getOpOperands()) { auto& out_spec = out_specs[return_operand.getOperandNumber()]; - ABSL_ASSIGN_OR_RETURN( - out_spec.layout, - GetLayoutForValue(return_operand.get(), client, - atom_program_executables, in_specs, symbol_table)); + ABSL_ASSIGN_OR_RETURN(out_spec.layout, + GetLayoutForValue(return_operand.get(), client, + atom_program_executables, in_specs, + symbol_table, device_list)); if (!out_spec.layout) { ABSL_ASSIGN_OR_RETURN(out_spec.layout, BuildDefaultLayout(out_spec, client)); } @@ -362,7 +397,7 @@ CompiledIfrtIrProgram::Create( absl::Status layout_status = PopulateLayouts( ifrt_ir_program->mlir_module, client, *atom_executable_map, - absl::MakeSpan(in_specs), absl::MakeSpan(out_specs)); + absl::MakeSpan(in_specs), absl::MakeSpan(out_specs), device_list); if (!layout_status.ok()) { for (auto& spec : in_specs) { spec.layout = nullptr; diff --git a/third_party/xla/xla/python/ifrt/ir/ifrt_ir_loaded_executable_test_lib.cc b/third_party/xla/xla/python/ifrt/ir/ifrt_ir_loaded_executable_test_lib.cc index 7730983c23967a..c99fb610d2a768 100644 --- a/third_party/xla/xla/python/ifrt/ir/ifrt_ir_loaded_executable_test_lib.cc +++ b/third_party/xla/xla/python/ifrt/ir/ifrt_ir_loaded_executable_test_lib.cc @@ -36,6 +36,7 @@ limitations under the License. #include "mlir/IR/OwningOpRef.h" #include "xla/pjrt/pjrt_compiler.h" #include "xla/pjrt/pjrt_executable.h" +#include "xla/pjrt/pjrt_layout.h" #include "xla/python/ifrt/array.h" #include "xla/python/ifrt/bundle.h" #include "xla/python/ifrt/device.h" @@ -863,6 +864,48 @@ module { std::move(device_list1))); } +TEST_F(IfrtIrLoadedExecutableTest, RemapSliceArray) { + if (GetNumDevices() < 2) { + GTEST_SKIP() << "Insufficient devices to run this test."; + } + std::string source = R"( +!array = !ifrt.array, + #ifrt.sharding_param<2x1 to [0] on 2>, [0,1]> +!array0 = !ifrt.array, + #ifrt.sharding_param<1x1 to [0] on 1>, [1]> +module { + func.func @main(%arg0: !array) -> !array0 + attributes {ifrt.function} { + %0, %ctrl_0 = ifrt.RemapArrays(%arg0) + mappings=[#ifrt.array_mapping<0, 0, [#ifrt.mapping<[1:2:1] to [0:1:1]>]>] + : (!array) -> !array0 + return %0 : !array0 + } +} + )"; + ASSERT_OK_AND_ASSIGN(DeviceListRef devices, PickDevices(2)); + ASSERT_OK_AND_ASSIGN(LoadedExecutableRef loaded_exec, + CompileProgram(source, devices)); + + std::vector data_shard0 = {0, 1}; + std::vector data_shard1 = {2, 3}; + DType dtype(DType::kS32); + Shape shard_shape({1, 2}); + ASSERT_OK_AND_ASSIGN(ArrayRef input, + CreateArray({data_shard0.data(), data_shard1.data()}, + Shape({2, 2}), shard_shape, dtype, devices)); + + ASSERT_OK_AND_ASSIGN( + LoadedExecutable::ExecuteResult result, + Execute(loaded_exec, absl::MakeSpan(&input, 1), devices)); + ASSERT_OK(result.status.Await()); + ASSERT_EQ(result.outputs.size(), 1); + ASSERT_OK_AND_ASSIGN(DeviceListRef out_devices, + client_->MakeDeviceList({devices->devices()[1]})); + ASSERT_NO_FATAL_FAILURE(AssertPerShardData( + result.outputs[0], dtype, shard_shape, {{2, 3}}, std::move(out_devices))); +} + TEST_F(IfrtIrLoadedExecutableTest, LoadedExecBinding) { ASSERT_OK_AND_ASSIGN(DeviceListRef devices, PickDevices(2)); std::string mhlo_source = R"( @@ -1263,6 +1306,16 @@ module { std::vector data = {0, 1}; DType dtype(DType::kS32); Shape shape({2, 1}); + ASSERT_OK_AND_ASSIGN( + std::vector> output_layouts, + ifrt_ir_executable->GetOutputLayouts()); + ASSERT_EQ(output_layouts.size(), 1); + ASSERT_OK_AND_ASSIGN(std::shared_ptr expected_layout, + client_->GetDefaultPjRtLayout( + dtype, shape.dims(), device->devices().front(), + /*memory_kind=*/MemoryKind())); + ASSERT_EQ(*output_layouts[0], *expected_layout); + ASSERT_OK_AND_ASSIGN(ArrayRef input, CreateArray({data.data()}, shape, /*shard_shape=*/shape, dtype, cpu_device)); diff --git a/third_party/xla/xla/python/ifrt/ir/program_interpreter.cc b/third_party/xla/xla/python/ifrt/ir/program_interpreter.cc index 2bff7853184ec3..6c64c794f7b901 100644 --- a/third_party/xla/xla/python/ifrt/ir/program_interpreter.cc +++ b/third_party/xla/xla/python/ifrt/ir/program_interpreter.cc @@ -22,6 +22,7 @@ limitations under the License. #include #include +#include "absl/container/btree_map.h" #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" #include "absl/functional/any_invocable.h" @@ -552,6 +553,104 @@ absl::StatusOr ProgramInterpreter::HandleOp( namespace { +absl::StatusOr ComputeDeviceListFromIntervals( + Client* client, const DeviceListRef& device_list, int64_t count, + absl::Span intervals) { + TF_RET_CHECK(count >= 0); + std::vector devices; + devices.reserve(count); + for (const IfrtIntervalAttr& interval : intervals) { + TF_RET_CHECK(interval.getStep() > 0); + int64_t index = interval.getStart(); + while (index < interval.getEnd()) { + TF_RET_CHECK(index >= 0 && index < device_list->size()); + devices.push_back(device_list->devices()[index]); + index += interval.getStep(); + } + } + TF_RET_CHECK(devices.size() == count); + return client->MakeDeviceList(devices); +} + +absl::StatusOr< + absl::flat_hash_map>> +ComputeInputDevicesForOutputMap(Client* client, + absl::Span input_specs, + absl::Span output_specs, + RemapArraysOp remap_op) { + // A list of intervals along with the sum of entries across all the intervals. + struct IntervalsAndCount { + std::vector intervals; + int64_t count = 0; + }; + + // Map from output array index to all its input contributors. + // + // The value is a map from input array index to the intervals of that input + // array that contribute to the given output. + // Using btree_map ensures deterministic iteration order across compiler runs. + absl::btree_map> + output_to_inputs_and_intervals; + for (const auto& array_mapping : remap_op.getMappings()) { + const auto array_mapping_attr = + llvm::cast(array_mapping); + int in_array = array_mapping_attr.getInArrayIndex(); + int out_array = array_mapping_attr.getOutArrayIndex(); + TF_RET_CHECK(in_array >= 0 && in_array < input_specs.size()); + TF_RET_CHECK(out_array >= 0 && out_array < output_specs.size()); + + IntervalsAndCount& intervals = + output_to_inputs_and_intervals[out_array][in_array]; + for (const auto& m : array_mapping_attr.getMappings()) { + const auto mapping_attr = llvm::cast(m); + IfrtIntervalAttr from_shards = mapping_attr.getFromShards(); + intervals.intervals.push_back(from_shards); + intervals.count += from_shards.size(); + } + } + + absl::flat_hash_map> + input_devices_for_output_map; + input_devices_for_output_map.reserve(output_specs.size()); + for (int out_array = 0; out_array < output_specs.size(); ++out_array) { + input_devices_for_output_map.insert({out_array, {}}); + } + + for (const auto& [out_array, input_intervals] : + output_to_inputs_and_intervals) { + TF_RET_CHECK(out_array >= 0 && out_array < output_specs.size()); + const DeviceListRef& out_devices = + output_specs[out_array].sharding->devices(); + std::vector& out_input_devices = + input_devices_for_output_map[out_array]; + for (const auto& [in_array, intervals] : input_intervals) { + TF_RET_CHECK(in_array >= 0 && in_array < input_specs.size()); + const DeviceListRef& in_devices = + input_specs[in_array].sharding->devices(); + TF_RET_CHECK(intervals.count >= 0 && + intervals.count <= out_devices->size()); + TF_RET_CHECK(intervals.count >= 0 && + intervals.count <= in_devices->size()); + DeviceListRef interval_device_list; + if (intervals.count == in_devices->size() && + intervals.intervals.size() == 1 && + intervals.intervals[0].getStart() == 0 && + intervals.intervals[0].getStep() == 1) { + interval_device_list = in_devices; + } else { + ABSL_ASSIGN_OR_RETURN( + interval_device_list, + ComputeDeviceListFromIntervals(client, in_devices, intervals.count, + intervals.intervals)); + } + out_input_devices.push_back(RemapPlan::InputDeviceRange{ + /*in_array=*/in_array, + /*input_devices=*/std::move(interval_device_list)}); + } + } + return input_devices_for_output_map; +} + struct RemapArraysOpState { std::string pretty_print; @@ -664,28 +763,6 @@ absl::StatusOr ProgramInterpreter::HandleOp( RemapArraysOpState state; state.pretty_print = PrettyPrint(remap_op); - // Construct the mappings of the remap plan. - std::vector mappings; - mappings.reserve(remap_op.getMappings().size()); - for (const auto& array_mapping : remap_op.getMappings()) { - const auto array_mapping_attr = - llvm::cast(array_mapping); - auto& mapping = mappings.emplace_back(); - mapping.in_array = array_mapping_attr.getInArrayIndex(); - mapping.out_array = array_mapping_attr.getOutArrayIndex(); - mapping.from.reserve(array_mapping_attr.getMappings().size()); - mapping.to.reserve(array_mapping_attr.getMappings().size()); - for (const auto& m : array_mapping_attr.getMappings()) { - const auto mapping_attr = llvm::cast(m); - auto from_shards = mapping_attr.getFromShards(); - auto to_shards = mapping_attr.getToShards(); - mapping.from.push_back(RemapPlan::Interval{ - from_shards.getStart(), from_shards.getEnd(), from_shards.getStep()}); - mapping.to.push_back(RemapPlan::Interval{ - to_shards.getStart(), to_shards.getEnd(), to_shards.getStep()}); - } - }; - // Get the input specs of the remap plan and the input arrays. std::vector input_specs; input_specs.reserve(remap_op.getInputs().size()); @@ -708,10 +785,13 @@ absl::StatusOr ProgramInterpreter::HandleOp( output_specs.push_back(std::move(spec)); } - ABSL_ASSIGN_OR_RETURN( - state.remap_plan, - RemapPlan::CreateOptimized(client_, std::move(input_specs), - std::move(output_specs), std::move(mappings))); + ABSL_ASSIGN_OR_RETURN(auto input_devices_for_output_map, + ComputeInputDevicesForOutputMap(client_, input_specs, + output_specs, remap_op)); + + state.remap_plan = RemapPlan(std::move(input_specs), std::move(output_specs), + std::move(input_devices_for_output_map)); + ABSL_RETURN_IF_ERROR(state.remap_plan.Validate()); state.remap_is_donated = remap_op.getDonated(); for (const auto output : remap_op.getOutputs()) { diff --git a/third_party/xla/xla/python/ifrt/remap_plan.cc b/third_party/xla/xla/python/ifrt/remap_plan.cc index 72bcf73aae949f..af2bce28090efc 100644 --- a/third_party/xla/xla/python/ifrt/remap_plan.cc +++ b/third_party/xla/xla/python/ifrt/remap_plan.cc @@ -20,9 +20,9 @@ limitations under the License. #include #include #include -#include #include +#include "absl/base/call_once.h" #include "absl/base/optimization.h" #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" @@ -33,6 +33,7 @@ limitations under the License. #include "absl/status/status_macros.h" #include "absl/status/statusor.h" #include "absl/strings/str_cat.h" +#include "absl/strings/str_format.h" #include "absl/strings/str_join.h" #include "absl/types/span.h" #include "xla/pjrt/pjrt_layout.h" @@ -46,7 +47,6 @@ limitations under the License. #include "xla/python/ifrt/shape.h" #include "xla/python/ifrt/sharding.h" #include "xla/status_macros.h" -#include "xla/util.h" namespace xla { namespace ifrt { @@ -134,16 +134,17 @@ void InputDeviceToOutputToProto( absl::Status CheckRange(int64_t num_shards, const RemapPlan::Interval& interval) { if (interval.start < 0 || interval.start > num_shards - 1) { - return InvalidArgument("start must be in [0, %d], but is %d", - num_shards - 1, interval.start); + return absl::InvalidArgumentError(absl::StrFormat( + "start must be in [0, %d], but is %d", num_shards - 1, interval.start)); } if (interval.step <= 0) { - return InvalidArgument("step must be positive, but is %d", interval.step); + return absl::InvalidArgumentError( + absl::StrFormat("step must be positive, but is %d", interval.step)); } if (interval.end < 0 || interval.end > num_shards + interval.step - 1) { - return InvalidArgument("end must be in [0, %d] if step is %d, but is %d", - num_shards + interval.step - 1, interval.step, - interval.end); + return absl::InvalidArgumentError(absl::StrFormat( + "end must be in [0, %d] if step is %d, but is %d", + num_shards + interval.step - 1, interval.step, interval.end)); } // The `end` bound above is necessary but not sufficient: with a large `step`, // the last stepped index can exceed `num_shards` while still satisfying @@ -153,9 +154,9 @@ absl::Status CheckRange(int64_t num_shards, const int64_t last_index = interval.end - 1 - (interval.end - 1 - interval.start) % interval.step; if (last_index >= num_shards) { - return InvalidArgument( + return absl::InvalidArgumentError(absl::StrFormat( "interval addresses shard %d, which is out of range [0, %d)", - last_index, num_shards); + last_index, num_shards)); } } return absl::OkStatus(); @@ -207,7 +208,8 @@ absl::StatusOr ComputeDeviceListFromIntervals( devices.reserve(count); for (const RemapPlan::Interval& interval : intervals) { if (interval.step <= 0) { - return InvalidArgument("step must be positive, but is %d", interval.step); + return absl::InvalidArgumentError( + absl::StrFormat("step must be positive, but is %d", interval.step)); } int64_t index = interval.start; while (index < interval.end) { @@ -264,20 +266,20 @@ ComputeInputDevicesForOutputMap(Client* client, for (int64_t i = 0; i < mappings.size(); ++i) { const RemapPlan::Mapping& mapping = mappings[i]; if (mapping.in_array < 0 || mapping.in_array >= input_specs.size()) { - return InvalidArgument( - "mappings[%d].in_array must be in [0, %d], but is %d", i, - input_specs.size() - 1, mapping.in_array); + return absl::InvalidArgumentError( + absl::StrFormat("mappings[%d].in_array must be in [0, %d], but is %d", + i, input_specs.size() - 1, mapping.in_array)); } if (mapping.out_array < 0 || mapping.out_array >= output_specs.size()) { - return InvalidArgument( + return absl::InvalidArgumentError(absl::StrFormat( "mappings[%d].out_array must be in [0, %d], but is %d", i, - output_specs.size() - 1, mapping.out_array); + output_specs.size() - 1, mapping.out_array)); } if (mapping.from.size() != mapping.to.size()) { - return InvalidArgument( + return absl::InvalidArgumentError(absl::StrFormat( "mappings[%d].from and mappings[%d].to must have the same number of " "intervals, but has %d and %d intervals", - i, i, mapping.from.size(), mapping.to.size()); + i, i, mapping.from.size(), mapping.to.size())); } const int64_t in_shards_count = input_specs[mapping.in_array] .sharding->devices() @@ -300,12 +302,14 @@ ComputeInputDevicesForOutputMap(Client* client, absl::flat_hash_map> input_devices_for_output_map; for (const auto& [out_array, input_intervals] : + // NOLINTNEXTLINE(*-custom-deterministic-iteration-order) output_to_inputs_and_intervals) { TF_RET_CHECK(out_array >= 0 && out_array < output_specs.size()); const DeviceListRef& out_devices = output_specs[out_array].sharding->devices(); auto [it, inserted] = input_devices_for_output_map.insert({out_array, {}}); TF_RET_CHECK(inserted); + // NOLINTNEXTLINE(*-custom-deterministic-iteration-order) for (const auto& [in_array, intervals] : input_intervals) { TF_RET_CHECK(in_array >= 0 && in_array < input_specs.size()); const DeviceListRef& in_devices = @@ -347,62 +351,147 @@ absl::StatusOr RemapPlan::CreateOptimized( namespace { -// A utility class that calculates the shard shape from an array spec. -class ShardShapeVector { - public: - static absl::StatusOr Create(const ArraySpec& spec) { - // Fast path for even shardings. - if (absl::StatusOr s = spec.sharding->GetShardShape(spec.shape); - s.ok()) { - return ShardShapeVector(*std::move(s)); - } - - ABSL_ASSIGN_OR_RETURN(auto shards, - spec.sharding->Disassemble( - spec.shape, SingleDeviceShardSemantics::kAllShards)); - std::vector shapes; - shapes.reserve(shards.size()); - for (auto& shard : shards) { - shapes.push_back(std::move(shard.first)); - } - return ShardShapeVector(std::move(shapes)); +// Validates array-level consistency between an input array spec and an output +// array spec. +absl::Status CheckArraySpecConsistency(int in_array, const ArraySpec& in_spec, + int out_array, + const ArraySpec& out_spec) { + if (in_spec.dtype != out_spec.dtype) { + return absl::InvalidArgumentError(absl::StrFormat( + "Input and output must have the same dtype: %v (input %d) vs. %v " + "(output %d)", + in_spec.dtype, in_array, out_spec.dtype, out_array)); + } + + if (in_spec.sharding == nullptr) { + return absl::InvalidArgumentError( + absl::StrFormat("Input array %d has null sharding", in_array)); + } + if (out_spec.sharding == nullptr) { + return absl::InvalidArgumentError( + absl::StrFormat("Output array %d has null sharding", out_array)); + } + + ABSL_ASSIGN_OR_RETURN(const Shape in_shard_shape, + in_spec.sharding->GetShardShape(in_spec.shape)); + ABSL_ASSIGN_OR_RETURN(const Shape out_shard_shape, + out_spec.sharding->GetShardShape(out_spec.shape)); + if (in_shard_shape != out_shard_shape) { + return absl::InvalidArgumentError(absl::StrFormat( + "Input and output must have the same shard shape: %v (input %d) vs. %v " + "(output %d)", + in_shard_shape, in_array, out_shard_shape, out_array)); + } + + if (in_spec.sharding->memory_kind() != out_spec.sharding->memory_kind()) { + return absl::InvalidArgumentError(absl::StrFormat( + "Input and output must have the same memory kind: %v (input %d) vs. %v " + "(output %d)", + in_spec.sharding->memory_kind(), in_array, + out_spec.sharding->memory_kind(), out_array)); + } + + const std::shared_ptr& in_layout = in_spec.layout; + const std::shared_ptr& out_layout = out_spec.layout; + if (in_layout != out_layout && + (in_layout == nullptr || out_layout == nullptr || + *in_layout != *out_layout)) { + return absl::InvalidArgumentError(absl::StrFormat( + "Input and output must have the same layout: %s (input %d) vs. %s " + "(output %d)", + in_layout != nullptr ? in_layout->ToString() : "", in_array, + out_layout != nullptr ? out_layout->ToString() : "", + out_array)); } - // Returns the shard shape of `index`-th shard. - const Shape& shard(int index) const { - if (auto* shape = std::get_if(&shapes_)) { - return *shape; - } - if (auto* shapes = std::get_if>(&shapes_)) { - return (*shapes)[index]; - } - LOG(FATAL) << "Unexpected shapes variant: " << shapes_.index(); - } - - private: - explicit ShardShapeVector(Shape shape) : shapes_(std::move(shape)) {} - - explicit ShardShapeVector(std::vector shapes) - : shapes_(std::move(shapes)) {} - - std::variant> shapes_; -}; + return absl::OkStatus(); +} } // namespace -absl::Status RemapPlan::Validate() const { +absl::Status RemapPlan::ValidateArraySpecsUncached() const { const int num_inputs = rep_->input_specs.size(); if (num_inputs == 0) { - return InvalidArgument("Must have at least one input"); + return absl::InvalidArgumentError("Must have at least one input"); } const int num_outputs = rep_->output_specs.size(); + if (num_outputs == 0) { + return absl::InvalidArgumentError("Must have at least one output"); + } + + for (int i = 0; i < num_inputs; ++i) { + if (rep_->input_specs[i].sharding == nullptr) { + return absl::InvalidArgumentError( + absl::StrFormat("Input array %d has null sharding", i)); + } + } + for (int i = 0; i < num_outputs; ++i) { + if (rep_->output_specs[i].sharding == nullptr) { + return absl::InvalidArgumentError( + absl::StrFormat("Output array %d has null sharding", i)); + } + } if (rep_->mappings.empty() && rep_->input_devices_for_output_map.empty()) { - return InvalidArgument( + return absl::InvalidArgumentError( "Must have at least one mapping or input_devices_for_output_map"); } + absl::flat_hash_set> checked_pairs; + for (int64_t i = 0; i < rep_->mappings.size(); ++i) { + const RemapPlan::Mapping& mapping = rep_->mappings[i]; + if (mapping.in_array < 0 || mapping.in_array >= num_inputs) { + return absl::InvalidArgumentError( + absl::StrFormat("mappings[%d].in_array must be in [0, %d], but is %d", + i, num_inputs - 1, mapping.in_array)); + } + if (mapping.out_array < 0 || mapping.out_array >= num_outputs) { + return absl::InvalidArgumentError(absl::StrFormat( + "mappings[%d].out_array must be in [0, %d], but is %d", i, + num_outputs - 1, mapping.out_array)); + } + if (checked_pairs.insert({mapping.in_array, mapping.out_array}).second) { + ABSL_RETURN_IF_ERROR(CheckArraySpecConsistency( + mapping.in_array, rep_->input_specs[mapping.in_array], + mapping.out_array, rep_->output_specs[mapping.out_array])); + } + } + + if (!rep_->input_devices_for_output_map.empty()) { + // NOLINTNEXTLINE(*-custom-deterministic-iteration-order) + for (const auto& [out_array, inputs] : rep_->input_devices_for_output_map) { + if (out_array < 0 || out_array >= num_outputs) { + return absl::InvalidArgumentError(absl::StrFormat( + "Output buffer index %d in `input_devices_for_output_map` is out " + "of range [0, %d]", + out_array, num_outputs - 1)); + } + for (const InputDeviceRange& range : inputs) { + if (range.in_array < 0 || range.in_array >= num_inputs) { + return absl::InvalidArgumentError(absl::StrFormat( + "Input buffer index %d in `input_devices_for_output_map` is out " + "of range [0, %d]", + range.in_array, num_inputs - 1)); + } + if (checked_pairs.insert({range.in_array, out_array}).second) { + ABSL_RETURN_IF_ERROR(CheckArraySpecConsistency( + range.in_array, rep_->input_specs[range.in_array], out_array, + rep_->output_specs[out_array])); + } + } + } + } + + return absl::OkStatus(); +} + +absl::Status RemapPlan::ValidateArrayShardMappingsUncached() const { + const int num_inputs = rep_->input_specs.size(); + const int num_outputs = rep_->output_specs.size(); + TF_RET_CHECK(num_inputs > 0); + TF_RET_CHECK(num_outputs > 0); + std::vector> in_used_buffers_list; std::vector> out_assigned_devices_list; absl::flat_hash_mapmappings.size(); ++i) { const RemapPlan::Mapping& mapping = rep_->mappings[i]; + TF_RET_CHECK(mapping.in_array >= 0 && mapping.in_array < num_inputs); + TF_RET_CHECK(mapping.out_array >= 0 && mapping.out_array < num_outputs); absl::flat_hash_set* in_device_set = rep_->input_devices_for_output_map.contains(mapping.out_array) ? &out_buffer_to_in_buffer_and_devices[mapping.out_array] [mapping.in_array] : nullptr; - if (mapping.in_array < 0 || mapping.in_array >= num_inputs) { - return InvalidArgument( - "mappings[%d].in_array must be in [0, %d], but is %d", i, - num_inputs - 1, mapping.in_array); - } - if (mapping.out_array < 0 || mapping.out_array >= num_outputs) { - return InvalidArgument( - "mappings[%d].out_array must be in [0, %d], but is %d", i, - num_outputs - 1, mapping.out_array); - } if (mapping.from.size() != mapping.to.size()) { - return InvalidArgument( + return absl::InvalidArgumentError(absl::StrFormat( "mappings[%d].from and mappings[%d].to must have the same number " "of intervals, but has %d and %d intervals", - i, i, mapping.from.size(), mapping.to.size()); - } - - const ArraySpec& input_spec = rep_->input_specs[mapping.in_array]; - const ArraySpec& output_spec = rep_->output_specs[mapping.out_array]; - - if (input_spec.dtype != output_spec.dtype) { - return InvalidArgument( - "Input and output must have the same dtype: %v (input %d) vs. %v " - "(output %d)", - input_spec.dtype, mapping.in_array, output_spec.dtype, - mapping.out_array); + i, i, mapping.from.size(), mapping.to.size())); } - const std::shared_ptr& in_layout = - input_spec.layout; - const std::shared_ptr& out_layout = - output_spec.layout; - if (in_layout != out_layout && - (!in_layout || !out_layout || *in_layout != *out_layout)) { - return InvalidArgument( - "Input and output must have the same layout: %s (input %d) vs. %s " - "(output %d)", - in_layout != nullptr ? in_layout->ToString() : "", - mapping.in_array, - out_layout != nullptr ? out_layout->ToString() : "", - mapping.out_array); - } - - ABSL_ASSIGN_OR_RETURN(const auto input_shard_shapes, - ShardShapeVector::Create(input_spec)); - ABSL_ASSIGN_OR_RETURN(const auto output_shard_shapes, - ShardShapeVector::Create(output_spec)); - std::vector& in_used_buffers = in_used_buffers_list[mapping.in_array]; absl::Span in_devices = rep_->input_specs[mapping.in_array] @@ -502,12 +552,12 @@ absl::Status RemapPlan::Validate() const { ABSL_RETURN_IF_ERROR(CheckRange(in_shards_count, in_interval)); ABSL_RETURN_IF_ERROR(CheckRange(out_shards_count, out_interval)); if (GetNumberOfSteps(in_interval) != GetNumberOfSteps(out_interval)) { - return InvalidArgument( + return absl::InvalidArgumentError(absl::StrFormat( "mappings[%d].from[%d] and mappings[%d].to[%d] must have the " "same number of steps, but were %d and %d (%s vs. %s)", i, s, i, s, GetNumberOfSteps(in_interval), GetNumberOfSteps(out_interval), in_interval.DebugString(), - out_interval.DebugString()); + out_interval.DebugString())); } int64_t in_shard = in_interval.start; @@ -516,38 +566,28 @@ absl::Status RemapPlan::Validate() const { TF_RET_CHECK(in_shard >= 0 && in_shard < in_shards_count); TF_RET_CHECK(out_shard >= 0 && out_shard < out_shards_count); if (in_used_buffers[in_shard]) { - return InvalidArgument( + return absl::InvalidArgumentError(absl::StrFormat( "Input array %d addressable shard %d is already used", - mapping.in_array, in_shard); + mapping.in_array, in_shard)); } in_used_buffers[in_shard] = true; if (in_device_set) { if (!in_device_set->insert(in_devices[in_shard]).second) { - return InvalidArgument( + return absl::InvalidArgumentError(absl::StrFormat( "Input device %s used more than once in mappings from input " "array %d to output array %d", in_devices[in_shard]->DebugString(), mapping.in_array, - mapping.out_array); + mapping.out_array)); } } if (out_assigned_devices[out_shard] != nullptr) { - return InvalidArgument( + return absl::InvalidArgumentError(absl::StrFormat( "Output array %d addressable shard %d is already assigned", - mapping.out_array, out_shard); + mapping.out_array, out_shard)); } out_assigned_devices[out_shard] = in_devices[in_shard]; - if (input_shard_shapes.shard(in_shard) != - output_shard_shapes.shard(out_shard)) { - return InvalidArgument( - "Output array %d addressable shard %d has a different shard " - "shape from the corresponding input shard: %v -> %v", - mapping.out_array, out_shard, - input_shard_shapes.shard(in_shard), - output_shard_shapes.shard(out_shard)); - } - in_shard += in_interval.step; out_shard += out_interval.step; } @@ -559,33 +599,30 @@ absl::Status RemapPlan::Validate() const { rep_->output_specs[i].sharding->devices()->AddressableDeviceList(); for (int out_shard = 0; out_shard < devices->size(); ++out_shard) { if (out_assigned_devices_list[i][out_shard] == nullptr) { - return InvalidArgument( + return absl::InvalidArgumentError(absl::StrFormat( "Output array %d addressable shard %d is unassigned", i, - out_shard); + out_shard)); } } if (out_assigned_devices_list[i] != devices->devices()) { - return InvalidArgument( + return absl::InvalidArgumentError(absl::StrFormat( "Output array %d addressable devices and sharding devices do not " "match: Expected %v, but got [%s]", i, *devices, absl::StrJoin(out_assigned_devices_list[i], ", ", [](std::string* s, Device* d) { absl::StrAppend(s, d->ToString()); - })); + }))); } } } if (!rep_->input_devices_for_output_map.empty()) { - if (num_outputs == 0) { - return InvalidArgument("Must have at least one output"); - } if (rep_->input_devices_for_output_map.size() != num_outputs) { - return InvalidArgument( + return absl::InvalidArgumentError(absl::StrFormat( "`input_devices_for_output_map` has %d outputs, but expected %d " "outputs", - rep_->input_devices_for_output_map.size(), num_outputs); + rep_->input_devices_for_output_map.size(), num_outputs)); } std::vector> in_device_sets; in_device_sets.reserve(num_inputs); @@ -597,48 +634,14 @@ absl::Status RemapPlan::Validate() const { } // NOLINTNEXTLINE(*-custom-deterministic-iteration-order) for (const auto& [out_array, inputs] : rep_->input_devices_for_output_map) { - if (out_array < 0 || out_array >= num_outputs) { - return InvalidArgument( - "Output buffer index %d in `input_devices_for_output_map` is out " - "of range [0, %d]", - out_array, num_outputs - 1); - } - const ArraySpec& output_spec = rep_->output_specs[out_array]; + TF_RET_CHECK(out_array >= 0 && out_array < num_outputs); for (const InputDeviceRange& range : inputs) { - if (range.in_array < 0 || range.in_array >= num_inputs) { - return InvalidArgument( - "Input buffer index %d in `input_devices_for_output_map` is out " - "of range [0, %d]", - range.in_array, num_inputs - 1); - } + TF_RET_CHECK(range.in_array >= 0 && range.in_array < num_inputs); if (range.input_devices == nullptr) { - return InvalidArgument( + return absl::InvalidArgumentError(absl::StrFormat( "Output buffer index %d in `input_devices_for_output_map` has " "null input_devices for input array %d", - out_array, range.in_array); - } - const ArraySpec& input_spec = rep_->input_specs[range.in_array]; - - if (input_spec.dtype != output_spec.dtype) { - return InvalidArgument( - "Input and output must have the same dtype: %v (input %d) vs. %v " - "(output %d)", - input_spec.dtype, range.in_array, output_spec.dtype, out_array); - } - - const std::shared_ptr& in_layout = - input_spec.layout; - const std::shared_ptr& out_layout = - output_spec.layout; - if (in_layout != out_layout && - (!in_layout || !out_layout || *in_layout != *out_layout)) { - return InvalidArgument( - "Input and output must have the same layout: %s (input %d) vs. " - "%s (output %d)", - in_layout != nullptr ? in_layout->ToString() : "", - range.in_array, - out_layout != nullptr ? out_layout->ToString() : "", - out_array); + out_array, range.in_array)); } const absl::flat_hash_set& in_device_set = @@ -646,11 +649,11 @@ absl::Status RemapPlan::Validate() const { for (Device* device : range.input_devices->AddressableDeviceList()->devices()) { if (!in_device_set.contains(device)) { - return InvalidArgument( + return absl::InvalidArgumentError(absl::StrFormat( "Output buffer index %d in `input_devices_for_output_map` " "references device %s from input array %d that is not in the " "input array's addressable device list", - out_array, device->DebugString(), range.in_array); + out_array, device->DebugString(), range.in_array)); } } } @@ -662,42 +665,42 @@ absl::Status RemapPlan::Validate() const { for (const auto& [out_array, inputs] : rep_->input_devices_for_output_map) { const auto out_it = out_buffer_to_in_buffer_and_devices.find(out_array); if (out_it == out_buffer_to_in_buffer_and_devices.end()) { - return InvalidArgument( + return absl::InvalidArgumentError(absl::StrFormat( "Output buffer index %d in `input_devices_for_output_map` but not " "in `mappings`", - out_array); + out_array)); } if (inputs.size() != out_it->second.size()) { - return InvalidArgument( + return absl::InvalidArgumentError(absl::StrFormat( "Output buffer index %d in `input_devices_for_output_map` has %d " "inputs, but `mappings` reference %d inputs", - out_array, inputs.size(), out_it->second.size()); + out_array, inputs.size(), out_it->second.size())); } for (const InputDeviceRange& range : inputs) { const auto in_it = out_it->second.find(range.in_array); if (in_it == out_it->second.end()) { - return InvalidArgument( + return absl::InvalidArgumentError(absl::StrFormat( "Output buffer index %d in `input_devices_for_output_map` " "references input array %d that is not present in `mappings`", - out_array, range.in_array); + out_array, range.in_array)); } if (in_it->second.size() != range.input_devices->AddressableDeviceList()->size()) { - return InvalidArgument( + return absl::InvalidArgumentError(absl::StrFormat( "Output buffer index %d in `input_devices_for_output_map` " "uses %d addressable devices from input array %d, but `mappings` " "contains %d addressable devices", out_array, range.input_devices->AddressableDeviceList()->size(), - range.in_array, in_it->second.size()); + range.in_array, in_it->second.size())); } for (const Device* const device : range.input_devices->AddressableDeviceList()->devices()) { if (!in_it->second.contains(device)) { - return InvalidArgument( + return absl::InvalidArgumentError(absl::StrFormat( "Output buffer index %d in `input_devices_for_output_map` " "references device %s from input array %d, but `mappings` does " "not reference that device", - out_array, device->DebugString(), range.in_array); + out_array, device->DebugString(), range.in_array)); } } } @@ -707,6 +710,22 @@ absl::Status RemapPlan::Validate() const { return absl::OkStatus(); } +absl::Status RemapPlan::ValidateArraySpecs() const { + absl::call_once(rep_->validate_array_specs_once, [this]() { + rep_->validate_array_specs_status = ValidateArraySpecsUncached(); + }); + return rep_->validate_array_specs_status; +} + +absl::Status RemapPlan::Validate() const { + ABSL_RETURN_IF_ERROR(ValidateArraySpecs()); + absl::call_once(rep_->validate_array_shard_mappings_once, [this]() { + rep_->validate_array_shard_mappings_status = + ValidateArrayShardMappingsUncached(); + }); + return rep_->validate_array_shard_mappings_status; +} + absl::StatusOr RemapPlan::FromProto(Client* client, const RemapPlanProto& proto) { const SerDesVersionNumber version_number(proto.version_number()); @@ -793,6 +812,7 @@ absl::Status RemapPlan::ToProto(RemapPlanProto& proto, proto.mutable_input_devices_for_output()->Reserve( rep_->input_devices_for_output_map.size()); for (const auto& [out_array, input_devices] : + // NOLINTNEXTLINE(*-custom-deterministic-iteration-order) rep_->input_devices_for_output_map) { InputDeviceToOutputToProto(version, out_array, input_devices, *proto.add_input_devices_for_output()); diff --git a/third_party/xla/xla/python/ifrt/remap_plan.h b/third_party/xla/xla/python/ifrt/remap_plan.h index 475ea85bd47f3d..6276e3e92d2f65 100644 --- a/third_party/xla/xla/python/ifrt/remap_plan.h +++ b/third_party/xla/xla/python/ifrt/remap_plan.h @@ -24,6 +24,7 @@ limitations under the License. #include #include "absl/base/attributes.h" +#include "absl/base/call_once.h" #include "absl/base/nullability.h" #include "absl/container/flat_hash_map.h" #include "absl/hash/hash.h" @@ -166,10 +167,17 @@ class RemapPlan { return rep_->input_devices_for_output_map; } - // Validates this plan against the requirements (see `RemapPlan` comment). - // This is a slow operation. It should not be performed repeatedly. - // Implementations of `Client::RemapArrays()` may bypass runtime checks on a - // plan's validity, delegating the role to this method. + // Validates array-level consistency (dtype, shard shape, memory kind, and + // layout) between input and output array pairs. The result will be cached + // within the plan. `Client::RemapArrays` implementations should at least do + // this validation. + absl::Status ValidateArraySpecs() const; + + // Validates this plan against all requirements, including array-level + // consistency (via `ValidateArraySpecs()`) and shard-level consistency (input + // array shards are correctly mapped to output array shards). This is a slow + // operation. The result will be cached within the plan. The users building a + // complex `RemapPlan` are strongly encouraged to call this method. absl::Status Validate() const; // Constructs `RemapPlan` from `RemapPlanProto`. Devices are looked up @@ -216,6 +224,18 @@ class RemapPlan { private: void Hash(absl::HashState state) const; + // Validates array-level consistency (dtype, shard shape, memory kind, and + // layout) between input and output array pairs. + absl::Status ValidateArraySpecsUncached() const; + + // Validates shard-level consistency (input array shards are correctly mapped + // to output array shards). + // + // Prerequisite: `ValidateArraySpecsUncached()` must have succeeded on this + // plan. This method assumes that array-level consistency, non-empty inputs, + // and array index bounds are already validated. + absl::Status ValidateArrayShardMappingsUncached() const; + struct Rep { // Specification of inputs. std::vector input_specs; @@ -248,12 +268,24 @@ class RemapPlan { static constexpr uint64_t kUnsetHash = 0; mutable std::atomic hash = kUnsetHash; + mutable absl::once_flag validate_array_specs_once; + mutable absl::Status validate_array_specs_status; + + mutable absl::once_flag validate_array_shard_mappings_once; + mutable absl::Status validate_array_shard_mappings_status; + Rep() = default; + Rep(std::vector input_specs, std::vector output_specs, + std::vector mappings) + : input_specs(std::move(input_specs)), + output_specs(std::move(output_specs)), + mappings(std::move(mappings)) {} + Rep(std::vector input_specs, std::vector output_specs, std::vector mappings, absl::flat_hash_map> - input_devices_for_output_map = {}) + input_devices_for_output_map) : input_specs(std::move(input_specs)), output_specs(std::move(output_specs)), mappings(std::move(mappings)), diff --git a/third_party/xla/xla/python/ifrt/remap_plan_test.cc b/third_party/xla/xla/python/ifrt/remap_plan_test.cc index 6c572f80a5944c..c2c125f1fbc7ab 100644 --- a/third_party/xla/xla/python/ifrt/remap_plan_test.cc +++ b/third_party/xla/xla/python/ifrt/remap_plan_test.cc @@ -44,8 +44,6 @@ limitations under the License. #include "xla/python/ifrt/serdes_version.h" #include "xla/python/ifrt/shape.h" #include "xla/python/ifrt/sharding.h" -#include "xla/tsl/lib/core/status_test_util.h" -#include "xla/tsl/platform/statusor.h" #include "xla/tsl/platform/test.h" namespace xla { @@ -80,6 +78,120 @@ class RemapPlanTest test_util::DeviceTestFixture fixture_; }; +TEST_P(RemapPlanTest, EmptyInputSpecs) { + std::vector output_specs; + output_specs.push_back( + ArraySpec{/*dtype=*/DType(DType::kS32), + /*shape=*/Shape({2, 3}), + /*sharding=*/ + ConcreteEvenSharding::Create(GetDevices({0}), MemoryKind(), + /*shape=*/Shape({2, 3}), + /*shard_shape=*/Shape({2, 3}))}); + RemapPlan plan(/*input_specs=*/{}, std::move(output_specs), + /*mappings=*/std::vector{}); + EXPECT_THAT( + plan.ValidateArraySpecs(), + absl_testing::StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("Must have at least one input"))); + EXPECT_THAT(plan.Validate(), absl_testing::StatusIs( + absl::StatusCode::kInvalidArgument, + HasSubstr("Must have at least one input"))); +} + +TEST_P(RemapPlanTest, EmptyOutputSpecs) { + std::vector input_specs; + input_specs.push_back( + ArraySpec{/*dtype=*/DType(DType::kS32), + /*shape=*/Shape({2, 3}), + /*sharding=*/ + ConcreteEvenSharding::Create(GetDevices({0}), MemoryKind(), + /*shape=*/Shape({2, 3}), + /*shard_shape=*/Shape({2, 3}))}); + RemapPlan plan(std::move(input_specs), /*output_specs=*/{}, + /*mappings=*/std::vector{}); + EXPECT_THAT( + plan.ValidateArraySpecs(), + absl_testing::StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("Must have at least one output"))); + EXPECT_THAT(plan.Validate(), absl_testing::StatusIs( + absl::StatusCode::kInvalidArgument, + HasSubstr("Must have at least one output"))); +} + +TEST_P(RemapPlanTest, NullInputSharding) { + std::vector input_specs; + input_specs.push_back( + ArraySpec{/*dtype=*/DType(DType::kS32), + /*shape=*/Shape({2, 3}), + /*sharding=*/ + ConcreteEvenSharding::Create(GetDevices({0}), MemoryKind(), + /*shape=*/Shape({2, 3}), + /*shard_shape=*/Shape({2, 3}))}); + input_specs.push_back(ArraySpec{/*dtype=*/DType(DType::kS32), + /*shape=*/Shape({2, 3}), + /*sharding=*/nullptr}); + std::vector output_specs; + output_specs.push_back( + ArraySpec{/*dtype=*/DType(DType::kS32), + /*shape=*/Shape({2, 3}), + /*sharding=*/ + ConcreteEvenSharding::Create(GetDevices({0}), MemoryKind(), + /*shape=*/Shape({2, 3}), + /*shard_shape=*/Shape({2, 3}))}); + std::vector mappings; + mappings.push_back(RemapPlan::Mapping{/*in_array=*/0, + /*out_array=*/0, + /*from=*/{RemapPlan::Interval{0, 1, 1}}, + /*to=*/{RemapPlan::Interval{0, 1, 1}}}); + RemapPlan plan(std::move(input_specs), std::move(output_specs), + std::move(mappings)); + EXPECT_THAT( + plan.ValidateArraySpecs(), + absl_testing::StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("Input array 1 has null sharding"))); + EXPECT_THAT( + plan.Validate(), + absl_testing::StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("Input array 1 has null sharding"))); +} + +TEST_P(RemapPlanTest, NullOutputSharding) { + std::vector input_specs; + input_specs.push_back( + ArraySpec{/*dtype=*/DType(DType::kS32), + /*shape=*/Shape({2, 3}), + /*sharding=*/ + ConcreteEvenSharding::Create(GetDevices({0}), MemoryKind(), + /*shape=*/Shape({2, 3}), + /*shard_shape=*/Shape({2, 3}))}); + std::vector output_specs; + output_specs.push_back( + ArraySpec{/*dtype=*/DType(DType::kS32), + /*shape=*/Shape({2, 3}), + /*sharding=*/ + ConcreteEvenSharding::Create(GetDevices({0}), MemoryKind(), + /*shape=*/Shape({2, 3}), + /*shard_shape=*/Shape({2, 3}))}); + output_specs.push_back(ArraySpec{/*dtype=*/DType(DType::kS32), + /*shape=*/Shape({2, 3}), + /*sharding=*/nullptr}); + std::vector mappings; + mappings.push_back(RemapPlan::Mapping{/*in_array=*/0, + /*out_array=*/0, + /*from=*/{RemapPlan::Interval{0, 1, 1}}, + /*to=*/{RemapPlan::Interval{0, 1, 1}}}); + RemapPlan plan(std::move(input_specs), std::move(output_specs), + std::move(mappings)); + EXPECT_THAT( + plan.ValidateArraySpecs(), + absl_testing::StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("Output array 1 has null sharding"))); + EXPECT_THAT( + plan.Validate(), + absl_testing::StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("Output array 1 has null sharding"))); +} + TEST_P(RemapPlanTest, EmptyMappings) { std::vector input_specs; input_specs.push_back( @@ -96,51 +208,96 @@ TEST_P(RemapPlanTest, EmptyMappings) { ConcreteEvenSharding::Create(GetDevices({0}), MemoryKind(), /*shape=*/Shape({2, 3}), /*shard_shape=*/Shape({2, 3}))}); - RemapPlan plan(std::move(input_specs), /*output_specs=*/{}, + std::vector output_specs; + output_specs.push_back( + ArraySpec{/*dtype=*/DType(DType::kS32), + /*shape=*/Shape({2, 3}), + /*sharding=*/ + ConcreteEvenSharding::Create(GetDevices({0}), MemoryKind(), + /*shape=*/Shape({2, 3}), + /*shard_shape=*/Shape({2, 3}))}); + RemapPlan plan(std::move(input_specs), std::move(output_specs), /*mappings=*/std::vector{}); + EXPECT_THAT( + plan.ValidateArraySpecs(), + absl_testing::StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("Must have at least one mapping"))); EXPECT_THAT( plan.Validate(), absl_testing::StatusIs(absl::StatusCode::kInvalidArgument, HasSubstr("Must have at least one mapping"))); } -TEST_P(RemapPlanTest, MixedDtype) { - ArraySpec array_spec_s32{ - /*dtype=*/DType(DType::kS32), - /*shape=*/Shape({2, 3}), - /*sharding=*/ - ConcreteEvenSharding::Create(GetDevices({0}), MemoryKind(), - /*shape=*/Shape({2, 3}), - /*shard_shape=*/Shape({2, 3}))}; - ArraySpec array_spec_f32{ - /*dtype=*/DType(DType::kF32), - /*shape=*/Shape({2, 3}), - /*sharding=*/ - ConcreteEvenSharding::Create(GetDevices({0}), MemoryKind(), - /*shape=*/Shape({2, 3}), - /*shard_shape=*/Shape({2, 3}))}; - +TEST_P(RemapPlanTest, InvalidInputArrayIndex) { std::vector input_specs; - input_specs.push_back(array_spec_s32); - input_specs.push_back(array_spec_f32); + input_specs.push_back( + ArraySpec{/*dtype=*/DType(DType::kS32), + /*shape=*/Shape({2, 3}), + /*sharding=*/ + ConcreteEvenSharding::Create(GetDevices({0}), MemoryKind(), + /*shape=*/Shape({2, 3}), + /*shard_shape=*/Shape({2, 3}))}); std::vector output_specs; - output_specs.push_back(array_spec_f32); - output_specs.push_back(array_spec_s32); - + output_specs.push_back( + ArraySpec{/*dtype=*/DType(DType::kS32), + /*shape=*/Shape({2, 3}), + /*sharding=*/ + ConcreteEvenSharding::Create(GetDevices({0}), MemoryKind(), + /*shape=*/Shape({2, 3}), + /*shard_shape=*/Shape({2, 3}))}); std::vector mappings; - mappings.push_back(RemapPlan::Mapping{/*in_array=*/0, - /*out_array=*/1, - /*from=*/{RemapPlan::Interval{0, 1, 1}}, - /*to=*/{RemapPlan::Interval{0, 1, 1}}}); - mappings.push_back(RemapPlan::Mapping{/*in_array=*/1, + mappings.push_back(RemapPlan::Mapping{/*in_array=*/1, // Invalid in_array /*out_array=*/0, /*from=*/{RemapPlan::Interval{0, 1, 1}}, /*to=*/{RemapPlan::Interval{0, 1, 1}}}); + RemapPlan plan(std::move(input_specs), std::move(output_specs), + std::move(mappings)); + EXPECT_THAT( + plan.ValidateArraySpecs(), + absl_testing::StatusIs( + absl::StatusCode::kInvalidArgument, + HasSubstr("mappings[0].in_array must be in [0, 0], but is 1"))); + EXPECT_THAT( + plan.Validate(), + absl_testing::StatusIs( + absl::StatusCode::kInvalidArgument, + HasSubstr("mappings[0].in_array must be in [0, 0], but is 1"))); +} - EXPECT_OK(RemapPlan::CreateOptimized(client(), std::move(input_specs), - std::move(output_specs), - std::move(mappings)) - .status()); +TEST_P(RemapPlanTest, InvalidOutputArrayIndex) { + std::vector input_specs; + input_specs.push_back( + ArraySpec{/*dtype=*/DType(DType::kS32), + /*shape=*/Shape({2, 3}), + /*sharding=*/ + ConcreteEvenSharding::Create(GetDevices({0}), MemoryKind(), + /*shape=*/Shape({2, 3}), + /*shard_shape=*/Shape({2, 3}))}); + std::vector output_specs; + output_specs.push_back( + ArraySpec{/*dtype=*/DType(DType::kS32), + /*shape=*/Shape({2, 3}), + /*sharding=*/ + ConcreteEvenSharding::Create(GetDevices({0}), MemoryKind(), + /*shape=*/Shape({2, 3}), + /*shard_shape=*/Shape({2, 3}))}); + std::vector mappings; + mappings.push_back(RemapPlan::Mapping{/*in_array=*/0, + /*out_array=*/1, // Invalid out_array + /*from=*/{RemapPlan::Interval{0, 1, 1}}, + /*to=*/{RemapPlan::Interval{0, 1, 1}}}); + RemapPlan plan(std::move(input_specs), std::move(output_specs), + std::move(mappings)); + EXPECT_THAT( + plan.ValidateArraySpecs(), + absl_testing::StatusIs( + absl::StatusCode::kInvalidArgument, + HasSubstr("mappings[0].out_array must be in [0, 0], but is 1"))); + EXPECT_THAT( + plan.Validate(), + absl_testing::StatusIs( + absl::StatusCode::kInvalidArgument, + HasSubstr("mappings[0].out_array must be in [0, 0], but is 1"))); } TEST_P(RemapPlanTest, InvalidOutputDtype) { @@ -167,6 +324,10 @@ TEST_P(RemapPlanTest, InvalidOutputDtype) { /*to=*/{RemapPlan::Interval{0, 1, 1}}}); RemapPlan plan(std::move(input_specs), std::move(output_specs), std::move(mappings)); + EXPECT_THAT(plan.ValidateArraySpecs(), + absl_testing::StatusIs( + absl::StatusCode::kInvalidArgument, + HasSubstr("Input and output must have the same dtype"))); EXPECT_THAT(plan.Validate(), absl_testing::StatusIs( absl::StatusCode::kInvalidArgument, @@ -208,12 +369,88 @@ TEST_P(RemapPlanTest, InvalidOutputDtypeFromMixedInputDtypes) { RemapPlan plan(std::move(input_specs), std::move(output_specs), std::move(mappings)); + EXPECT_THAT(plan.ValidateArraySpecs(), + absl_testing::StatusIs( + absl::StatusCode::kInvalidArgument, + HasSubstr("Input and output must have the same dtype"))); EXPECT_THAT(plan.Validate(), absl_testing::StatusIs( absl::StatusCode::kInvalidArgument, HasSubstr("Input and output must have the same dtype"))); } +TEST_P(RemapPlanTest, InvalidShardShape) { + std::vector input_specs; + input_specs.push_back( + ArraySpec{/*dtype=*/DType(DType::kS32), + /*shape=*/Shape({2, 3}), + /*sharding=*/ + ConcreteEvenSharding::Create(GetDevices({0}), MemoryKind(), + /*shape=*/Shape({2, 3}), + /*shard_shape=*/Shape({2, 3}))}); + std::vector output_specs; + output_specs.push_back( + ArraySpec{/*dtype=*/DType(DType::kS32), + /*shape=*/Shape({3, 2}), + /*sharding=*/ + ConcreteEvenSharding::Create(GetDevices({0}), MemoryKind(), + /*shape=*/Shape({3, 2}), + /*shard_shape=*/Shape({3, 2}))}); + std::vector mappings; + mappings.push_back(RemapPlan::Mapping{/*in_array=*/0, + /*out_array=*/0, + /*from=*/{RemapPlan::Interval{0, 1, 1}}, + /*to=*/{RemapPlan::Interval{0, 1, 1}}}); + RemapPlan plan(std::move(input_specs), std::move(output_specs), + std::move(mappings)); + EXPECT_THAT( + plan.ValidateArraySpecs(), + absl_testing::StatusIs( + absl::StatusCode::kInvalidArgument, + HasSubstr("Input and output must have the same shard shape"))); + EXPECT_THAT( + plan.Validate(), + absl_testing::StatusIs( + absl::StatusCode::kInvalidArgument, + HasSubstr("Input and output must have the same shard shape"))); +} + +TEST_P(RemapPlanTest, InvalidMemoryKind) { + std::vector input_specs; + input_specs.push_back(ArraySpec{ + /*dtype=*/DType(DType::kS32), + /*shape=*/Shape({2, 3}), + /*sharding=*/ + ConcreteEvenSharding::Create(GetDevices({0}), MemoryKind("host"), + /*shape=*/Shape({2, 3}), + /*shard_shape=*/Shape({2, 3}))}); + std::vector output_specs; + output_specs.push_back( + ArraySpec{/*dtype=*/DType(DType::kS32), + /*shape=*/Shape({2, 3}), + /*sharding=*/ + ConcreteEvenSharding::Create(GetDevices({0}), MemoryKind(), + /*shape=*/Shape({2, 3}), + /*shard_shape=*/Shape({2, 3}))}); + std::vector mappings; + mappings.push_back(RemapPlan::Mapping{/*in_array=*/0, + /*out_array=*/0, + /*from=*/{RemapPlan::Interval{0, 1, 1}}, + /*to=*/{RemapPlan::Interval{0, 1, 1}}}); + RemapPlan plan(std::move(input_specs), std::move(output_specs), + std::move(mappings)); + EXPECT_THAT( + plan.ValidateArraySpecs(), + absl_testing::StatusIs( + absl::StatusCode::kInvalidArgument, + HasSubstr("Input and output must have the same memory kind"))); + EXPECT_THAT( + plan.Validate(), + absl_testing::StatusIs( + absl::StatusCode::kInvalidArgument, + HasSubstr("Input and output must have the same memory kind"))); +} + TEST_P(RemapPlanTest, InvalidLayout) { std::vector input_specs; input_specs.push_back(ArraySpec{ @@ -246,6 +483,10 @@ TEST_P(RemapPlanTest, InvalidLayout) { /*to=*/{RemapPlan::Interval{0, 1, 1}}}); RemapPlan plan(std::move(input_specs), std::move(output_specs), std::move(mappings)); + EXPECT_THAT(plan.ValidateArraySpecs(), + absl_testing::StatusIs( + absl::StatusCode::kInvalidArgument, + HasSubstr("Input and output must have the same layout"))); EXPECT_THAT(plan.Validate(), absl_testing::StatusIs( absl::StatusCode::kInvalidArgument, @@ -269,84 +510,23 @@ TEST_P(RemapPlanTest, ValidLayoutFromDifferentLayoutObjects) { output_specs.push_back(ArraySpec{ /*dtype=*/DType(DType::kS32), /*shape=*/Shape({2, 3}), - /*sharding=*/ - ConcreteEvenSharding::Create(GetDevices({0}), MemoryKind(), - /*shape=*/Shape({2, 3}), - /*shard_shape=*/Shape({2, 3})), - /*layout=*/ - std::make_shared( - xla::LayoutUtil::MakeAscendingLayout(2)), // same layout - }); - std::vector mappings; - mappings.push_back(RemapPlan::Mapping{/*in_array=*/0, - /*out_array=*/0, - /*from=*/{RemapPlan::Interval{0, 1, 1}}, - /*to=*/{RemapPlan::Interval{0, 1, 1}}}); - RemapPlan plan(std::move(input_specs), std::move(output_specs), - std::move(mappings)); - EXPECT_OK(plan.Validate()); -} - -TEST_P(RemapPlanTest, InvalidInputArrayIndex) { - std::vector input_specs; - input_specs.push_back( - ArraySpec{/*dtype=*/DType(DType::kS32), - /*shape=*/Shape({2, 3}), - /*sharding=*/ - ConcreteEvenSharding::Create(GetDevices({0}), MemoryKind(), - /*shape=*/Shape({2, 3}), - /*shard_shape=*/Shape({2, 3}))}); - std::vector output_specs; - output_specs.push_back( - ArraySpec{/*dtype=*/DType(DType::kS32), - /*shape=*/Shape({2, 3}), - /*sharding=*/ - ConcreteEvenSharding::Create(GetDevices({0}), MemoryKind(), - /*shape=*/Shape({2, 3}), - /*shard_shape=*/Shape({2, 3}))}); - std::vector mappings; - mappings.push_back(RemapPlan::Mapping{/*in_array=*/1, // Invalid in_array - /*out_array=*/0, - /*from=*/{RemapPlan::Interval{0, 1, 1}}, - /*to=*/{RemapPlan::Interval{0, 1, 1}}}); - RemapPlan plan(std::move(input_specs), std::move(output_specs), - std::move(mappings)); - EXPECT_THAT( - plan.Validate(), - absl_testing::StatusIs( - absl::StatusCode::kInvalidArgument, - HasSubstr("mappings[0].in_array must be in [0, 0], but is 1"))); -} - -TEST_P(RemapPlanTest, InvalidOutputArrayIndex) { - std::vector input_specs; - input_specs.push_back( - ArraySpec{/*dtype=*/DType(DType::kS32), - /*shape=*/Shape({2, 3}), - /*sharding=*/ - ConcreteEvenSharding::Create(GetDevices({0}), MemoryKind(), - /*shape=*/Shape({2, 3}), - /*shard_shape=*/Shape({2, 3}))}); - std::vector output_specs; - output_specs.push_back( - ArraySpec{/*dtype=*/DType(DType::kS32), - /*shape=*/Shape({2, 3}), - /*sharding=*/ - ConcreteEvenSharding::Create(GetDevices({0}), MemoryKind(), - /*shape=*/Shape({2, 3}), - /*shard_shape=*/Shape({2, 3}))}); + /*sharding=*/ + ConcreteEvenSharding::Create(GetDevices({0}), MemoryKind(), + /*shape=*/Shape({2, 3}), + /*shard_shape=*/Shape({2, 3})), + /*layout=*/ + std::make_shared( + xla::LayoutUtil::MakeAscendingLayout(2)), // same layout + }); std::vector mappings; mappings.push_back(RemapPlan::Mapping{/*in_array=*/0, - /*out_array=*/1, // Invalid out_array + /*out_array=*/0, /*from=*/{RemapPlan::Interval{0, 1, 1}}, /*to=*/{RemapPlan::Interval{0, 1, 1}}}); RemapPlan plan(std::move(input_specs), std::move(output_specs), std::move(mappings)); - EXPECT_THAT( - plan.Validate(), - absl_testing::StatusIs( - absl::StatusCode::kInvalidArgument, - HasSubstr("mappings[0].out_array must be in [0, 0], but is 1"))); + EXPECT_OK(plan.ValidateArraySpecs()); + EXPECT_OK(plan.Validate()); } TEST_P(RemapPlanTest, InvalidIntervalCount) { @@ -609,175 +789,6 @@ TEST_P(RemapPlanTest, InvalidOutputDevices) { "sharding devices do not match"))); } -TEST_P(RemapPlanTest, CheckOneInputToOneOutput) { - ArraySpec dummy_spec = GetDummySpec(); - - std::vector mappings = { - RemapPlan::Mapping{/*in_array=*/0, - /*out_array=*/0, - /*from=*/{RemapPlan::Interval{0, 2, 1}}, - /*to=*/{RemapPlan::Interval{0, 2, 1}}}}; - RemapPlan plan({dummy_spec}, {dummy_spec}, std::move(mappings)); - - TF_EXPECT_OK( - plan.CheckArrayCopySemantics(xla::ifrt::ArrayCopySemantics::kReuseInput)); - TF_EXPECT_OK(plan.CheckArrayCopySemantics( - xla::ifrt::ArrayCopySemantics::kDonateInput)); -} - -TEST_P(RemapPlanTest, CheckOneInputToMultipleOutputs) { - ArraySpec dummy_spec = GetDummySpec(); - - std::vector mappings = { - RemapPlan::Mapping{/*in_array=*/0, - /*out_array=*/0, - /*from=*/{RemapPlan::Interval{0, 2, 1}}, - /*to=*/{RemapPlan::Interval{0, 2, 1}}}, - RemapPlan::Mapping{/*in_array=*/0, - /*out_array=*/1, - /*from=*/{RemapPlan::Interval{0, 2, 1}}, - /*to=*/{RemapPlan::Interval{0, 2, 1}}}}; - RemapPlan plan({dummy_spec}, {dummy_spec, dummy_spec}, std::move(mappings)); - - TF_EXPECT_OK( - plan.CheckArrayCopySemantics(xla::ifrt::ArrayCopySemantics::kReuseInput)); - TF_EXPECT_OK(plan.CheckArrayCopySemantics( - xla::ifrt::ArrayCopySemantics::kDonateInput)); -} - -TEST_P(RemapPlanTest, CheckMultipleInputsToOneOutput) { - ArraySpec dummy_spec = GetDummySpec(); - - std::vector mappings = { - RemapPlan::Mapping{/*in_array=*/0, - /*out_array=*/0, - /*from=*/{RemapPlan::Interval{0, 2, 1}}, - /*to=*/{RemapPlan::Interval{0, 2, 1}}}, - RemapPlan::Mapping{/*in_array=*/1, - /*out_array=*/0, - /*from=*/{RemapPlan::Interval{0, 2, 1}}, - /*to=*/{RemapPlan::Interval{0, 2, 1}}}}; - RemapPlan plan({dummy_spec, dummy_spec}, {dummy_spec}, std::move(mappings)); - - EXPECT_THAT( - plan.CheckArrayCopySemantics(xla::ifrt::ArrayCopySemantics::kReuseInput), - absl_testing::StatusIs( - absl::StatusCode::kInvalidArgument, - HasSubstr("kDonateInput is required if multiple inputs are " - "mapped to one output"))); - TF_EXPECT_OK(plan.CheckArrayCopySemantics( - xla::ifrt::ArrayCopySemantics::kDonateInput)); -} - -TEST_P(RemapPlanTest, InvalidInputDevicesForOutputMap) { - ArraySpec dummy_spec = GetDummySpec(); - - std::vector input_specs = {dummy_spec, dummy_spec}; - std::vector output_specs = {dummy_spec}; - - std::vector mappings = { - RemapPlan::Mapping{/*in_array=*/0, - /*out_array=*/0, - /*from=*/{RemapPlan::Interval{0, 1, 1}}, - /*to=*/{RemapPlan::Interval{0, 1, 1}}}, - RemapPlan::Mapping{/*in_array=*/1, - /*out_array=*/0, - /*from=*/{RemapPlan::Interval{1, 2, 1}}, - /*to=*/{RemapPlan::Interval{1, 2, 1}}}}; - - { - absl::flat_hash_map> - input_devices_for_output_map; - input_devices_for_output_map.insert({1, {}}); - RemapPlan plan(input_specs, output_specs, mappings, - std::move(input_devices_for_output_map)); - EXPECT_THAT(plan.Validate(), - absl_testing::StatusIs(absl::StatusCode::kInvalidArgument, - HasSubstr("Output buffer index 1"))); - } - - { - absl::flat_hash_map> - input_devices_for_output_map; - input_devices_for_output_map.insert( - {0, {{1, dummy_spec.sharding->devices()}}}); - RemapPlan plan(input_specs, output_specs, mappings, - std::move(input_devices_for_output_map)); - EXPECT_THAT(plan.Validate(), - absl_testing::StatusIs(absl::StatusCode::kInvalidArgument, - HasSubstr("has 1 inputs"))); - } - - { - absl::flat_hash_map> - input_devices_for_output_map; - input_devices_for_output_map.insert( - {0, - {{3, dummy_spec.sharding->devices()}, - {4, dummy_spec.sharding->devices()}}}); - RemapPlan plan(input_specs, output_specs, mappings, - std::move(input_devices_for_output_map)); - EXPECT_THAT(plan.Validate(), - absl_testing::StatusIs(absl::StatusCode::kInvalidArgument, - HasSubstr("Input buffer index 3"))); - } - - { - std::vector three_input_specs = {dummy_spec, dummy_spec, - dummy_spec}; - absl::flat_hash_map> - input_devices_for_output_map; - input_devices_for_output_map.insert( - {0, {{0, GetDevices({0})}, {2, dummy_spec.sharding->devices()}}}); - RemapPlan plan(three_input_specs, output_specs, mappings, - std::move(input_devices_for_output_map)); - EXPECT_THAT(plan.Validate(), - absl_testing::StatusIs(absl::StatusCode::kInvalidArgument, - HasSubstr("not present in `mappings`"))); - } - - { - absl::flat_hash_map> - input_devices_for_output_map; - input_devices_for_output_map.insert( - {0, {{0, GetDevices({1})}, {1, dummy_spec.sharding->devices()}}}); - RemapPlan plan(input_specs, output_specs, mappings, - std::move(input_devices_for_output_map)); - EXPECT_THAT( - plan.Validate(), - absl_testing::StatusIs(absl::StatusCode::kInvalidArgument, - HasSubstr("does not reference that device"))); - } - - { - std::vector multi_output_specs = {dummy_spec, dummy_spec}; - std::vector multi_mappings = { - RemapPlan::Mapping{/*in_array=*/0, - /*out_array=*/0, - /*from=*/{RemapPlan::Interval{0, 2, 1}}, - /*to=*/{RemapPlan::Interval{0, 2, 1}}}, - RemapPlan::Mapping{/*in_array=*/1, - /*out_array=*/1, - /*from=*/{RemapPlan::Interval{0, 2, 1}}, - /*to=*/{RemapPlan::Interval{0, 2, 1}}}}; - absl::flat_hash_map> - input_devices_for_output_map; - input_devices_for_output_map.insert( - {0, {{0, dummy_spec.sharding->devices()}}}); - RemapPlan plan(input_specs, multi_output_specs, multi_mappings, - std::move(input_devices_for_output_map)); - EXPECT_THAT(plan.Validate(), - absl_testing::StatusIs( - absl::StatusCode::kInvalidArgument, - HasSubstr("has 1 outputs, but expected 2 outputs"))); - } - - ASSERT_OK(RemapPlan::CreateOptimized(client(), std::move(input_specs), - std::move(output_specs), - std::move(mappings)) - .status()); -} - TEST_P(RemapPlanTest, InputDevicesForOutputMapWithoutMappings) { ArraySpec dummy_spec = GetDummySpec(); @@ -872,15 +883,184 @@ TEST_P(RemapPlanTest, InvalidInputDevicesForOutputMapWithoutMappings) { } { - absl::flat_hash_map> map; - map.insert({0, {{0, GetDevices({2})}}}); - RemapPlan plan(input_specs, output_specs, std::move(map)); + absl::flat_hash_map> map; + map.insert({0, {{0, GetDevices({2})}}}); + RemapPlan plan(input_specs, output_specs, std::move(map)); + EXPECT_THAT( + plan.Validate(), + absl_testing::StatusIs( + absl::StatusCode::kInvalidArgument, + HasSubstr("not in the input array's addressable device list"))); + } +} + +TEST_P(RemapPlanTest, InvalidInputDevicesForOutputMap) { + ArraySpec dummy_spec = GetDummySpec(); + + std::vector input_specs = {dummy_spec, dummy_spec}; + std::vector output_specs = {dummy_spec}; + + std::vector mappings = { + RemapPlan::Mapping{/*in_array=*/0, + /*out_array=*/0, + /*from=*/{RemapPlan::Interval{0, 1, 1}}, + /*to=*/{RemapPlan::Interval{0, 1, 1}}}, + RemapPlan::Mapping{/*in_array=*/1, + /*out_array=*/0, + /*from=*/{RemapPlan::Interval{1, 2, 1}}, + /*to=*/{RemapPlan::Interval{1, 2, 1}}}}; + + { + absl::flat_hash_map> + input_devices_for_output_map; + input_devices_for_output_map.insert({1, {}}); + RemapPlan plan(input_specs, output_specs, mappings, + std::move(input_devices_for_output_map)); + EXPECT_THAT(plan.Validate(), + absl_testing::StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("Output buffer index 1"))); + } + + { + absl::flat_hash_map> + input_devices_for_output_map; + input_devices_for_output_map.insert( + {0, {{1, dummy_spec.sharding->devices()}}}); + RemapPlan plan(input_specs, output_specs, mappings, + std::move(input_devices_for_output_map)); + EXPECT_THAT(plan.Validate(), + absl_testing::StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("has 1 inputs"))); + } + + { + absl::flat_hash_map> + input_devices_for_output_map; + input_devices_for_output_map.insert( + {0, + {{3, dummy_spec.sharding->devices()}, + {4, dummy_spec.sharding->devices()}}}); + RemapPlan plan(input_specs, output_specs, mappings, + std::move(input_devices_for_output_map)); + EXPECT_THAT(plan.Validate(), + absl_testing::StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("Input buffer index 3"))); + } + + { + std::vector three_input_specs = {dummy_spec, dummy_spec, + dummy_spec}; + absl::flat_hash_map> + input_devices_for_output_map; + input_devices_for_output_map.insert( + {0, {{0, GetDevices({0})}, {2, dummy_spec.sharding->devices()}}}); + RemapPlan plan(three_input_specs, output_specs, mappings, + std::move(input_devices_for_output_map)); + EXPECT_THAT(plan.Validate(), + absl_testing::StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("not present in `mappings`"))); + } + + { + absl::flat_hash_map> + input_devices_for_output_map; + input_devices_for_output_map.insert( + {0, {{0, GetDevices({1})}, {1, dummy_spec.sharding->devices()}}}); + RemapPlan plan(input_specs, output_specs, mappings, + std::move(input_devices_for_output_map)); EXPECT_THAT( plan.Validate(), - absl_testing::StatusIs( - absl::StatusCode::kInvalidArgument, - HasSubstr("not in the input array's addressable device list"))); + absl_testing::StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("does not reference that device"))); + } + + { + std::vector multi_output_specs = {dummy_spec, dummy_spec}; + std::vector multi_mappings = { + RemapPlan::Mapping{/*in_array=*/0, + /*out_array=*/0, + /*from=*/{RemapPlan::Interval{0, 2, 1}}, + /*to=*/{RemapPlan::Interval{0, 2, 1}}}, + RemapPlan::Mapping{/*in_array=*/1, + /*out_array=*/1, + /*from=*/{RemapPlan::Interval{0, 2, 1}}, + /*to=*/{RemapPlan::Interval{0, 2, 1}}}}; + absl::flat_hash_map> + input_devices_for_output_map; + input_devices_for_output_map.insert( + {0, {{0, dummy_spec.sharding->devices()}}}); + RemapPlan plan(input_specs, multi_output_specs, multi_mappings, + std::move(input_devices_for_output_map)); + EXPECT_THAT(plan.Validate(), + absl_testing::StatusIs( + absl::StatusCode::kInvalidArgument, + HasSubstr("has 1 outputs, but expected 2 outputs"))); } + + ASSERT_OK(RemapPlan::CreateOptimized(client(), std::move(input_specs), + std::move(output_specs), + std::move(mappings)) + .status()); +} + +TEST_P(RemapPlanTest, CheckOneInputToOneOutput) { + ArraySpec dummy_spec = GetDummySpec(); + + std::vector mappings = { + RemapPlan::Mapping{/*in_array=*/0, + /*out_array=*/0, + /*from=*/{RemapPlan::Interval{0, 2, 1}}, + /*to=*/{RemapPlan::Interval{0, 2, 1}}}}; + RemapPlan plan({dummy_spec}, {dummy_spec}, std::move(mappings)); + + EXPECT_OK( + plan.CheckArrayCopySemantics(xla::ifrt::ArrayCopySemantics::kReuseInput)); + EXPECT_OK(plan.CheckArrayCopySemantics( + xla::ifrt::ArrayCopySemantics::kDonateInput)); +} + +TEST_P(RemapPlanTest, CheckOneInputToMultipleOutputs) { + ArraySpec dummy_spec = GetDummySpec(); + + std::vector mappings = { + RemapPlan::Mapping{/*in_array=*/0, + /*out_array=*/0, + /*from=*/{RemapPlan::Interval{0, 2, 1}}, + /*to=*/{RemapPlan::Interval{0, 2, 1}}}, + RemapPlan::Mapping{/*in_array=*/0, + /*out_array=*/1, + /*from=*/{RemapPlan::Interval{0, 2, 1}}, + /*to=*/{RemapPlan::Interval{0, 2, 1}}}}; + RemapPlan plan({dummy_spec}, {dummy_spec, dummy_spec}, std::move(mappings)); + + EXPECT_OK( + plan.CheckArrayCopySemantics(xla::ifrt::ArrayCopySemantics::kReuseInput)); + EXPECT_OK(plan.CheckArrayCopySemantics( + xla::ifrt::ArrayCopySemantics::kDonateInput)); +} + +TEST_P(RemapPlanTest, CheckMultipleInputsToOneOutput) { + ArraySpec dummy_spec = GetDummySpec(); + + std::vector mappings = { + RemapPlan::Mapping{/*in_array=*/0, + /*out_array=*/0, + /*from=*/{RemapPlan::Interval{0, 2, 1}}, + /*to=*/{RemapPlan::Interval{0, 2, 1}}}, + RemapPlan::Mapping{/*in_array=*/1, + /*out_array=*/0, + /*from=*/{RemapPlan::Interval{0, 2, 1}}, + /*to=*/{RemapPlan::Interval{0, 2, 1}}}}; + RemapPlan plan({dummy_spec, dummy_spec}, {dummy_spec}, std::move(mappings)); + + EXPECT_THAT( + plan.CheckArrayCopySemantics(xla::ifrt::ArrayCopySemantics::kReuseInput), + absl_testing::StatusIs( + absl::StatusCode::kInvalidArgument, + HasSubstr("kDonateInput is required if multiple inputs are " + "mapped to one output"))); + EXPECT_OK(plan.CheckArrayCopySemantics( + xla::ifrt::ArrayCopySemantics::kDonateInput)); } TEST_P(RemapPlanTest, CheckArrayCopySemanticsWithoutMappings) { @@ -890,9 +1070,9 @@ TEST_P(RemapPlanTest, CheckArrayCopySemanticsWithoutMappings) { absl::flat_hash_map> map; map.insert({0, {{0, dummy_spec.sharding->devices()}}}); RemapPlan plan({dummy_spec}, {dummy_spec}, std::move(map)); - TF_EXPECT_OK(plan.CheckArrayCopySemantics( + EXPECT_OK(plan.CheckArrayCopySemantics( xla::ifrt::ArrayCopySemantics::kReuseInput)); - TF_EXPECT_OK(plan.CheckArrayCopySemantics( + EXPECT_OK(plan.CheckArrayCopySemantics( xla::ifrt::ArrayCopySemantics::kDonateInput)); } @@ -906,11 +1086,107 @@ TEST_P(RemapPlanTest, CheckArrayCopySemanticsWithoutMappings) { absl::StatusCode::kInvalidArgument, HasSubstr("kDonateInput is required if multiple inputs are " "mapped to one output"))); - TF_EXPECT_OK(plan.CheckArrayCopySemantics( + EXPECT_OK(plan.CheckArrayCopySemantics( xla::ifrt::ArrayCopySemantics::kDonateInput)); } } +TEST_P(RemapPlanTest, MixedDtype) { + ArraySpec array_spec_s32{ + /*dtype=*/DType(DType::kS32), + /*shape=*/Shape({2, 3}), + /*sharding=*/ + ConcreteEvenSharding::Create(GetDevices({0}), MemoryKind(), + /*shape=*/Shape({2, 3}), + /*shard_shape=*/Shape({2, 3}))}; + ArraySpec array_spec_f32{ + /*dtype=*/DType(DType::kF32), + /*shape=*/Shape({2, 3}), + /*sharding=*/ + ConcreteEvenSharding::Create(GetDevices({0}), MemoryKind(), + /*shape=*/Shape({2, 3}), + /*shard_shape=*/Shape({2, 3}))}; + + std::vector input_specs; + input_specs.push_back(array_spec_s32); + input_specs.push_back(array_spec_f32); + std::vector output_specs; + output_specs.push_back(array_spec_f32); + output_specs.push_back(array_spec_s32); + + std::vector mappings; + mappings.push_back(RemapPlan::Mapping{/*in_array=*/0, + /*out_array=*/1, + /*from=*/{RemapPlan::Interval{0, 1, 1}}, + /*to=*/{RemapPlan::Interval{0, 1, 1}}}); + mappings.push_back(RemapPlan::Mapping{/*in_array=*/1, + /*out_array=*/0, + /*from=*/{RemapPlan::Interval{0, 1, 1}}, + /*to=*/{RemapPlan::Interval{0, 1, 1}}}); + + EXPECT_OK(RemapPlan::CreateOptimized(client(), std::move(input_specs), + std::move(output_specs), + std::move(mappings)) + .status()); +} + +TEST_P(RemapPlanTest, CreateOptimizedIntervalEndExceedsNumShards) { + ArraySpec spec{ + /*dtype=*/DType(DType::kS32), + /*shape=*/Shape({4, 6}), + /*sharding=*/ + ConcreteEvenSharding::Create(GetDevices({0, 1, 2, 3}), MemoryKind(), + /*shape=*/Shape({4, 6}), + /*shard_shape=*/Shape({2, 3}))}; + + std::vector input_specs; + input_specs.push_back(spec); + std::vector output_specs; + output_specs.push_back(spec); + + std::vector mappings; + mappings.push_back(RemapPlan::Mapping{/*in_array=*/0, + /*out_array=*/0, + /*from=*/{RemapPlan::Interval{0, 5, 2}}, + /*to=*/{RemapPlan::Interval{0, 2, 1}}}); + + EXPECT_THAT( + RemapPlan::CreateOptimized(client(), std::move(input_specs), + std::move(output_specs), std::move(mappings)), + absl_testing::StatusIs( + absl::StatusCode::kInvalidArgument, + HasSubstr("interval addresses shard 4, which is out of " + "range [0, 4)"))); +} + +TEST_P(RemapPlanTest, CreateOptimizedNegativeStartNoCrash) { + ArraySpec spec{ + /*dtype=*/DType(DType::kS32), + /*shape=*/Shape({4, 6}), + /*sharding=*/ + ConcreteEvenSharding::Create(GetDevices({0, 1, 2, 3}), MemoryKind(), + /*shape=*/Shape({4, 6}), + /*shard_shape=*/Shape({2, 3}))}; + + std::vector input_specs; + input_specs.push_back(spec); + std::vector output_specs; + output_specs.push_back(spec); + + std::vector mappings; + mappings.push_back( + RemapPlan::Mapping{/*in_array=*/0, + /*out_array=*/0, + /*from=*/{RemapPlan::Interval{-1, 0, 1}}, + /*to=*/{RemapPlan::Interval{0, 1, 1}}}); + + EXPECT_THAT( + RemapPlan::CreateOptimized(client(), std::move(input_specs), + std::move(output_specs), std::move(mappings)), + absl_testing::StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("start must be in [0, 3], but is -1"))); +} + TEST_P(RemapPlanTest, Hash) { std::vector plans; plans.push_back(RemapPlan()); @@ -978,63 +1254,6 @@ TEST_P(RemapPlanTest, Hash) { EXPECT_TRUE(absl::VerifyTypeImplementsAbslHashCorrectly(plans)); } -TEST_P(RemapPlanTest, CreateOptimizedIntervalEndExceedsNumShards) { - ArraySpec spec{ - /*dtype=*/DType(DType::kS32), - /*shape=*/Shape({4, 6}), - /*sharding=*/ - ConcreteEvenSharding::Create(GetDevices({0, 1, 2, 3}), MemoryKind(), - /*shape=*/Shape({4, 6}), - /*shard_shape=*/Shape({2, 3}))}; - - std::vector input_specs; - input_specs.push_back(spec); - std::vector output_specs; - output_specs.push_back(spec); - - std::vector mappings; - mappings.push_back(RemapPlan::Mapping{/*in_array=*/0, - /*out_array=*/0, - /*from=*/{RemapPlan::Interval{0, 5, 2}}, - /*to=*/{RemapPlan::Interval{0, 2, 1}}}); - - EXPECT_THAT( - RemapPlan::CreateOptimized(client(), std::move(input_specs), - std::move(output_specs), std::move(mappings)), - absl_testing::StatusIs( - absl::StatusCode::kInvalidArgument, - HasSubstr("interval addresses shard 4, which is out of " - "range [0, 4)"))); -} - -TEST_P(RemapPlanTest, CreateOptimizedNegativeStartNoCrash) { - ArraySpec spec{ - /*dtype=*/DType(DType::kS32), - /*shape=*/Shape({4, 6}), - /*sharding=*/ - ConcreteEvenSharding::Create(GetDevices({0, 1, 2, 3}), MemoryKind(), - /*shape=*/Shape({4, 6}), - /*shard_shape=*/Shape({2, 3}))}; - - std::vector input_specs; - input_specs.push_back(spec); - std::vector output_specs; - output_specs.push_back(spec); - - std::vector mappings; - mappings.push_back( - RemapPlan::Mapping{/*in_array=*/0, - /*out_array=*/0, - /*from=*/{RemapPlan::Interval{-1, 0, 1}}, - /*to=*/{RemapPlan::Interval{0, 1, 1}}}); - - EXPECT_THAT( - RemapPlan::CreateOptimized(client(), std::move(input_specs), - std::move(output_specs), std::move(mappings)), - absl_testing::StatusIs(absl::StatusCode::kInvalidArgument, - HasSubstr("start must be in [0, 3], but is -1"))); -} - INSTANTIATE_TEST_SUITE_P(NumDevices, RemapPlanTest, testing::Values(test_util::DeviceTestParam{ /*num_devices=*/4, @@ -1103,14 +1322,15 @@ TEST_P(RemapPlanSerDesTest, ToFromProto) { RemapPlan plan(std::move(input_specs), std::move(output_specs), std::move(mappings), std::move(input_devices_for_output_map)); - TF_ASSERT_OK_AND_ASSIGN(RemapPlanProto plan_proto, plan.ToProto(version())); - TF_ASSERT_OK_AND_ASSIGN(RemapPlan plan_copy, - RemapPlan::FromProto(client(), plan_proto)); + ASSERT_OK_AND_ASSIGN(RemapPlanProto plan_proto, plan.ToProto(version())); + ASSERT_OK_AND_ASSIGN(RemapPlan plan_copy, + RemapPlan::FromProto(client(), plan_proto)); EXPECT_THAT(plan_copy.mappings(), ElementsAreArray(plan.mappings())); ASSERT_EQ(plan.input_devices_for_output_map().size(), plan_copy.input_devices_for_output_map().size()); for (const auto& [out_array, input_devices] : + // NOLINTNEXTLINE(*-custom-deterministic-iteration-order) plan.input_devices_for_output_map()) { ASSERT_TRUE(plan_copy.input_devices_for_output_map().contains(out_array)); const auto& copy_input_devices = diff --git a/third_party/xla/xla/python/ifrt/rtti.h b/third_party/xla/xla/python/ifrt/rtti.h index 656c1203a4f06d..d72aecbc2abc1d 100644 --- a/third_party/xla/xla/python/ifrt/rtti.h +++ b/third_party/xla/xla/python/ifrt/rtti.h @@ -89,23 +89,22 @@ class RTTIExtends : public ParentT { static_assert(ParentT::kDepth + 1 < internal::kMaxRTTIDepth, "Exceeded maximum supported IFRT RTTI inheritance depth."); - // TODO(hyeontaek): Add a `static_assert` that `&ThisT::ID is different from - // `&ParentT::ID`. While not defining `ID` is handy when dynamic casting is - // not expected (e.g., classes defined for testing), it is error-prone in - // general. Thus, it would be safer to require `ID` to be defined. Note that - // we cannot enforce this requirement yet because the user code must be first - // migrated to define `ID`. - RTTIExtends& operator=(const RTTIExtends&) = default; RTTIExtends& operator=(RTTIExtends&&) = default; // Returns the static class ID for `ThisT`. - static const void* classID() { return &ThisT::ID; } + static const void* classID() { + static_assert(&ThisT::ID != &ParentT::ID, + "ThisT must define its own ID: static char ID;"); + return &ThisT::ID; + } private: template static constexpr internal::RTTITypeInfo MakeTypeInfo( std::index_sequence) { + static_assert(&ThisT::ID != &ParentT::ID, + "ThisT must define its own ID: static char ID;"); return {&ThisT::ID, {(Is < kDepth ? ParentT::kTypeInfo.ancestors[Is] : (Is == kDepth ? &ThisT::ID : nullptr))...}}; diff --git a/third_party/xla/xla/python/ifrt/rtti_nc_test.cc b/third_party/xla/xla/python/ifrt/rtti_nc_test.cc index 808a9cbb45f79a..d4776e5e0f60de 100644 --- a/third_party/xla/xla/python/ifrt/rtti_nc_test.cc +++ b/third_party/xla/xla/python/ifrt/rtti_nc_test.cc @@ -56,6 +56,21 @@ class Depth7 : public RTTIExtends { static char ID; // NOLINT }; +class BaseWithoutId : public RTTIExtends {}; +// expected-error@* {{ThisT must define its own ID}} +[[maybe_unused]] BaseWithoutId base_without_id; + +class BaseWithId : public RTTIExtends { + public: + static char ID; // NOLINT +}; + +[[maybe_unused]] char BaseWithId::ID = 0; + +class DerivedWithoutId : public RTTIExtends {}; +// expected-error@* {{ThisT must define its own ID}} +[[maybe_unused]] DerivedWithoutId derived_without_id; + } // namespace namespace { @@ -101,6 +116,14 @@ TEST(RttiNcTest, UnrelatedType) { dyn_cast(&derived_b); } +class UninstantiatedWithoutId + : public RTTIExtends {}; + +TEST(RttiNcTest, MissingId) { + // expected-error@* {{ThisT must define its own ID}} + UninstantiatedWithoutId::classID(); +} + } // namespace } // namespace ifrt diff --git a/third_party/xla/xla/python/ifrt/sharding_spec.cc b/third_party/xla/xla/python/ifrt/sharding_spec.cc index a0d3b02d82f454..802da07f3dfb45 100644 --- a/third_party/xla/xla/python/ifrt/sharding_spec.cc +++ b/third_party/xla/xla/python/ifrt/sharding_spec.cc @@ -785,9 +785,9 @@ absl::StatusOr ShardingParamShardingSpec::GetShardShape( const Shape& shape) const { if (shape.dims().size() != sharding_param_.dim_shards().size()) { return absl::InvalidArgumentError(absl::StrFormat( - "Numbers of dimensions don't match. From Shape %d vs from " - "ShardingParam %d", - shape.dims().size(), sharding_param_.dim_shards().size())); + "Numbers of dimensions don't match. From Shape %v vs from " + "ShardingParam %s", + shape, sharding_param_.DebugString())); } std::vector dims; dims.reserve(shape.dims().size()); diff --git a/third_party/xla/xla/python/ifrt/sharding_spec_test.cc b/third_party/xla/xla/python/ifrt/sharding_spec_test.cc index 86d45252180f5b..bf485336a60590 100644 --- a/third_party/xla/xla/python/ifrt/sharding_spec_test.cc +++ b/third_party/xla/xla/python/ifrt/sharding_spec_test.cc @@ -650,7 +650,8 @@ TEST_P(ShardingParamShardingSpecTest, GetShardShape) { absl_testing::StatusIs( tsl::error::INVALID_ARGUMENT, HasSubstr("Numbers of dimensions don't match. From " - "Shape 3 vs from ShardingParam 2"))); + "Shape [6,6,6] vs from ShardingParam " + "2x3 to [1, 0] on 3x2"))); } TEST_P(ShardingParamShardingSpecTest, HasSamePartitioning) { @@ -707,7 +708,8 @@ TEST_P(ShardingParamShardingSpecTest, DisassembleFailsWhenRankNotMatch) { absl_testing::StatusIs( tsl::error::INVALID_ARGUMENT, HasSubstr("Numbers of dimensions don't match. From " - "Shape 3 vs from ShardingParam 2"))); + "Shape [6,6,6] vs from ShardingParam " + "2x3 to [1, 0] on 3x2"))); } TEST_P(ShardingParamShardingSpecTest, DisassembleFailsForUnevenSharding) { diff --git a/third_party/xla/xla/python/ifrt/sharding_test.cc b/third_party/xla/xla/python/ifrt/sharding_test.cc index 0c403771a16579..b7a539084d321f 100644 --- a/third_party/xla/xla/python/ifrt/sharding_test.cc +++ b/third_party/xla/xla/python/ifrt/sharding_test.cc @@ -1061,7 +1061,8 @@ TEST_P(ShardingParamShardingTest, GetShardShape) { EXPECT_THAT(sharding->GetShardShape(Shape({6, 6, 6})), StatusIs(tsl::error::INVALID_ARGUMENT, HasSubstr("Numbers of dimensions don't match. From " - "Shape 3 vs from ShardingParam 2"))); + "Shape [6,6,6] vs from ShardingParam " + "2x3 to [1, 0] on 3x2"))); } TEST_P(ShardingParamShardingTest, GetShardShapeWithUnreducedAxes) { @@ -1077,7 +1078,8 @@ TEST_P(ShardingParamShardingTest, GetShardShapeWithUnreducedAxes) { EXPECT_THAT(sharding->GetShardShape(Shape({6, 6, 6})), StatusIs(tsl::error::INVALID_ARGUMENT, HasSubstr("Numbers of dimensions don't match. From " - "Shape 3 vs from ShardingParam 2"))); + "Shape [6,6,6] vs from ShardingParam " + "2x1 to [1, 0] on 3x2 unreduced [1]"))); } TEST_P(ShardingParamShardingTest, HasSamePartitioning) { @@ -1201,7 +1203,8 @@ TEST_P(ShardingParamShardingTest, DisassembleFailsWhenRankNotMatch) { Shape({6, 6, 6}), SingleDeviceShardSemantics::kAllShards), StatusIs(tsl::error::INVALID_ARGUMENT, HasSubstr("Numbers of dimensions don't match. From " - "Shape 3 vs from ShardingParam 2"))); + "Shape [6,6,6] vs from ShardingParam " + "2x3 to [1, 0] on 3x2"))); } TEST_P(ShardingParamShardingTest, DisassembleFailsForUnevenSharding) { diff --git a/third_party/xla/xla/python/ifrt/user_context_test_util.h b/third_party/xla/xla/python/ifrt/user_context_test_util.h index 63c381c82ce42a..f1ff2e910ca6e4 100644 --- a/third_party/xla/xla/python/ifrt/user_context_test_util.h +++ b/third_party/xla/xla/python/ifrt/user_context_test_util.h @@ -42,7 +42,7 @@ class TestUserContext : public RTTIExtends { std::string DebugString() const override { return debug_string_; } - // No new `ID` is not defined because tests below do not exercise RTTI. + inline static char ID = 0; // NOLINT private: explicit TestUserContext(UserContextId id, std::string debug_string) diff --git a/third_party/xla/xla/python/ifrt_proxy/client/memory.h b/third_party/xla/xla/python/ifrt_proxy/client/memory.h index 1e9c696fd9b551..7480583d25d332 100644 --- a/third_party/xla/xla/python/ifrt_proxy/client/memory.h +++ b/third_party/xla/xla/python/ifrt_proxy/client/memory.h @@ -58,6 +58,8 @@ class Memory : public RTTIExtends { absl::string_view DebugString() const override { return debug_string_; } absl::string_view ToString() const override { return to_string_; } + inline static char ID = 0; // NOLINT + private: friend class Client; // For `devices_` initialization. diff --git a/third_party/xla/xla/python/ifrt_proxy/server/host_callback.cc b/third_party/xla/xla/python/ifrt_proxy/server/host_callback.cc index 5786914354ca74..9b07d2894530e6 100644 --- a/third_party/xla/xla/python/ifrt_proxy/server/host_callback.cc +++ b/third_party/xla/xla/python/ifrt_proxy/server/host_callback.cc @@ -143,6 +143,8 @@ RemoteLoadedHostCallback::~RemoteLoadedHostCallback() { } } +char RemoteLoadedHostCallback::ID = 0; + absl::Status RemoteLoadedHostCallback::Execute(void** result_ptrs, void** operand_ptrs) { if (queue_ == nullptr) { diff --git a/third_party/xla/xla/python/ifrt_proxy/server/host_callback.h b/third_party/xla/xla/python/ifrt_proxy/server/host_callback.h index bd9f76f8e683a0..6c6f870f03aad4 100644 --- a/third_party/xla/xla/python/ifrt_proxy/server/host_callback.h +++ b/third_party/xla/xla/python/ifrt_proxy/server/host_callback.h @@ -112,6 +112,8 @@ class RemoteLoadedHostCallback // deserialized into `RmeoteLoadedHostCallback` using `CreateFromSerialized`. absl::StatusOr Serialize() const override; + static char ID; // NOLINT + private: // Implements the interface required by `xla::HostCallback`. absl::Status Execute(void** result_ptrs, void** operand_ptrs); diff --git a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec.cc b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec.cc index d92d8a643426c6..d55e2e2d866013 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec.cc +++ b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec.cc @@ -94,7 +94,8 @@ std::unique_ptr HloShardingSpec::Create( CHECK_EQ(num_shards, xla_hlo_sharding.num_devices()) << "`num_shards` and `xla_hlo_sharding`'s `num_devices` does not " "match: " - << num_shards << " vs. " << xla_hlo_sharding.num_devices(); + << num_shards << " vs. " << xla_hlo_sharding.num_devices() + << "; sharding=" << xla_hlo_sharding.ToString(); } return std::unique_ptr( new HloShardingSpec(num_shards, std::move(xla_hlo_sharding))); @@ -120,8 +121,9 @@ absl::StatusOr HloShardingSpec::ToSharding( DeviceListRef devices, MemoryKind memory_kind) const { if (devices->size() != num_shards()) { return absl::InvalidArgumentError(absl::StrFormat( - "HloShardingSpec requires %d devices, but received %d devices", - num_shards(), devices->size())); + "HloShardingSpec requires %d devices, but received %d devices; " + "sharding=%s", + num_shards(), devices->size(), xla_hlo_sharding_.ToString())); } std::shared_ptr spec = std::static_pointer_cast(weak_from_this().lock()); @@ -140,9 +142,9 @@ absl::StatusOr HloShardingSpec::GetShardShape(const Shape& shape) const { } if (shape.dims().size() != xla_hlo_sharding_.TiledDataRank()) { return absl::InvalidArgumentError(absl::StrFormat( - "Numbers of dimensions don't match. From Shape %d vs from " - "HloSharding %d", - shape.dims().size(), xla_hlo_sharding_.TiledDataRank())); + "Numbers of dimensions don't match. From Shape %v vs from " + "HloSharding %s", + shape, xla_hlo_sharding_.ToString())); } const absl::Span sharding_dims = xla_hlo_sharding_.dimensions(); @@ -236,11 +238,15 @@ absl::StatusOr> HloShardingSpec::IndexDomains( if (xla_hlo_sharding_.IsManual()) { return absl::InvalidArgumentError( - "Manual sharding does not support IndexDomains"); + absl::StrFormat("Manual sharding does not support IndexDomains: " + "sharding=%s", + xla_hlo_sharding_.ToString())); } if (xla_hlo_sharding_.IsUnreduced()) { return absl::InvalidArgumentError( - "Unreduced sharding does not support IndexDomains"); + absl::StrFormat("Unreduced sharding does not support IndexDomains: " + "sharding=%s", + xla_hlo_sharding_.ToString())); } if (xla_hlo_sharding_.IsReplicatedOrSingleDevice()) { // Fast path for a fully replicated or maximal sharding. @@ -293,16 +299,20 @@ absl::StatusOr> HloShardingSpec::UniqueIndexDomains(const Shape& shape) const { if (xla_hlo_sharding_.IsManual()) { return absl::InvalidArgumentError( - "Manual sharding does not support UniqueIndexDomains"); + absl::StrFormat("Manual sharding does not support UniqueIndexDomains: " + "sharding=%s", + xla_hlo_sharding_.ToString())); } if (xla_hlo_sharding_.IsUnreduced()) { - return absl::InvalidArgumentError( - "Unreduced sharding does not support UniqueIndexDomains"); + return absl::InvalidArgumentError(absl::StrFormat( + "Unreduced sharding does not support UniqueIndexDomains: sharding=%s", + xla_hlo_sharding_.ToString())); } if (xla_hlo_sharding_.HasNonReplicatedSubgroup()) { - return absl::InvalidArgumentError( + return absl::InvalidArgumentError(absl::StrFormat( "Non-replicated subgroup (e.g., manual or unreduced subgroup) sharding " - "does not support UniqueIndexDomains"); + "does not support UniqueIndexDomains: sharding=%s", + xla_hlo_sharding_.ToString())); } if (xla_hlo_sharding_.IsReplicatedOrSingleDevice()) { absl::call_once(unique_shard_indices_once_, [this] { @@ -342,8 +352,8 @@ HloShardingSpec::UniqueIndexDomains(const Shape& shape) const { if (num_shards_ % num_unique_tiles != 0) { return absl::InvalidArgumentError(absl::StrFormat( "HloShardingSpec has %d shards, but HloSharding has %d unique tiles, " - "which is not a divisor of the number of shards", - num_shards_, num_unique_tiles)); + "which is not a divisor of the number of shards; sharding=%s", + num_shards_, num_unique_tiles, xla_hlo_sharding_.ToString())); } const int64_t num_replicas = num_shards_ / num_unique_tiles; @@ -379,16 +389,21 @@ absl::StatusOr> HloShardingSpec::ShardToUniqueIndexDomainIndex() const { if (xla_hlo_sharding_.IsManual()) { return absl::InvalidArgumentError( - "Manual sharding does not support ShardToUniqueIndexDomainIndex"); + absl::StrFormat("Manual sharding does not support " + "ShardToUniqueIndexDomainIndex: sharding=%s", + xla_hlo_sharding_.ToString())); } if (xla_hlo_sharding_.IsUnreduced()) { return absl::InvalidArgumentError( - "Unreduced sharding does not support ShardToUniqueIndexDomainIndex"); + absl::StrFormat("Unreduced sharding does not support " + "ShardToUniqueIndexDomainIndex: sharding=%s", + xla_hlo_sharding_.ToString())); } if (xla_hlo_sharding_.HasNonReplicatedSubgroup()) { - return absl::InvalidArgumentError( + return absl::InvalidArgumentError(absl::StrFormat( "Non-replicated subgroup (e.g., manual or unreduced subgroup) sharding " - "does not support ShardToUniqueIndexDomainIndex"); + "does not support ShardToUniqueIndexDomainIndex: sharding=%s", + xla_hlo_sharding_.ToString())); } if (xla_hlo_sharding_.IsReplicatedOrSingleDevice()) { absl::call_once(shard_to_unique_index_domain_index_once_, [this] { @@ -401,8 +416,8 @@ HloShardingSpec::ShardToUniqueIndexDomainIndex() const { if (num_shards_ % num_unique_tiles != 0) { return absl::InvalidArgumentError(absl::StrFormat( "HloShardingSpec has %d shards, but HloSharding has %d unique tiles, " - "which is not a divisor of the number of shards", - num_shards_, num_unique_tiles)); + "which is not a divisor of the number of shards; sharding=%s", + num_shards_, num_unique_tiles, xla_hlo_sharding_.ToString())); } const int64_t num_replicas = num_shards_ / num_unique_tiles; diff --git a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec_test.cc b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec_test.cc index 214bff8445f9e3..042fa2e78a6d7b 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec_test.cc +++ b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_spec_test.cc @@ -152,7 +152,8 @@ TEST_F(HloShardingSpecTest, GetShardShape) { absl_testing::StatusIs( tsl::error::INVALID_ARGUMENT, HasSubstr("Numbers of dimensions don't match. From " - "Shape 3 vs from HloSharding 2"))); + "Shape [6,6,6] vs from HloSharding " + "{devices=[2,3]<=[6]}"))); } TEST_F(HloShardingSpecTest, HasSamePartitioning) { diff --git a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_test.cc b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_test.cc index d22c8afd621538..3c4e99af150a2c 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_test.cc +++ b/third_party/xla/xla/python/pjrt_ifrt/xla_sharding_test.cc @@ -169,7 +169,8 @@ TEST_P(HloShardingTest, GetShardShape) { absl_testing::StatusIs( tsl::error::INVALID_ARGUMENT, HasSubstr("Numbers of dimensions don't match. From " - "Shape 3 vs from HloSharding 2"))); + "Shape [6,6,6] vs from HloSharding " + "{devices=[2,3]<=[6]}"))); } TEST_P(HloShardingTest, HasSamePartitioning) { diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD index 39c41bbce90a62..8ae6293d347548 100644 --- a/third_party/xla/xla/service/BUILD +++ b/third_party/xla/xla/service/BUILD @@ -3215,6 +3215,7 @@ cc_library( "//xla/tsl/platform:statusor", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/log", + "@com_google_absl//absl/log:check", "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", diff --git a/third_party/xla/xla/service/buffer_assignment.cc b/third_party/xla/xla/service/buffer_assignment.cc index f78ac0da41006f..aa650b8502f043 100644 --- a/third_party/xla/xla/service/buffer_assignment.cc +++ b/third_party/xla/xla/service/buffer_assignment.cc @@ -2862,6 +2862,7 @@ absl::Status BufferAssigner::AssignBuffersWithSequentialOrdering( int64_t alignment = assignment->color_alignment_(color); HeapSimulator::Options options; options.alloc_constants = opts_.allocate_buffers_for_constants; + options.view_color = opts_.dus_view_color; auto private_stacks_it = private_stacks.find(color); if (private_stacks_it != private_stacks.end()) { // For private stack colors, we collect all of the buffers that are @@ -2925,6 +2926,7 @@ absl::Status BufferAssigner::AssignBuffersWithSequentialOrdering( int64_t alignment = assignment->color_alignment_(color); HeapSimulator::Options options; options.buffers_to_assign = &color_map[color]; + options.view_color = opts_.dus_view_color; HeapSimulator::Result result; ABSL_ASSIGN_OR_RETURN( result, HeapSimulator::Run( @@ -3191,6 +3193,15 @@ BufferAssigner::CreateAssignment( HloLiveRange::Run(schedule, *alias_analysis, module->entry_computation(), true)); + // A view base's storage is read through the view by the view's consumers, + // so its live range must reach the last transitive reader; allocation reuse + // decisions below consult these live ranges. + if (opts_.dus_view_color.has_value()) { + ExtendViewBaseLiveRanges(hlo_live_range.get(), + alias_analysis->dataflow_analysis(), + *opts_.dus_view_color); + } + VLOG(1) << "Assigning buffers to module " << module->name(); XLA_VLOG_LINES(3, module->ToString()); XLA_VLOG_LINES(3, alias_analysis->ToString()); diff --git a/third_party/xla/xla/service/buffer_assignment_test.cc b/third_party/xla/xla/service/buffer_assignment_test.cc index 1c79ec2dbc3705..78e530bea29e66 100644 --- a/third_party/xla/xla/service/buffer_assignment_test.cc +++ b/third_party/xla/xla/service/buffer_assignment_test.cc @@ -241,6 +241,35 @@ class BufferAssignmentTest : public HloHardwareIndependentTestBase { return std::move(assignment).value(); } + // Like RunBufferAssignmentWithDusViewColor, but with the module's own + // sequential schedule, so temp buffers are packed by the heap simulator + // (which is where view base live range extension matters). + std::unique_ptr + RunSequentialBufferAssignmentWithDusViewColor( + HloModule* module, absl::string_view view_instruction_name, + BufferValue::Color view_color, + std::optional dus_view_color) { + BufferAssigner::Options opts; + opts.allocate_buffers_for_constants = true; + opts.dus_view_color = dus_view_color; + opts.colorer = [name = std::string(view_instruction_name), view_color]( + HloAliasAnalysis* alias_analysis, const HloOrdering&) { + for (HloValue* value : alias_analysis->dataflow_analysis().values()) { + value->set_color(value->instruction()->name() == name + ? view_color + : BufferValue::Color(0)); + } + return absl::OkStatus(); + }; + absl::StatusOr> assignment = + BufferAssigner::Run( + module, std::make_unique(module->schedule()), + &BufferSizeBytes, &alias_info_, + [](LogicalBuffer::Color) { return 1; }, std::move(opts)); + CHECK_OK(assignment.status()); + return std::move(assignment).value(); + } + std::unique_ptr RunBufferAssignmentWithInstructionSequence( HloModule* module, absl::Span instruction_sequence, int64_t alignment = 1, @@ -3027,6 +3056,52 @@ ENTRY e { EXPECT_TRUE(assignment->HasTopLevelAllocation(view)); } +// A view is an address into its base's buffer, so consumers of the view read +// the BASE's storage at their own (later) schedule times. The heap simulator +// must keep the base buffer reserved until the view's last transitive reader; +// without that extension the base's slot is recycled for an unrelated temp +// defined between the view and the reader, and the reader loads the temp's +// bytes. +TEST_F(BufferAssignmentTest, DusViewBaseLiveRangeExtendsToViewReaders) { + const char* const kHlo = R"( +HloModule DusViewBaseLiveRange, is_scheduled=true + +ENTRY e { + p0 = f32[64]{0} parameter(0) + base = f32[64]{0} negate(p0) + view = f32[64]{0:S(7)} custom-call(base), custom_call_target="view" + filler = f32[64]{0} exponential(p0) + ROOT out = f32[64]{0} add(view, filler) +} +)"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(kHlo)); + const HloInstruction* base = FindInstruction(module.get(), "base"); + const HloInstruction* filler = FindInstruction(module.get(), "filler"); + + constexpr BufferValue::Color kViewColor = 7; + std::unique_ptr assignment = + RunSequentialBufferAssignmentWithDusViewColor(module.get(), "view", + kViewColor, kViewColor); + + // `base` is only directly used by the view, while `out` reads base's + // storage through the view two schedule steps later. `filler`, live in + // between, must not be packed on top of base's bytes. + ASSERT_OK_AND_ASSIGN(BufferAllocation::Slice base_slice, + assignment->GetUniqueTopLevelSlice(base)); + ASSERT_OK_AND_ASSIGN(BufferAllocation::Slice filler_slice, + assignment->GetUniqueTopLevelSlice(filler)); + const bool same_allocation = + base_slice.allocation() == filler_slice.allocation(); + const bool bytes_overlap = + same_allocation && + base_slice.offset() < filler_slice.offset() + filler_slice.size() && + filler_slice.offset() < base_slice.offset() + base_slice.size(); + EXPECT_FALSE(bytes_overlap) + << "view base was recycled while still read through the view: base " + << base_slice.ToString() << " vs filler " << filler_slice.ToString(); +} + class WhileBufferAssignmentTest : public HloHardwareIndependentTestBase { protected: std::unique_ptr BuildWhileConditionComputation( diff --git a/third_party/xla/xla/service/collective_ops_utils.cc b/third_party/xla/xla/service/collective_ops_utils.cc index d84d7e990093c2..6d4b12090534d0 100644 --- a/third_party/xla/xla/service/collective_ops_utils.cc +++ b/third_party/xla/xla/service/collective_ops_utils.cc @@ -1029,8 +1029,6 @@ bool IsOneShotRaggedAllToAllWithNcclEnabled(const DebugOptions& opts) { return opts.xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl(); } -namespace { - std::optional GetCollectiveOpType( const HloInstruction* instruction) { if (!IsNonFusionCollective(instruction)) { @@ -1068,8 +1066,6 @@ std::optional GetCollectiveOpType( } } -} // namespace - NcclSymmetricBuffersSpec::NcclSymmetricBuffersSpec( const DebugOptions& debug_options) { for (const auto& filter_proto : diff --git a/third_party/xla/xla/service/collective_ops_utils.h b/third_party/xla/xla/service/collective_ops_utils.h index 72b2aa73505866..65419f60f2d2ee 100644 --- a/third_party/xla/xla/service/collective_ops_utils.h +++ b/third_party/xla/xla/service/collective_ops_utils.h @@ -324,6 +324,11 @@ class NcclSymmetricBuffersSpec { bool IsNcclSymmetricBuffersEnabledForCollective( const HloInstruction* instruction, const DebugOptions& opts); +// Returns the CollectiveOpType corresponding to the given instruction, or +// std::nullopt if the instruction is not a recognized collective operation. +std::optional GetCollectiveOpType( + const HloInstruction* instruction); + //===----------------------------------------------------------------------===// // Async collective configuration. //===----------------------------------------------------------------------===// diff --git a/third_party/xla/xla/service/collective_opt_utils.cc b/third_party/xla/xla/service/collective_opt_utils.cc index 7bd7c4153712b9..37b13f35e3daa4 100644 --- a/third_party/xla/xla/service/collective_opt_utils.cc +++ b/third_party/xla/xla/service/collective_opt_utils.cc @@ -76,8 +76,10 @@ bool IsTableLookup(const HloInstruction* hlo) { } std::optional GetScalarInt64Value(const HloInstruction* constant) { - CHECK_EQ(constant->opcode(), HloOpcode::kConstant); - CHECK(ShapeUtil::IsEffectiveScalar(constant->shape())); + if (constant == nullptr || constant->opcode() != HloOpcode::kConstant || + !ShapeUtil::IsEffectiveScalar(constant->shape())) { + return std::nullopt; + } absl::InlinedVector multi_index( constant->shape().dimensions().size()); return constant->literal().GetIntegralAsS64(multi_index); @@ -238,8 +240,10 @@ const HloInstruction* BacktrackToBase( if (cur->opcode() == HloOpcode::kClamp) { // For some clamp ops it's possible to prove they are no-ops at compile // time. - std::optional lower_bound = GetScalarInt64Value(cur->operand(0)); - std::optional upper_bound = GetScalarInt64Value(cur->operand(2)); + std::optional lower_bound = + GetScalarInt64Value(BacktrackToBase(cur->operand(0), cache)); + std::optional upper_bound = + GetScalarInt64Value(BacktrackToBase(cur->operand(2), cache)); std::optional> range = GetKnownRange(cur->operand(1)); if (lower_bound.has_value() && upper_bound.has_value() && @@ -297,8 +301,8 @@ std::optional EvaluateOffset( if (!inner) { return std::nullopt; } - auto lower = GetScalarInt64Value(expr->operand(0)); - auto upper = GetScalarInt64Value(expr->operand(2)); + auto lower = GetScalarInt64Value(BacktrackToBase(expr->operand(0), cache)); + auto upper = GetScalarInt64Value(BacktrackToBase(expr->operand(2), cache)); if (!lower || !upper) { return std::nullopt; } diff --git a/third_party/xla/xla/service/collective_opt_utils_test.cc b/third_party/xla/xla/service/collective_opt_utils_test.cc index 6ab42c464c2b4a..95ba1576be67a5 100644 --- a/third_party/xla/xla/service/collective_opt_utils_test.cc +++ b/third_party/xla/xla/service/collective_opt_utils_test.cc @@ -961,5 +961,85 @@ TEST_F(MatchDsPadAllGatherTest, MatchDsPadAllGatherNoAllGather) { EXPECT_FALSE(MatchDsPadAllGather(root, &pad_hlo, &ag_hlo)); } +using MatchWithDynamicSliceClampReshapeTest = HloHardwareIndependentTestBase; + +TEST_F(MatchWithDynamicSliceClampReshapeTest, ClampWithReshapedConstantBounds) { + constexpr absl::string_view hlo_string = R"( + HloModule module + + ENTRY entry { + param = f32[16,10] parameter(0) + ag = f32[32,10] all-gather(param), dimensions={0}, + replica_groups={{0,1}} + pid = u32[] replica-id() + const_16 = u32[] constant(16) + mul = u32[] multiply(pid, const_16) + lower_const = u32[1] constant({0}) + lower_bound = u32[] reshape(lower_const) + upper_const = u32[1] constant({16}) + upper_bound = u32[] reshape(upper_const) + clamped_offset = u32[] clamp(lower_bound, mul, upper_bound) + zero = s32[] constant(0) + ROOT ds = f32[16,10] dynamic-slice(ag, clamped_offset, zero), + dynamic_slice_sizes={16,10} + } + )"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnUnverifiedModule(hlo_string)); + const HloInstruction* ag = + module->entry_computation()->GetInstructionWithName("ag"); + const HloChannelInstruction* ag_instr = Cast(ag); + + std::optional spec = MatchWithDynamicSlice( + ag_instr, /*num_partitions=*/1, /*num_replicas=*/2, + /*allow_multiple_split_dims=*/false, + /*allow_intervening_reshape=*/false, /*min_rank=*/0, + HloPredicateIsOp, + HloPredicateIsOp, + /*is_constrain_layout=*/false, /*use_global_device_ids=*/false, + /*is_cross_module=*/false, /*allow_intervening_bitcast=*/false, + /*allow_multiple_users=*/false); + + EXPECT_TRUE(spec.has_value()); +} + +TEST_F(MatchWithDynamicSliceClampReshapeTest, ClampWithNonConstantBounds) { + constexpr absl::string_view hlo_string = R"( + HloModule module + + ENTRY entry { + param = f32[16,10] parameter(0) + lower_param = u32[] parameter(1) + upper_param = u32[] parameter(2) + ag = f32[32,10] all-gather(param), dimensions={0}, + replica_groups={{0,1}} + pid = u32[] replica-id() + const_16 = u32[] constant(16) + mul = u32[] multiply(pid, const_16) + clamped_offset = u32[] clamp(lower_param, mul, upper_param) + zero = s32[] constant(0) + ROOT ds = f32[16,10] dynamic-slice(ag, clamped_offset, zero), + dynamic_slice_sizes={16,10} + } + )"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnUnverifiedModule(hlo_string)); + const HloInstruction* ag = + module->entry_computation()->GetInstructionWithName("ag"); + const HloChannelInstruction* ag_instr = Cast(ag); + + std::optional spec = MatchWithDynamicSlice( + ag_instr, /*num_partitions=*/1, /*num_replicas=*/2, + /*allow_multiple_split_dims=*/false, + /*allow_intervening_reshape=*/false, /*min_rank=*/0, + HloPredicateIsOp, + HloPredicateIsOp, + /*is_constrain_layout=*/false, /*use_global_device_ids=*/false, + /*is_cross_module=*/false, /*allow_intervening_bitcast=*/false, + /*allow_multiple_users=*/false); + + EXPECT_FALSE(spec.has_value()); +} + } // namespace } // namespace xla diff --git a/third_party/xla/xla/service/conditional_code_motion.cc b/third_party/xla/xla/service/conditional_code_motion.cc index 8262065d6d1579..d83f03394f085f 100644 --- a/third_party/xla/xla/service/conditional_code_motion.cc +++ b/third_party/xla/xla/service/conditional_code_motion.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include #include +#include #include #include #include @@ -43,6 +44,8 @@ limitations under the License. #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_instructions.h" #include "xla/hlo/ir/hlo_opcode.h" +#include "xla/hlo/ir/hlo_original_value.h" +#include "xla/hlo/ir/hlo_original_value_util.h" #include "xla/hlo/ir/hlo_print_options.h" #include "xla/hlo/pass/hlo_pass_pipeline.h" #include "xla/hlo/transforms/simplifiers/hlo_dce.h" @@ -59,6 +62,39 @@ namespace xla { namespace conditional_opt { +void UpdateInstructionOriginalValue(HloInstruction* instruction, + HloInstruction* operand_source = nullptr) { + if (instruction->original_value() == nullptr && operand_source == nullptr) { + return; + } + if (operand_source != nullptr) { + CopyOriginalValue(operand_source, instruction, /*clone=*/true, + /*issue_warning=*/false); + return; + } + if (instruction->opcode() == HloOpcode::kTuple || + instruction->opcode() == HloOpcode::kGetTupleElement) { + instruction->set_original_value( + OriginalValue::CreateFromInstruction(instruction)); + return; + } + if (!instruction->original_value()->IsTuple()) { + auto old_ov = instruction->original_value(); + auto new_ov = std::make_shared(instruction->shape()); + new_ov->mutable_tree()->CopySubtreeFrom(old_ov->tree(), {}, {0}); + instruction->set_original_value(new_ov); + return; + } + absl::flat_hash_map old_to_new_map; + for (const auto& [shape_index, _] : + instruction->original_value()->original_arrays()) { + if (!shape_index.empty()) { + old_to_new_map[shape_index[0]] = shape_index[0]; + } + } + CopyOriginalValue(instruction, instruction, old_to_new_map); +} + HloInstruction* CloneNestedTuples(HloInstruction* tuple) { if (!tuple->shape().IsTuple()) { return tuple; @@ -788,6 +824,7 @@ absl::StatusOr ConditionalCodeMotion::MoveInstructionOut( HloInstruction* new_root = conditional->branch_computation(0)->root_instruction(); *conditional->mutable_shape() = new_root->shape(); + conditional_opt::UpdateInstructionOriginalValue(conditional); // Keep conditional instruction sharding consistent with the branches. Note // that this sharding could be lost after this pass. conditional->copy_sharding(new_root); @@ -920,6 +957,7 @@ absl::StatusOr ConditionalCodeMotion::MoveUserInstructionsIn( HloInstruction* new_root = computation->root_instruction(); new_root->AppendOperand(new_op); *new_root->mutable_shape()->add_tuple_shapes() = new_op->shape(); + conditional_opt::UpdateInstructionOriginalValue(new_root); VLOG(2) << "Extending conditional root " << i << " : " << new_root->ToString() << "\n"; } @@ -939,6 +977,7 @@ absl::StatusOr ConditionalCodeMotion::MoveUserInstructionsIn( HloInstruction* new_root = conditional->branch_computation(0)->root_instruction(); *conditional->mutable_shape() = new_root->shape(); + conditional_opt::UpdateInstructionOriginalValue(conditional); // Keep conditional instruction sharding consistent with the branches. Note // that this sharding could be lost after this pass. conditional->copy_sharding(new_root); @@ -949,6 +988,7 @@ absl::StatusOr ConditionalCodeMotion::MoveUserInstructionsIn( VLOG(2) << "Resetting shape of user: " << user->ToString() << "\n"; *user->mutable_shape() = conditional->shape().tuple_shapes(user->tuple_index()); + conditional_opt::UpdateInstructionOriginalValue(user); } } } @@ -1016,7 +1056,7 @@ class MoveOperandIntoBranch { bool UpdateParamShape( std::vector>& matching_tuple_indices, const Shape* param_shape, HloInstruction*& branch_param, - HloInstruction*& param_tuple) { + HloInstruction*& param_tuple, HloInstruction* user = nullptr) { bool used = false; // Update shapes from branch parameter to the immediate tuple user. for (int64_t matching_index = matching_tuple_indices.size() - 1; @@ -1038,6 +1078,7 @@ class MoveOperandIntoBranch { } used = false; *branch_param->mutable_shape() = *param_shape; + conditional_opt::UpdateInstructionOriginalValue(branch_param, user); const Shape* new_param_shape = nullptr; for (auto param_users : gte_users) { if (param_users.empty()) { @@ -1065,11 +1106,13 @@ class MoveOperandIntoBranch { param_shape = new_param_shape; VLOG(1) << "new_param_shape: " << param_shape->ToString(); *param_user->mutable_shape() = *new_param_shape; + conditional_opt::UpdateInstructionOriginalValue(param_user); VLOG(1) << "branch parameter: " << param_user->ToString(); used = true; } else { VLOG(1) << "new_param_shape=" << new_param_shape->ToString(); *param_user->mutable_shape() = *new_param_shape; + conditional_opt::UpdateInstructionOriginalValue(param_user); CHECK_OK(param_user->ReplaceAllUsesWith(branch_param)); } } @@ -1132,6 +1175,7 @@ class MoveOperandIntoBranch { user->mutable_shape()->mutable_tuple_shapes()->push_back( new_input->shape()); } + conditional_opt::UpdateInstructionOriginalValue(user); int64_t nesting_index = 1; for (auto user_now = user->users()[0]; nesting_index < matching_tuple_indices.size() && @@ -1143,6 +1187,7 @@ class MoveOperandIntoBranch { *user_now->mutable_shape()->mutable_tuple_shapes(opd_index) = user->shape(); } + conditional_opt::UpdateInstructionOriginalValue(user_now); VLOG(2) << "Done replacing tuple:" << user->ToString(); CHECK_EQ(user_now->user_count(), 1); } @@ -1161,6 +1206,8 @@ class MoveOperandIntoBranch { if (user == cond) { auto new_input = input->AddInstruction(HloInstruction::CreateTuple(new_operands)); + new_input->set_original_value( + OriginalValue::CreateFromInstruction(new_input)); for (int64_t i = 0; i < new_operands.size(); ++i) { op_map_[new_operands[i]] = i; } @@ -1186,6 +1233,7 @@ class MoveOperandIntoBranch { "directly."; VLOG(5) << branch_comp->ToString() << "\n"; *branch_param->mutable_shape() = *param_shape; + conditional_opt::UpdateInstructionOriginalValue(branch_param, user); if (branch_param == branch_comp->root_instruction()) { VLOG(2) << "Cloning root user"; auto new_user = @@ -1197,7 +1245,7 @@ class MoveOperandIntoBranch { } } else { if (!UpdateParamShape(matching_tuple_indices, param_shape, branch_param, - param_tuple)) { + param_tuple, user)) { VLOG(2) << "instruction is not used in this branch."; continue; } diff --git a/third_party/xla/xla/service/conditional_code_motion_test.cc b/third_party/xla/xla/service/conditional_code_motion_test.cc index 2136e4f5dc3899..ca348f572facad 100644 --- a/third_party/xla/xla/service/conditional_code_motion_test.cc +++ b/third_party/xla/xla/service/conditional_code_motion_test.cc @@ -30,6 +30,7 @@ limitations under the License. #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_opcode.h" +#include "xla/hlo/ir/hlo_original_value.h" #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" #include "xla/hlo/testlib/test.h" #include "xla/hlo/transforms/simplifiers/hlo_dce.h" @@ -2639,6 +2640,49 @@ ENTRY main { } } +TEST_F(ConditionalCodeMotionTest, OriginalValuePreservedOnMoveIn) { + absl::string_view hlo_string = R"( +HloModule TestModule + +on_true { + arg_tuple.1 = (f32[10]) parameter(0) + get-tuple-element.1 = f32[10] get-tuple-element(arg_tuple.1), index=0 + add.1 = f32[10] add(get-tuple-element.1, get-tuple-element.1) + ROOT tuple.3 = (f32[10]) tuple(add.1) +} + +on_false { + arg_tuple.2 = (f32[10]) parameter(0) + get-tuple-element.2 = f32[10] get-tuple-element(arg_tuple.2), index=0 + mul.1 = f32[10] multiply(get-tuple-element.2, get-tuple-element.2) + ROOT tuple.4 = (f32[10]) tuple(mul.1) +} + +ENTRY main { + pred.1 = pred[] parameter(0) + tuple.1 = (f32[10]) parameter(1) + tuple.2 = (f32[10]) parameter(2) + conditional = (f32[10]) conditional(pred.1, tuple.1, tuple.2), + true_computation=on_true, false_computation=on_false + get-first-index = f32[10] get-tuple-element(conditional), index=0 + ROOT pow.1 = f32[10] power(get-first-index, get-first-index) +} +)"; + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + for (HloComputation* comp : module->computations()) { + for (HloInstruction* inst : comp->instructions()) { + inst->set_original_value(OriginalValue::CreateFromInstruction(inst)); + } + } + + ConditionalCodeMotion pass(true, true); + ASSERT_OK_AND_ASSIGN(bool changed, pass.Run(module.get())); + EXPECT_TRUE(changed); + + EXPECT_OK(verifier().Run(module.get()).status()); +} + } // namespace conditional_opt } // namespace xla diff --git a/third_party/xla/xla/service/cpu/BUILD b/third_party/xla/xla/service/cpu/BUILD index 99dc48e2494892..7e00afbb5ead37 100644 --- a/third_party/xla/xla/service/cpu/BUILD +++ b/third_party/xla/xla/service/cpu/BUILD @@ -193,6 +193,7 @@ cc_library( "//xla/backends/cpu/runtime:function_library", "//xla/backends/cpu/runtime:kernel_thunk", "//xla/backends/cpu/runtime:thunk", + "//xla/backends/cpu/runtime:thunk_proto_cc", "//xla/backends/cpu/runtime:thunk_proto_cc_impl", "//xla/backends/cpu/runtime:thunk_proto_serdes", "//xla/backends/cpu/transforms:library_rewriter", @@ -321,6 +322,7 @@ cc_library( "//xla/tsl/concurrency:executor", "//xla/tsl/platform:env", "//xla/tsl/platform:errors", + "//xla/tsl/platform:logging", "//xla/tsl/platform:status", "//xla/tsl/platform:statusor", "//xla/tsl/protobuf:error_codes_proto_impl_cc", @@ -701,7 +703,6 @@ cc_library( hdrs = ["ir_emitter.h"], copts = tsl_copts(), deps = [ - ":backend_config_proto_cc", ":cpu_instruction_fusion", ":cpu_options", ":cpu_runtime", @@ -773,14 +774,12 @@ cc_library( srcs = ["ir_function.cc"], hdrs = ["ir_function.h"], deps = [ - ":cpu_runtime", ":ir_emission_utils", "//xla:shape_util", "//xla:status_macros", "//xla:types", "//xla/service:hlo_module_config", "//xla/service/llvm_ir:llvm_util", - "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", @@ -1043,7 +1042,6 @@ cc_library( copts = runtime_copts(), visibility = ["//visibility:public"], deps = [ - "//xla/tsl/framework/contraction:eigen_contraction_kernel_no_mkl", "@com_google_absl//absl/base:core_headers", "@eigen_archive//:eigen3", ], diff --git a/third_party/xla/xla/service/cpu/cpu_compiler.cc b/third_party/xla/xla/service/cpu/cpu_compiler.cc index 39a6808562aa7c..66f0170f55c5a2 100644 --- a/third_party/xla/xla/service/cpu/cpu_compiler.cc +++ b/third_party/xla/xla/service/cpu/cpu_compiler.cc @@ -753,7 +753,7 @@ absl::Status CpuCompiler::RunHloPassesThroughLayoutAssn( pipeline.AddPass(); pipeline.AddPass(); - // The TopkDecomposer generates a compare op with type=TOTALORDER and must + // The TopkDecomposer generates a compare op with order=TOTAL and must // run before the ComparisonExpander which rewrites such comparisons. pipeline.AddPass([&](const HloInstruction* instr) { return instr->opcode() == HloOpcode::kTopK; diff --git a/third_party/xla/xla/service/cpu/cpu_compiler_test.cc b/third_party/xla/xla/service/cpu/cpu_compiler_test.cc index ad581e1851f0e1..5d50791b0c2f2c 100644 --- a/third_party/xla/xla/service/cpu/cpu_compiler_test.cc +++ b/third_party/xla/xla/service/cpu/cpu_compiler_test.cc @@ -127,7 +127,7 @@ TEST_F(CpuCompilerTest, PermutationSortConvertedToScatter) { p.0.rhs = f32[] parameter(1) p.1.lhs = s32[] parameter(2) p.1.rhs = s32[] parameter(3) - ROOT lt = pred[] compare(p.0.lhs, p.0.rhs), direction=LT, type=TOTALORDER + ROOT lt = pred[] compare(p.0.lhs, p.0.rhs), direction=LT, order=TOTAL } compare2 { diff --git a/third_party/xla/xla/service/cpu/cpu_runtime.h b/third_party/xla/xla/service/cpu/cpu_runtime.h index c2b34ff178dc79..23498e5b68af7f 100644 --- a/third_party/xla/xla/service/cpu/cpu_runtime.h +++ b/third_party/xla/xla/service/cpu/cpu_runtime.h @@ -108,8 +108,6 @@ inline constexpr absl::string_view inline constexpr absl::string_view kReleaseOutfeedBufferAfterPopulationSymbolName = "__xla_cpu_runtime_ReleaseOutfeedBufferAfterPopulation"; -inline constexpr absl::string_view kParallelForkJoinSymbolName = - "__xla_cpu_runtime_ParallelForkJoin"; inline constexpr absl::string_view kPrintfToStderrSymbolName = "__xla_cpu_runtime_PrintfToStderr"; inline constexpr absl::string_view kStatusIsSuccessSymbolName = diff --git a/third_party/xla/xla/service/cpu/ir_emitter.cc b/third_party/xla/xla/service/cpu/ir_emitter.cc index 1780319a8145cb..35bf9cf364d5ca 100644 --- a/third_party/xla/xla/service/cpu/ir_emitter.cc +++ b/third_party/xla/xla/service/cpu/ir_emitter.cc @@ -81,7 +81,6 @@ limitations under the License. #include "xla/primitive_util.h" #include "xla/service/buffer_assignment.h" #include "xla/service/collective_ops_utils.h" -#include "xla/service/cpu/backend_config.pb.h" #include "xla/service/cpu/cpu_instruction_fusion.h" #include "xla/service/cpu/cpu_options.h" #include "xla/service/cpu/cpu_runtime.h" @@ -2194,45 +2193,10 @@ absl::Status IrEmitter::HandleFusion(HloInstruction* fusion) { absl::Status IrEmitter::HandleCall(HloInstruction* call) { HloComputation* computation = call->to_apply(); - llvm::Function* call_ir_function = FindOrDie( - emitted_functions_, ComputationToEmit{computation, allow_reassociation_}); ABSL_RETURN_IF_ERROR(EmitTargetAddressForOp(call)); - auto backend_config_or = - computation->root_instruction()->backend_config(); - if (backend_config_or.ok() && - !backend_config_or->outer_dimension_partitions().empty()) { - // Having a nonempty set of 'outer_dimension_partitions' means that this - // computation has been specially selected to be parallelized (one where the - // root instruction is trivially parallelizable, like elementwise addition - // of two tensors). The LLVM function generated for this computation accepts - // an additional set of loop bounds, allowing the caller to control the - // subset of the output that is generated by each call. - - std::vector call_args = GetArrayFunctionCallArguments( - {}, b(), computation->name(), - /*return_value_buffer=*/emitted_value_[call], - /*exec_run_options_arg=*/GetExecutableRunOptionsArgument(), - /*buffer_table_arg=*/GetBufferTableArgument(), - /*status_arg=*/GetStatusArgument(), - /*profile_counters_arg=*/GetProfileCountersArgument()); - - // The parallel fork/join runtime will call the generated function once for - // each partition in parallel, using an appropriate set of loop bounds for - // each call such that it only generates one partition of the output. - HloInstruction* root = computation->root_instruction(); - ABSL_RETURN_IF_ERROR(EmitCallToParallelForkJoin( - call_args, root->shape(), - backend_config_or->outer_dimension_partitions(), b(), call_ir_function, - computation->name())); - - if (ComputationTransitivelyContainsCustomCall(computation)) { - EmitEarlyReturnIfErrorStatus(); - } - } else { - EmitGlobalCall(*computation, computation->name()); - } + EmitGlobalCall(*computation, computation->name()); return absl::OkStatus(); } diff --git a/third_party/xla/xla/service/cpu/ir_function.cc b/third_party/xla/xla/service/cpu/ir_function.cc index f163c2b975b4b0..f74a2de721cca4 100644 --- a/third_party/xla/xla/service/cpu/ir_function.cc +++ b/third_party/xla/xla/service/cpu/ir_function.cc @@ -16,21 +16,15 @@ limitations under the License. #include "xla/service/cpu/ir_function.h" #include -#include #include -#include #include -#include "absl/log/check.h" -#include "absl/status/status.h" #include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" #include "absl/types/span.h" #include "llvm/IR/Function.h" #include "llvm/IR/Value.h" -#include "xla/service/cpu/cpu_runtime.h" #include "xla/service/llvm_ir/llvm_util.h" -#include "xla/shape_partition.h" #include "xla/status_macros.h" namespace xla { @@ -225,106 +219,5 @@ std::vector GetArrayFunctionCallArguments( buffer_table_arg, status_arg, profile_counters_arg}; } -// Emits a call to a runtime fork/join function which dispatches parallel -// calls to 'parallel_function' (and joins threads before returning). -absl::Status EmitCallToParallelForkJoin( - const std::vector& arguments, const Shape& shape, - absl::Span dimension_partition_counts, - llvm::IRBuilderBase* b, llvm::Function* parallel_function, - absl::string_view name) { - llvm::Module* module = b->GetInsertBlock()->getModule(); - - // Build ParallelForkJoin function type. - std::vector compute_function_params = - GetComputeFunctionParams(module); - // Number of parallel compute functions. - compute_function_params.push_back(b->getInt32Ty()); - // Array of partitions. There is an array element for each - // partition x partition_dim x 2 (for dimension start and limit). - compute_function_params.push_back( - llvm::PointerType::get(module->getContext(), 0)); - // Number of partitioned most-major dimensions in 'shape'. - compute_function_params.push_back(b->getInt32Ty()); - // Function pointer for compute function to be dispatched in parallel. - compute_function_params.push_back( - llvm::PointerType::get(module->getContext(), 0)); - - llvm::FunctionType* fork_join_type = llvm::FunctionType::get( - /*Result=*/llvm::Type::getVoidTy(module->getContext()), - /*Params=*/compute_function_params, - /*isVarArg=*/false); - - llvm::Function* fork_join_func = llvm::dyn_cast( - module - ->getOrInsertFunction(runtime::kParallelForkJoinSymbolName, - fork_join_type) - .getCallee()); - fork_join_func->setCallingConv(llvm::CallingConv::C); - fork_join_func->setDoesNotThrow(); - - // Add common compute function arguments. - std::vector fork_join_arguments(arguments); - - // Create ShapePartitionIterator to generate all partitions of 'shape'. - ShapePartitionIterator partition_iterator(shape, dimension_partition_counts); - const int64_t num_partitions = partition_iterator.GetTotalPartitionCount(); - // Add argument specifying the number of parallel partitions. - fork_join_arguments.push_back(b->getInt32(num_partitions)); - - // The number of partitioned most-major dimensions in 'shape'. - const int32_t num_partitioned_dims = dimension_partition_counts.size(); - // A dimension partition consists of two elements: [start_index, limit_index). - const int32_t dim_partition_size = 2; - // Calculate array partition stride. - const int32_t array_partition_stride = - num_partitioned_dims * dim_partition_size; - // Calculate the total number of elements in the partition array. - const int32_t partition_array_size = - dim_partition_size * num_partitioned_dims * num_partitions; - - // Store dimension partition values as llvm constants in 'partitions'. - // See comments in runtime_fork_join.cc for array layout description. - std::vector partitions(partition_array_size); - for (int32_t i = 0; i < num_partitions; ++i) { - std::vector> dim_partitions = - partition_iterator.GetPartition(i); - CHECK_EQ(num_partitioned_dims, dim_partitions.size()); - const int32_t partition_index = i * array_partition_stride; - for (int32_t j = 0; j < num_partitioned_dims; ++j) { - const std::pair& dim_partition = dim_partitions[j]; - const int32_t index = partition_index + j * dim_partition_size; - // Store partition [dim_start, dim_limit) intervals for each dimension. - partitions[index] = b->getInt64(dim_partition.first); - partitions[index + 1] = - b->getInt64(dim_partition.first + dim_partition.second); - } - } - - // Create global variable out of dimension partitions in 'partitions'. - llvm::ArrayType* partitions_array_type = - llvm::ArrayType::get(b->getInt64Ty(), partition_array_size); - llvm::Constant* partitions_array = - llvm::ConstantArray::get(partitions_array_type, partitions); - llvm::GlobalVariable* global_partitions_array = new llvm::GlobalVariable( - /*M=*/*module, - /*Ty=*/partitions_array_type, - /*isConstant=*/true, - /*Linkage=*/llvm::GlobalValue::PrivateLinkage, - /*Initializer=*/partitions_array, - /*Name=*/ - absl::StrCat(name, "_parallel_dimension_partitions")); - - // Add argument specifying parallel dimension partitions. - fork_join_arguments.push_back(global_partitions_array); - // Add argument specifying the number of partitioned most-major dimensions. - fork_join_arguments.push_back(b->getInt32(num_partitioned_dims)); - // Add argument for parallel compute function pointer. - fork_join_arguments.push_back(parallel_function); - // Emit call to parallel fork/join. - b->CreateCall(fork_join_func, fork_join_arguments); - - return absl::OkStatus(); -} - } // namespace cpu } // namespace xla diff --git a/third_party/xla/xla/service/cpu/ir_function.h b/third_party/xla/xla/service/cpu/ir_function.h index e67258fd3abb2f..0db12ede9ce1c4 100644 --- a/third_party/xla/xla/service/cpu/ir_function.h +++ b/third_party/xla/xla/service/cpu/ir_function.h @@ -139,14 +139,6 @@ std::vector GetArrayFunctionCallArguments( llvm::Value* exec_run_options_arg, llvm::Value* buffer_table_arg, llvm::Value* status_arg, llvm::Value* profile_counters_arg); -// Emits a call to a runtime fork/join function which dispatches parallel -// calls to 'parallel_function' (and joins threads before returning). -absl::Status EmitCallToParallelForkJoin( - const std::vector& arguments, const Shape& shape, - absl::Span dimension_partition_counts, - llvm::IRBuilderBase* b, llvm::Function* parallel_function, - absl::string_view name); - } // namespace cpu } // namespace xla diff --git a/third_party/xla/xla/service/elemental_ir_emitter.cc b/third_party/xla/xla/service/elemental_ir_emitter.cc index e49674dde9c412..9342e17e2bfdb0 100644 --- a/third_party/xla/xla/service/elemental_ir_emitter.cc +++ b/third_party/xla/xla/service/elemental_ir_emitter.cc @@ -1376,16 +1376,13 @@ absl::StatusOr ElementalIrEmitter::EmitComplexUnaryOp( ABSL_ASSIGN_OR_RETURN(llvm::Value * sin_b, EmitSin(component_type, b)); llvm::Value* imag_numerator = FMul(four, FMul(cos_b, sin_b)); - // About "x^2 is a better approximation than Expm1(x) + Expm1(x) - // for small values of x": this statement is not - // accurate. Previously, Expm1(x) implementation had accuracy - // issues for small x (where it was supposed to stand out in - // accuracy!), but after resolving these issues (see - // openxla/xla#10376), using precomputed exp_2a_m1 and - // exp_neg_2a_m1 is accurate enough and we'll save a few - // instructions. - - auto exp_sum_m2 = FAdd(exp_2a_m1, exp_neg_2a_m1); + // Computing the denominator's exponential term as expm1(2a) + expm1(-2a) + // suffers from catastrophic cancellation as a -> 0. Instead, use the + // identity (e^(2a) - 1)(e^(-2a) - 1) = 2 - 2*cosh(2a), which gives + // 2*(cosh(2a) - 1) = -expm1(2a) * expm1(-2a) to avoid cancellation + // and reuse exp_2a_m1 and exp_neg_2a_m1. + llvm::Value* exp_product = FMul(exp_2a_m1, exp_neg_2a_m1); + llvm::Value* exp_sum_m2 = FMul(neg_one, exp_product); llvm::Value* denom = FAdd(exp_sum_m2, two_cos_2b_p2); // As `a` grows toward +inf and -inf, the real numerator will grow towards diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD index 1f6a4aee2d3fce..bb4371e44bc69f 100644 --- a/third_party/xla/xla/service/gpu/BUILD +++ b/third_party/xla/xla/service/gpu/BUILD @@ -2150,6 +2150,7 @@ xla_test( "h100": [ "full", ], + "gb300": ["broken"], # TODO(b/505574653): re-enable when fixed. }, backends = ["gpu"], shard_count = 2, diff --git a/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass.cc b/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass.cc index 02337aaa512f05..796a9d40aa7dd5 100644 --- a/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass.cc +++ b/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass.cc @@ -290,6 +290,7 @@ ConfigAssigner::Options GetConfigAssignerOptions( debug_options.xla_gpu_use_new_autotune_cache_format(); options.compile_all_supported_configs = debug_options.xla_compile_all_supported_configs(); + options.force_config = debug_options.xla_force_config(); return options; } @@ -313,6 +314,7 @@ CodegenOrchestrator::Options GetCodegenOrchestratorOptions( options.allow_reg_spills_fn = [](const HloInstruction&, autotuner::Backend) { return false; }; } + options.candidate_configs_file = debug_options.xla_candidate_configs_file(); return options; } diff --git a/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass_test.cc b/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass_test.cc index 56178f086a5476..6c0853e3e0b2ec 100644 --- a/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass_test.cc +++ b/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass_test.cc @@ -999,6 +999,30 @@ TEST_F(ConfigAssignerPassTest, CudnnFusionForbidsSpills) { EXPECT_FALSE(options.allow_reg_spills_fn(*instr, autotuner::Backend::CUDNN)); } +TEST_F(ConfigAssignerPassTest, ForceConfigPropagatesToConfigAssignerOptions) { + DebugOptions debug_options = GetDebugOptionsForTest(); + const std::string forced_config = + "backend: TRITON\n" + "backend_config {\n" + " triton {\n" + " block_m: 64\n" + " block_n: 64\n" + " }\n" + "}"; + debug_options.set_xla_force_config(forced_config); + auto options = + GetConfigAssignerOptions(debug_options, /*is_deviceless=*/false); + EXPECT_EQ(options.force_config, forced_config); +} + +TEST_F(ConfigAssignerPassTest, + CandidateConfigsFilePropagatesToCodegenOrchestratorOptions) { + DebugOptions debug_options = GetDebugOptionsForTest(); + debug_options.set_xla_candidate_configs_file("/tmp/candidates.pbtxt"); + auto options = GetCodegenOrchestratorOptions(debug_options); + EXPECT_EQ(options.candidate_configs_file, "/tmp/candidates.pbtxt"); +} + TEST_F(ConfigAssignerPassTest, CustomFusionForbidsSpills) { auto options = GetCodegenOrchestratorOptions(GetDebugOptionsForTest()); diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc index a14075073a6126..e218d76eb19931 100644 --- a/third_party/xla/xla/service/gpu/gpu_compiler.cc +++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc @@ -708,7 +708,7 @@ absl::Status RunPreSPMDPartitionerPasses( pre_spmd_pipeline.AddPass(); - // The TopkDecomposer generates a compare op with type=TOTALORDER and must + // The TopkDecomposer generates a compare op with order=TOTAL and must // run before the ComparisonExpander which rewrites such comparisons. pre_spmd_pipeline.AddPass([&](const HloInstruction* instr) { return instr->opcode() == HloOpcode::kTopK; diff --git a/third_party/xla/xla/service/gpu/gpu_compiler_test.cc b/third_party/xla/xla/service/gpu/gpu_compiler_test.cc index 6edf4434711d1b..c67fb55fabf3b3 100644 --- a/third_party/xla/xla/service/gpu/gpu_compiler_test.cc +++ b/third_party/xla/xla/service/gpu/gpu_compiler_test.cc @@ -1745,25 +1745,30 @@ m { // Define a test-specific enum for expected TopK implementations. enum class TopKImpl { - kCustomKernel, // Custom GPU kernel - kSelectK, // raft::select_k - kSort // Fallback Sort+Slice + kCustomKernel, // Custom GPU kernel + kSelectK, // raft::select_k + kSort, // Fallback Sort+Slice + kSelectKWithU64Adapter, // pack_to_u64 + raft::select_k(u64) + + // unpack_from_u64 + kSortWithS32Adapter // pack_to_s32 + Sort(s32)+Slice + unpack_from_s32 }; // Test fixture for verifying GPU TopK lowering to SelectK or custom kernel. class GpuCompilerSelectKTest : public GpuCompilerTest, - public ::testing::WithParamInterface> {}; + public ::testing::WithParamInterface< + std::tuple> {}; // Test lowering of TopK to different GPU implementations // (CustomKernel, raft::select_k, or Sort+Slice (LLVM/CUBSort)). TEST_P(GpuCompilerSelectKTest, SelectKOrCustomKernelThunk) { - auto [n, k, expected_impl] = GetParam(); + auto [dtype, n, k, is_stable, expected_impl] = GetParam(); bool is_rocm = device_description().gpu_compute_capability().IsRocm(); bool is_oneapi = device_description().gpu_compute_capability().IsOneAPI(); - if (is_rocm && expected_impl == TopKImpl::kSelectK) { + if (is_rocm && (expected_impl == TopKImpl::kSelectK || + expected_impl == TopKImpl::kSelectKWithU64Adapter)) { GTEST_SKIP() << "raft::select_k is not supported in ROCm."; } // TODO(intel-tf): Remove this check once TopK specialization for SYCL/oneAPI @@ -1777,11 +1782,11 @@ TEST_P(GpuCompilerSelectKTest, SelectKOrCustomKernelThunk) { HloModule m ENTRY main { - p = f32[8,$0]{1,0} parameter(0) - ROOT t = (f32[8,$1]{1,0}, s32[8,$1]{1,0}) topk(p), k=$1, largest=true, is_stable=false + p = $0[8,$1]{1,0} parameter(0) + ROOT t = ($0[8,$2]{1,0}, s32[8,$2]{1,0}) topk(p), k=$2, largest=true, is_stable=$3 } )", - n, k); + dtype, n, k, is_stable); // Configure module with debug options. HloModuleConfig config; @@ -1833,6 +1838,12 @@ ENTRY main { if (kinds.size() == 1) { // LLVM EXPECT_THAT(kinds, ElementsAre(Thunk::Kind::kCommandBuffer)); + } else if (kinds.size() == 4 && kinds[0] == Thunk::Kind::kCopy) { + // LLVM Bitonic sort (unbundled, with input copy) + EXPECT_THAT(kinds, + ElementsAre(Thunk::Kind::kCopy, Thunk::Kind::kCustomKernel, + Thunk::Kind::kCustomKernel, + Thunk::Kind::kCustomKernel)); } else if (kinds.size() == 4) { // CUB sort via FFI custom call EXPECT_THAT(kinds, ElementsAre(Thunk::Kind::kCustomKernel, @@ -1845,6 +1856,19 @@ ENTRY main { break; } + case TopKImpl::kSelectKWithU64Adapter: + EXPECT_THAT(kinds, + ElementsAre(Thunk::Kind::kCustomKernel, Thunk::Kind::kSelectK, + Thunk::Kind::kCustomKernel)); + break; + + case TopKImpl::kSortWithS32Adapter: { + EXPECT_THAT(kinds, ElementsAre(Thunk::Kind::kCustomKernel, + Thunk::Kind::kCustomKernel, + Thunk::Kind::kCustomKernel)); + break; + } + default: FAIL() << "Unexpected TopKImpl: " << static_cast(expected_impl); } @@ -2575,16 +2599,52 @@ INSTANTIATE_TEST_SUITE_P( }); auto SelectKTestParams() { - // Depending on N and K, XLA chooses different TopK implementations: - // CustomKernel, raft::select_k, or Sort+Slice. + // Depending on dtype, N, K and is_stable flag, XLA chooses different TopK + // implementations: + // CustomKernel, raft::select_k, select_k_with_u64_adapter, or Sort+Slice. // The heuristic for selecting between TopK CustomKernel and // raft::matrix::select_k was developed as part of the initial research // described in b/409009349. - return ::testing::Values(std::make_tuple(1023, 4, TopKImpl::kSelectK), - std::make_tuple(1024, 4, TopKImpl::kCustomKernel), - std::make_tuple(1024, 16, TopKImpl::kSelectK), - std::make_tuple(8192, 24, TopKImpl::kSelectK), - std::make_tuple(8192, 512, TopKImpl::kSort)); + return ::testing::Values( + // dtype, n, k, is_stable, expected_impl + std::make_tuple("f32", 1023, 4, false, TopKImpl::kSelectK), + std::make_tuple("f32", 1023, 4, true, TopKImpl::kSort), + std::make_tuple("f32", 1024, 4, false, TopKImpl::kCustomKernel), + std::make_tuple("f32", 1024, 4, true, TopKImpl::kCustomKernel), + std::make_tuple("f32", 1024, 16, false, TopKImpl::kSelectK), + std::make_tuple("f32", 1024, 16, true, TopKImpl::kCustomKernel), + std::make_tuple("f32", 8192, 24, false, TopKImpl::kSelectK), + std::make_tuple("f32", 8192, 24, true, TopKImpl::kSelectKWithU64Adapter), + std::make_tuple("f32", 8192, 512, false, TopKImpl::kSort), + std::make_tuple("f32", 8192, 512, true, TopKImpl::kSort), + // f32: exact upper bound of max_k + std::make_tuple("f32", 8192, 128, false, TopKImpl::kSelectK), + std::make_tuple("f32", 8192, 128, true, TopKImpl::kSelectKWithU64Adapter), + // f32: just over the upper bound + std::make_tuple("f32", 8192, 129, false, TopKImpl::kSort), + std::make_tuple("f32", 8192, 129, true, TopKImpl::kSort), + // bf16 and size <= 2**16 - use sort(s32) + slice. + std::make_tuple("bf16", 1023, 4, false, TopKImpl::kSortWithS32Adapter), + std::make_tuple("bf16", 1023, 4, true, TopKImpl::kSortWithS32Adapter), + std::make_tuple("bf16", 1024, 4, false, TopKImpl::kSortWithS32Adapter), + std::make_tuple("bf16", 1024, 4, true, TopKImpl::kSortWithS32Adapter), + std::make_tuple("bf16", 1024, 16, false, TopKImpl::kSortWithS32Adapter), + std::make_tuple("bf16", 1024, 16, true, TopKImpl::kSortWithS32Adapter), + // bf16 and size > 2**16. + std::make_tuple("bf16", 65540, 16, false, TopKImpl::kSelectK), + std::make_tuple("bf16", 65540, 16, true, TopKImpl::kCustomKernel), + std::make_tuple("bf16", 65540, 24, false, TopKImpl::kSelectK), + std::make_tuple("bf16", 65540, 24, true, + TopKImpl::kSelectKWithU64Adapter), + std::make_tuple("bf16", 65540, 512, false, TopKImpl::kSort), + std::make_tuple("bf16", 65540, 512, true, TopKImpl::kSort), + // bf16: exact upper bound of max_k for batch=8 + std::make_tuple("bf16", 65540, 128, false, TopKImpl::kSelectK), + std::make_tuple("bf16", 65540, 128, true, + TopKImpl::kSelectKWithU64Adapter), + // bf16: just over the upper bound + std::make_tuple("bf16", 65540, 129, false, TopKImpl::kSort), + std::make_tuple("bf16", 65540, 129, true, TopKImpl::kSort)); } // Instantiate the test suite with (n, k, expected_kind) pairs. INSTANTIATE_TEST_SUITE_P(SelectKOrCustomKernel, GpuCompilerSelectKTest, diff --git a/third_party/xla/xla/service/gpu/model/triton_emitter_constraints.cc b/third_party/xla/xla/service/gpu/model/triton_emitter_constraints.cc index ff4683114da0fc..45dff5f691692c 100644 --- a/third_party/xla/xla/service/gpu/model/triton_emitter_constraints.cc +++ b/third_party/xla/xla/service/gpu/model/triton_emitter_constraints.cc @@ -406,7 +406,8 @@ Decision VerifyTritonConstraints(const TiledHloComputation& tiled_computation, const int64_t shared_memory_limit = device_info.shared_memory_per_block_optin(); for (const TiledHloInstruction* inst : all_instructions) { - if (inst->hlo()->opcode() != HloOpcode::kTranspose) { + if (inst->hlo()->opcode() != HloOpcode::kTranspose || + inst->operands().empty()) { continue; } // The transposed operand (operand 0) is the tile that is staged in shared diff --git a/third_party/xla/xla/service/heap_simulator/heap_simulator.cc b/third_party/xla/xla/service/heap_simulator/heap_simulator.cc index 847bcd49250dc4..90e9347a50bd38 100644 --- a/third_party/xla/xla/service/heap_simulator/heap_simulator.cc +++ b/third_party/xla/xla/service/heap_simulator/heap_simulator.cc @@ -342,6 +342,17 @@ absl::Status HeapSimulator::RunComputation( auto& buffer_live_ranges = hlo_live_range->buffer_live_ranges(); + // A value used as the base of a "view" (a value colored options_.view_color, + // an address into the base's buffer with no storage of its own) is read + // through the view by the view's consumers at later schedule times. Extend + // the base's live range to the view's last transitive reader before the + // define/free events are laid out, so the buffer cannot be recycled while a + // reader still loads from it. + if (options_.view_color.has_value()) { + ExtendViewBaseLiveRanges(hlo_live_range, dataflow_analysis, + *options_.view_color); + } + for (const HloValue* value : dataflow_analysis.values()) { // Ignore buffers that are not tracked. if (!buffer_live_ranges.contains(value)) { diff --git a/third_party/xla/xla/service/heap_simulator/heap_simulator.h b/third_party/xla/xla/service/heap_simulator/heap_simulator.h index 187e06dd731fb5..9beb6d0c34768f 100644 --- a/third_party/xla/xla/service/heap_simulator/heap_simulator.h +++ b/third_party/xla/xla/service/heap_simulator/heap_simulator.h @@ -144,6 +144,13 @@ class HeapSimulator { // If 'buffers_to_assign' is provided, only those buffers are assigned // offsets, otherwise all buffers defined by the instructions are assigned. const absl::flat_hash_set* buffers_to_assign; + // Memory space color marking "view" values (address stand-ins aliasing + // into their first operand's buffer, see + // BufferAssigner::Options::dus_view_color). When set, a value used as a + // view's base is kept live until the view's last transitive reader: those + // readers read the value's buffer through the view, so it must not be + // recycled before them. + std::optional view_color; }; // Returns the minimum memory required to compute an HLO module where all diff --git a/third_party/xla/xla/service/hlo_schedule_test.cc b/third_party/xla/xla/service/hlo_schedule_test.cc index c6567096e7eda9..fe755b49f1311d 100644 --- a/third_party/xla/xla/service/hlo_schedule_test.cc +++ b/third_party/xla/xla/service/hlo_schedule_test.cc @@ -17,6 +17,7 @@ limitations under the License. #include #include +#include #include #include #include @@ -666,6 +667,103 @@ ENTRY %test (arg.0: (f32[], f32[])) -> f32[] { ASSERT_OK(module->schedule().Verify()); } +TEST_F(HloScheduleTest, UpdateScheduleWithReplacedOperandsScheduledLater) { + // Test that when an instruction's operands are replaced with instructions + // that were scheduled later (or newly added instructions depending on + // instructions scheduled later), Update() properly invalidates and reorders + // them in topological order. + const std::string module_str = R"( +HloModule m, is_scheduled=true + +ENTRY %test { + %c0 = f32[] constant(1.0) + %x1 = f32[] negate(%c0) + %inst_a = f32[] add(%x1, %c0) + %x0 = f32[] copy(%c0) + ROOT %inst_b = f32[] add(%inst_a, %x0) +} +)"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(module_str)); + HloComputation* entry = module->entry_computation(); + HloInstruction* c0 = entry->GetInstructionWithName("c0"); + HloInstruction* x0 = entry->GetInstructionWithName("x0"); + HloInstruction* x1 = entry->GetInstructionWithName("x1"); + HloInstruction* inst_a = entry->GetInstructionWithName("inst_a"); + + // Create inst_c = add(x0, c0) and replace uses of x1 with inst_c. + HloInstruction* inst_c = entry->AddInstruction( + HloInstruction::CreateBinary(c0->shape(), HloOpcode::kAdd, x0, c0)); + ASSERT_OK(x1->ReplaceAllUsesWith(inst_c)); + ASSERT_OK(entry->RemoveInstruction(x1)); + + // Before Update(), the schedule has inst_a before x0 (and before inst_c). + ASSERT_IS_NOT_OK(module->schedule().Verify()); + ASSERT_OK(module->schedule().Update()); + ASSERT_OK(module->schedule().Verify()); + + // Verify inst_a is scheduled after x0 and inst_c. + const auto& seq = module->schedule().sequence(entry).instructions(); + auto pos = [&](const HloInstruction* inst) { + return std::distance(seq.begin(), absl::c_find(seq, inst)); + }; + EXPECT_LT(pos(x0), pos(inst_c)); + EXPECT_LT(pos(inst_c), pos(inst_a)); +} + +TEST_F(HloScheduleTest, UpdateSchedulePreservesOrderWhenNewOperandAdded) { + // Tests that when an instruction's operand is replaced with a newly added + // instruction, Update() preserves the relative schedule order among existing + // scheduled instructions rather than prematurely promoting the consumer ASAP. + // + // This models the regression observed in JAX polydiv + // (lax_numpy_test_gpu_b200) where CopyInsertion added a copy for an initial + // buffer and rewired the first in-place slice update. A naive invalidation + // approach that evicts any instruction with an unscheduled operand caused the + // slice update to be emitted prematurely before intermediate instructions, + // disrupting the in-place execution sequence. + const std::string module_str = R"( +HloModule m, is_scheduled=true + +ENTRY %test { + %c0 = f32[] constant(1.0) + %x0 = f32[] negate(%c0) + %interm = f32[] negate(%x0) + %inst_a = f32[] add(%c0, %x0) + ROOT %inst_b = f32[] add(%inst_a, %interm) +} +)"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(module_str)); + HloComputation* entry = module->entry_computation(); + HloInstruction* c0 = entry->GetInstructionWithName("c0"); + HloInstruction* interm = entry->GetInstructionWithName("interm"); + HloInstruction* inst_a = entry->GetInstructionWithName("inst_a"); + + // In the original schedule, %interm is scheduled before %inst_a. + const auto& orig_seq = module->schedule().sequence(entry).instructions(); + auto get_pos = [](const std::vector& seq, + const HloInstruction* inst) { + return std::distance(seq.begin(), absl::c_find(seq, inst)); + }; + ASSERT_LT(get_pos(orig_seq, interm), get_pos(orig_seq, inst_a)); + + // Add new_copy = copy(c0) and replace operand 0 of inst_a (%c0) with + // new_copy. + HloInstruction* new_copy = entry->AddInstruction( + HloInstruction::CreateUnary(c0->shape(), HloOpcode::kCopy, c0)); + ASSERT_OK(inst_a->ReplaceOperandWith(0, new_copy)); + + ASSERT_OK(module->schedule().Update()); + ASSERT_OK(module->schedule().Verify()); + + // new_copy must be placed before inst_a, but inst_a must not be prematurely + // promoted before %interm; the original relative ordering must be preserved. + const auto& new_seq = module->schedule().sequence(entry).instructions(); + EXPECT_LT(get_pos(new_seq, new_copy), get_pos(new_seq, inst_a)); + EXPECT_LT(get_pos(new_seq, interm), get_pos(new_seq, inst_a)); +} + std::unique_ptr BuildBenchmarkModule(int64_t num_instructions) { HloModuleConfig config; auto module = std::make_unique("bm_module", config); diff --git a/third_party/xla/xla/service/latency_hiding_scheduler.cc b/third_party/xla/xla/service/latency_hiding_scheduler.cc index 5db9e349a950ee..9bf2702667d0f9 100644 --- a/third_party/xla/xla/service/latency_hiding_scheduler.cc +++ b/third_party/xla/xla/service/latency_hiding_scheduler.cc @@ -366,10 +366,10 @@ bool LatencyEstimator::IsAsyncPair(const HloGraphNode& from, bool LatencyEstimator::IsP2pPair(const HloGraphNode& from, const HloGraphNode& target) const { - return (from.GetInstr().opcode() == HloOpcode::kSend && - target.GetInstr().opcode() == HloOpcode::kSendDone) || - (from.GetInstr().opcode() == HloOpcode::kRecv && - target.GetInstr().opcode() == HloOpcode::kRecvDone); + return (from.GetOpcode() == HloOpcode::kSend && + target.GetOpcode() == HloOpcode::kSendDone) || + (from.GetOpcode() == HloOpcode::kRecv && + target.GetOpcode() == HloOpcode::kRecvDone); } std::optional @@ -1736,7 +1736,7 @@ bool ReadySetLt::AIsBetterThanB(DefaultSchedulerCore::ScheduleCandidate& a, } } if (an->IsSupportedAsyncDone() && bn->IsSupportedAsyncDone() && - an->GetInstr().opcode() == bn->GetInstr().opcode()) { + an->GetOpcode() == bn->GetOpcode()) { const HloGraphNode& start_an = sched_state.sched_graph->GetNode(an->GetInstr().operand(0)); const HloGraphNode& start_bn = @@ -2977,6 +2977,11 @@ HloScheduleGraph::HloScheduleGraph( DCHECK_EQ(n, GetNodePtr(instr)); n->instr_ = instr; n->opcode_ = instr->opcode(); + n->is_host_transfer_ = + (n->opcode_ == HloOpcode::kSend || n->opcode_ == HloOpcode::kSendDone || + n->opcode_ == HloOpcode::kRecv || + n->opcode_ == HloOpcode::kRecvDone) && + static_cast(instr)->is_host_transfer(); n->original_position_ = current_pos; current_pos++; @@ -3476,6 +3481,44 @@ void HloScheduleGraph::AnnotateGraph( } } +bool DefaultSchedulerCore::DefaultSchedulingInstructionCrossesOverlapLimit( + const SchedulingState& sched_state, const HloGraphNode* node) { + if (!node->HasRecursiveResources()) { + return false; + } + const HloInstruction& instr = node->GetInstr(); + const bool is_nested_sync_comp = !instr.called_computations().empty() && + instr.opcode() != HloOpcode::kAsyncStart && + instr.opcode() != HloOpcode::kAsyncDone; + + auto& num_resources_needed = node->GetRecursiveResources(); + // NOLINTNEXTLINE(*-custom-deterministic-iteration-order) + for (const auto& [resource, count] : num_resources_needed) { + auto it = sched_state.max_concurrent_resource.find(resource); + if (it == sched_state.max_concurrent_resource.end()) { + continue; + } + if (is_nested_sync_comp && + sched_state.async_tracker->IsInorderResource(resource)) { + int64_t total_capacity = + sched_state.async_tracker->GetNumAvailableResources(resource); + if (it->second < total_capacity) { + VLOG(5) << "In-order resource " << resource + << " currently has outer in-flight operations (available " + << it->second << " < total " << total_capacity + << "). Cannot schedule nested computation " << instr.name(); + return true; + } + } + if (count > it->second) { + VLOG(5) << "Cross overlap limit for resource: " << resource + << " count: " << count << " limit: " << it->second; + return true; + } + } + return false; +} + absl::Status DefaultSchedulerCore::InitializeScheduler( const HloModule* module) { module_ = module; @@ -3545,24 +3588,7 @@ absl::Status DefaultSchedulerCore::InitializeScheduler( if (!scheduling_instruction_crosses_overlap_limit_) { scheduling_instruction_crosses_overlap_limit_ = - [](const SchedulingState& sched_state, const HloGraphNode* node) { - if (!node->HasRecursiveResources()) { - return false; - } - auto& num_resources_needed = node->GetRecursiveResources(); - for (const auto& [resource, count] : num_resources_needed) { - auto it = sched_state.max_concurrent_resource.find(resource); - if (it == sched_state.max_concurrent_resource.end()) { - continue; - } - if (count > it->second) { - VLOG(5) << "Cross overlap limit for resource: " << resource - << " count: " << count << " limit: " << it->second; - return true; - } - } - return false; - }; + DefaultSchedulingInstructionCrossesOverlapLimit; is_default_scheduling_instruction_crosses_overlap_limit_ = true; } return absl::OkStatus(); diff --git a/third_party/xla/xla/service/latency_hiding_scheduler.h b/third_party/xla/xla/service/latency_hiding_scheduler.h index 532fd04afb1c2a..d353d222b81163 100644 --- a/third_party/xla/xla/service/latency_hiding_scheduler.h +++ b/third_party/xla/xla/service/latency_hiding_scheduler.h @@ -47,6 +47,7 @@ limitations under the License. #include "xla/hlo/analysis/hlo_reachability.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/ir/hlo_instructions.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/hlo/ir/hlo_schedule.h" #include "xla/hlo/pass/hlo_pass_interface.h" @@ -122,6 +123,11 @@ enum class ResourceHazardType { // ops that are valuable for selective overlap. kSelective = 3, kUnshareable = 4, + // An in-order resource (e.g. hardware FIFO command queue) where operations + // are processed sequentially. Instructions calling a synchronous nested + // computation using this resource cannot be scheduled while any outer + // asynchronous operation on this resource is in flight. + kInOrder = 5, }; template >> @@ -427,6 +433,13 @@ class AsyncTracker { // Returns whether the provided node occupies a selective resource. bool OccupiesSelectiveResource(const HloGraphNode* node) const; + // Returns whether the resource is an in-order resource (e.g. hardware FIFO + // command queue) where operations are processed sequentially and cannot be + // interleaved across nested computation boundaries. + bool IsInorderResource(int64_t resource_type) const { + return GetResourceHazardType(resource_type) == ResourceHazardType::kInOrder; + } + inline CanonicalAsyncOp GetCanonicalAsyncOp(const HloInstruction& hlo) const { return get_canonical_async_op_(hlo); } @@ -688,6 +701,10 @@ class HloGraphNode { explicit HloGraphNode(const HloInstruction* i, int64_t original_position) : instr_(i), opcode_(i->opcode()), original_position_(original_position) { InitBitFields(); + is_host_transfer_ = + (opcode_ == HloOpcode::kSend || opcode_ == HloOpcode::kSendDone || + opcode_ == HloOpcode::kRecv || opcode_ == HloOpcode::kRecvDone) && + static_cast(i)->is_host_transfer(); } static void UpdateOrAddDependency(HloGraphNode* from, HloGraphNode* to, @@ -761,6 +778,7 @@ class HloGraphNode { } const HloInstruction& GetInstr() const { return *instr_; } HloOpcode GetOpcode() const { return opcode_; } + bool IsHostTransfer() const { return is_host_transfer_; } bool IsScheduled() const { return scheduled_; } int32_t GetIndegree() const { return indegree_; } int32_t GetOutdegree() const { return outdegree_; } @@ -1053,14 +1071,16 @@ class HloGraphNode { // Opcode of instr_, copied here for better cache behavior (so we can look at // the opcode without having to touch another cache line). HloOpcode opcode_; + // If multiple nodes are there with force_delay_ = true, the one with the + // lowest delay priority will be scheduled first. + int force_delay_priority_ = 0; // Some of the booleans are looked at very often, so we avoid making them // bitfields // Force the scheduling of the nodes with attribute set as late as possible. bool force_delay_ = false; - // If multiple nodes are there with force_delay_ = true, the one with the - // lowest delay priority will be scheduled first. - int force_delay_priority_ = 0; + // Whether the instruction is a host transfer (send/recv). + bool is_host_transfer_ = false; // Force the scheduling of the nodes with attribute set as early as possible. bool force_early_ = false; // If has_rare_ is false, then all the fields in rare can assumed to be @@ -1946,6 +1966,8 @@ class DefaultSchedulerCore : public SchedulerCore { static bool DeleteOccupierFromResource( HloGraphNode::TimeCost current_time, HloEdge& edge, std::vector>& occupiers); + static bool DefaultSchedulingInstructionCrossesOverlapLimit( + const SchedulingState& sched_state, const HloGraphNode* node); int64_t GetMemoryPeak() override { return module_pressure_state_->GetMemoryPeak(); } diff --git a/third_party/xla/xla/service/memory_space_assignment/algorithm.cc b/third_party/xla/xla/service/memory_space_assignment/algorithm.cc index 84a9d65f6cd55d..cedbe5dda63579 100644 --- a/third_party/xla/xla/service/memory_space_assignment/algorithm.cc +++ b/third_party/xla/xla/service/memory_space_assignment/algorithm.cc @@ -2868,43 +2868,6 @@ std::optional MsaAlgorithm::GetLatestSourceOperandScheduleTime( return latest_source_operand_time; } -int64_t ViewExtendedTransitiveUseTime( - const HloInstruction* view, int64_t view_color, - const absl::flat_hash_map& - instruction_schedule) { - CHECK(!view->shape().IsTuple() && view->shape().has_layout() && - view->shape().layout().memory_space() == view_color) - << "not a view: " << view->ToString(); - auto is_view_colored = [view_color](const HloInstruction* instruction) { - return instruction->shape().has_layout() && - instruction->shape().layout().memory_space() == view_color; - }; - int64_t use_time = -1; - absl::flat_hash_set visited = {view}; - std::vector worklist = {view}; - while (!worklist.empty()) { - const HloInstruction* current = worklist.back(); - worklist.pop_back(); - auto time_it = instruction_schedule.find(current); - if (time_it != instruction_schedule.end()) { - use_time = std::max(use_time, time_it->second); - } - for (const HloInstruction* user : current->users()) { - if (is_view_colored(user)) { - if (visited.insert(user).second) { - worklist.push_back(user); - } - } else { - auto user_time_it = instruction_schedule.find(user); - if (user_time_it != instruction_schedule.end()) { - use_time = std::max(use_time, user_time_it->second); - } - } - } - } - return use_time; -} - namespace { // Computes each value's [first_use_time, last_use_time] interval. When diff --git a/third_party/xla/xla/service/memory_space_assignment/algorithm.h b/third_party/xla/xla/service/memory_space_assignment/algorithm.h index 71b7f2f69eac19..8df0efe92ca420 100644 --- a/third_party/xla/xla/service/memory_space_assignment/algorithm.h +++ b/third_party/xla/xla/service/memory_space_assignment/algorithm.h @@ -127,18 +127,6 @@ struct AllocationSegmentContext { bool only_extend_existing_allocation; }; -// Returns the latest schedule time at which `view` (a value colored -// `view_color`, see Options::dus_view_color) still has its underlying storage -// read through it: the max schedule time over the transitive closure of the -// view's readers, following users that are themselves view colored. Exposed -// for testing. -// -// REQUIRES: view->shape().IsTuple() == false. -int64_t ViewExtendedTransitiveUseTime( - const HloInstruction* view, int64_t view_color, - const absl::flat_hash_map& - instruction_schedule); - // Compare asynchronous copies such that an earlier start time has the same or // earlier end time and an earlier end time has the same or earlier start time. bool operator<(const AsynchronousCopy& a, const AsynchronousCopy& b); diff --git a/third_party/xla/xla/service/scan_loop_accumulator_input_unification.cc b/third_party/xla/xla/service/scan_loop_accumulator_input_unification.cc index 0ac01d68091a36..4b6823d1e5fe44 100644 --- a/third_party/xla/xla/service/scan_loop_accumulator_input_unification.cc +++ b/third_party/xla/xla/service/scan_loop_accumulator_input_unification.cc @@ -22,6 +22,7 @@ limitations under the License. #include #include "absl/container/flat_hash_set.h" +#include "absl/log/check.h" #include "absl/log/log.h" #include "absl/status/status_macros.h" #include "absl/status/statusor.h" @@ -78,39 +79,35 @@ FindAccumulatorInputPairs(const HloDataflowAnalysis& dataflow_analysis, // Finding the accumulator instructions std::vector possible_acc; - for (int64_t param_idx = 0; - param_idx < while_instr->while_init()->operand_count(); ++param_idx) { - for (HloInstruction* gte : body_param->users()) { - if (!Match(gte, match::GetTupleElement().WithTupleIndex(param_idx))) { - continue; - } - if (gte->operand(0) != body_param) { - continue; - } + for (HloInstruction* gte : body_param->users()) { + if (gte->opcode() != HloOpcode::kGetTupleElement) { + continue; + } + int64_t param_idx = gte->tuple_index(); + // HloVerifier ensures that param_idx >= 0. + CHECK_GE(param_idx, 0); - // The accumulator should only be used exactly once as the operand of - // dynamic-update-slice. - if (gte->user_count() > 1 || gte->user_count() == 0) { - continue; - } - HloInstruction* gte_user = gte->users().at(0); - if (MatchShapeCoveringDynamicIndexInstruction( - gte_user, gte, HloOpcode::kDynamicUpdateSlice, config) - .has_value()) { - // The accumulator should be written at the same index - if (computation->root_instruction()->mutable_operand(param_idx) == - gte_user) { - possible_acc.push_back(gte); - VLOG(3) << "accumulator index: " << param_idx << " = " << gte->name(); - } + // The accumulator should only be used exactly once as the operand of + // dynamic-update-slice. + if (gte->user_count() != 1) { + continue; + } + HloInstruction* gte_user = gte->users().at(0); + if (MatchShapeCoveringDynamicIndexInstruction( + gte_user, gte, HloOpcode::kDynamicUpdateSlice, config) + .has_value()) { + // The accumulator should be written at the same index + if (computation->root_instruction()->operand(param_idx) == gte_user) { + possible_acc.push_back(gte); + VLOG(3) << "accumulator index: " << param_idx << " = " << gte->name(); } } } - // If operand is actually an operand of the instr, returns the index of the - // operand, otherwise returns -1. - auto operand_index = [](HloInstruction* instr, - HloInstruction* operand) -> int64_t { + // Version of HloInstruction::operand_index that returns -1 if the operand is + // not found instead of crashing. + auto find_operand_index = [](const HloInstruction* instr, + const HloInstruction* operand) -> int64_t { for (int64_t i = 0; i < instr->operand_count(); ++i) { if (operand == instr->operand(i)) { return i; @@ -119,19 +116,14 @@ FindAccumulatorInputPairs(const HloDataflowAnalysis& dataflow_analysis, return -1; }; - // Returns the first GTE instruction in the parent computation of the tuple - // with the form of get-tuple-element(tuple), index=idx + // Returns the first GTE instruction in the users of the tuple with the form + // of get-tuple-element(tuple), index=idx auto find_gte_instr = [](HloInstruction* tuple, int64_t idx) -> HloInstruction* { - for (HloInstruction* instr : tuple->parent()->MakeInstructionPostOrder()) { - HloInstruction* operand; - if (Match(instr, match::GetTupleElement() - .WithOperand(0, match::Op(&operand)) - .WithTupleIndex(idx))) { - if (operand != tuple) { - continue; - } - return instr; + for (HloInstruction* user : tuple->users()) { + if (user->opcode() == HloOpcode::kGetTupleElement && + user->tuple_index() == idx) { + return user; } } return nullptr; @@ -154,8 +146,8 @@ FindAccumulatorInputPairs(const HloDataflowAnalysis& dataflow_analysis, if (acc_gte_outer_body == nullptr) { continue; } - int64_t idx = - operand_index(outer_while_body->root_instruction(), acc_gte_outer_body); + int64_t idx = find_operand_index(outer_while_body->root_instruction(), + acc_gte_outer_body); VLOG(3) << "Accumulator output of the scan in the outer body = " << acc_gte_outer_body->name() << ", index = " << idx; if (idx == -1) { @@ -174,18 +166,24 @@ FindAccumulatorInputPairs(const HloDataflowAnalysis& dataflow_analysis, // Find the corresponding gte in the body of the inner loop int64_t input_idx_inner = - operand_index(while_instr->while_init(), input_gte_outer); + find_operand_index(while_instr->while_init(), input_gte_outer); + if (input_idx_inner == -1) { + continue; + } HloInstruction* input_gte_inner = find_gte_instr(computation->parameter_instruction(0), input_idx_inner); - + // An unused loop input is technically valid, skip it. + if (input_gte_inner == nullptr) { + continue; + } if (!LoopIndexIsReadOnly(dataflow_analysis, while_instr, input_idx_inner)) { continue; } VLOG(3) << "Input parameter scan body = " << input_gte_inner->name() << ", index = " << input_gte_inner->tuple_index(); - // Input must have to users, one is the dynamic-slice and the other is the + // Input must have two users, one is the dynamic-slice and the other is the // root of the loop body. if (input_gte_inner->user_count() != 2) { continue; diff --git a/third_party/xla/xla/service/select_and_scatter_expander_test.cc b/third_party/xla/xla/service/select_and_scatter_expander_test.cc index c8f870fc091bc5..4223443d41c129 100644 --- a/third_party/xla/xla/service/select_and_scatter_expander_test.cc +++ b/third_party/xla/xla/service/select_and_scatter_expander_test.cc @@ -29,7 +29,7 @@ constexpr absl::string_view kModuleStr = %ge_F32.v3 (lhs: f32[], rhs: f32[]) -> pred[] { %lhs = f32[] parameter(0) %rhs = f32[] parameter(1) - ROOT %greater-than-or-equal-to = pred[] compare(f32[] %lhs, f32[] %rhs), direction=GE, type=TOTALORDER + ROOT %greater-than-or-equal-to = pred[] compare(f32[] %lhs, f32[] %rhs), direction=GE, order=TOTAL } %add_F32.v3 (lhs.1: f32[], rhs.1: f32[]) -> f32[] { diff --git a/third_party/xla/xla/service/spmd/fft_handler.cc b/third_party/xla/xla/service/spmd/fft_handler.cc index 0d9ca5d1bde171..388805f7008b59 100644 --- a/third_party/xla/xla/service/spmd/fft_handler.cc +++ b/third_party/xla/xla/service/spmd/fft_handler.cc @@ -366,7 +366,8 @@ absl::Status SpmdPartitioningVisitor::HandleFft(HloInstruction* hlo) { } // Support partition at the last dimension only. - if (!hlo->has_sharding() || + if (!hlo->has_sharding() || hlo->sharding().IsReplicated() || + hlo->sharding().dimensions().empty() || hlo->sharding().dimensions().back() != num_partitions_) { return DefaultAction(hlo); } diff --git a/third_party/xla/xla/service/spmd/gather_scatter_handler.cc b/third_party/xla/xla/service/spmd/gather_scatter_handler.cc index 9d22a72e8d29a3..be6880248d0afb 100644 --- a/third_party/xla/xla/service/spmd/gather_scatter_handler.cc +++ b/third_party/xla/xla/service/spmd/gather_scatter_handler.cc @@ -45,6 +45,7 @@ limitations under the License. #include "xla/literal.h" #include "xla/literal_util.h" #include "xla/service/collective_ops_utils.h" +#include "xla/service/spmd/shardy/constants.h" #include "xla/service/spmd/spmd_partitioner.h" #include "xla/service/spmd/spmd_partitioner_util.h" #include "xla/shape.h" @@ -655,6 +656,9 @@ absl::StatusOr PartitionGatherTrivialSlicedOperandDimensions( auto filtered = b->AddInstruction(HloInstruction::CreateTernary( pgather->shape(), HloOpcode::kSelect, broadcast_filter, CreateZero(pgather->shape(), b), pgather)); + if (gather->frontend_attributes().map().contains(sdy::kHasUnreducedAxes)) { + filtered->add_frontend_attribute(sdy::kHasUnreducedAxes, "true"); + } // All-reduce along trivially sliced dimensions. auto ar = operand.state().partitioner->AllReduceAlongShardingDims( b, filtered, original_operand_sharding, operand.state().next_channel_id, @@ -1088,6 +1092,9 @@ absl::Status SpmdPartitioningVisitor::HandleGatherWithoutConflicts( HloInstruction* filtered = b->AddInstruction(HloInstruction::CreateTernary( pgather->shape(), HloOpcode::kSelect, broadcast_filter, CreateZero(pgather->shape(), b), pgather)); + if (hlo->frontend_attributes().map().contains(sdy::kHasUnreducedAxes)) { + filtered->add_frontend_attribute(sdy::kHasUnreducedAxes, "true"); + } HloInstruction* ar = operand.state().partitioner->AllReduceAlongShardingDims( b, filtered, operand.sharding(), operand.state().next_channel_id, diff --git a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_manual_reduction_collectives.cc b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_manual_reduction_collectives.cc index 74f201123ae1a9..b7f562f58fbdbe 100644 --- a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_manual_reduction_collectives.cc +++ b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_manual_reduction_collectives.cc @@ -264,6 +264,9 @@ getAxesCoordinateAndSize(OpBuilder& builder, mlir::Location loc, void convertShardedToUnreduced(sdy::ShardedToUnreducedOp op, mlir::IRRewriter& rewriter) { TensorShardingAttr outSharding = op.getOutSharding(); + // We intentionally do not support non-sum reductions (e.g. `MAX`, `MIN`) + // because the operation semantics are specific to `SUM` and are not intended + // to be used with other reduction ops without an explicit use case. CHECK_EQ(outSharding.getReductionOp(), sdy::ReductionOp::SUM); MeshAttr mesh = outSharding.getMesh(op); // If the mesh does not have iota device ids, we need an extra step to convert @@ -340,6 +343,9 @@ void convertShardedToUnreduced(sdy::ShardedToUnreducedOp op, void convertReplicatedToUnreduced(sdy::ReplicatedToUnreducedOp op, mlir::IRRewriter& rewriter) { TensorShardingAttr outSharding = op.getOutSharding(); + // We intentionally do not support non-sum reductions (e.g. `MAX`, `MIN`) + // because the operation semantics are specific to `SUM` and are not intended + // to be used with other reduction ops without an explicit use case. CHECK_EQ(outSharding.getReductionOp(), sdy::ReductionOp::SUM); MeshAttr mesh = outSharding.getMesh(op); diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc b/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc index 4f264492bc84d8..82b26126abba40 100644 --- a/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc +++ b/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc @@ -9880,6 +9880,28 @@ ENTRY entry { } } +TEST_P(SpmdPartitioningTest, GatherPartitionedOnTrivialSliceDimsUnreduced) { + absl::string_view hlo_string = R"( +HloModule module + +ENTRY entry { + %input = f32[17,9] parameter(0), sharding={devices=[2,1]<=[2]} + %indices = s32[2,3] parameter(1), sharding={replicated} + ROOT %gather = f32[2,3,9] gather(%input, %indices), offset_dims={2}, + collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, + slice_sizes={1,9}, sharding={unreduced} +})"; + + SpmdPartitionerOptions options; + for (bool need_resolve_conflicts : {true, false}) { + options.need_resolve_conflicts = need_resolve_conflicts; + TF_ASSERT_OK_AND_ASSIGN( + auto module, + PartitionComputation(hlo_string, /*num_devices=*/2, options)); + EXPECT_EQ(FindInstruction(module.get(), HloOpcode::kAllReduce), nullptr); + } +} + TEST_P(SpmdPartitioningTest, GatherPartitionedOnTrivialSliceDims_PartialReplicate) { absl::string_view hlo_string = R"( @@ -13063,6 +13085,30 @@ ENTRY entry { EXPECT_TRUE(has_fft); } +TEST_P(SpmdPartitioningTest, Fft3DReplicatedShardingDoesNotCrash) { + // For an FFT instruction with replicated sharding, gspmd should not attempt + // to index into the empty sharding.dimensions(). + absl::string_view hlo_string = R"( +HloModule module + +ENTRY entry { + constant = c64[1,1,8] constant({{{(0,0),(1,1),(2,2),(3,3),(4,4),(5,5),(6,6),(7,7)}}}), + sharding={replicated} + ROOT fft = c64[1,1,8] fft(c64[1,1,8] constant), fft_type=FFT, fft_length={8}, + sharding={replicated} +} +)"; + + TF_ASSERT_OK_AND_ASSIGN(auto module, + PartitionComputation(hlo_string, /*num_devices=*/2)); + bool has_fft = false; + for (const HloInstruction* instr : + module->entry_computation()->instructions()) { + if (instr->opcode() == HloOpcode::kFft) has_fft = true; + } + EXPECT_TRUE(has_fft); +} + TEST_P(SpmdPartitioningTest, DotInputsAreIdentical) { absl::string_view hlo_string = R"( HloModule module @@ -17350,7 +17396,7 @@ region_695.22546 { Arg_3.22550 = s32[] parameter(3) Arg_0.22547 = bf16[] parameter(0) Arg_1.22548 = bf16[] parameter(1) - ROOT compare.22551 = pred[] compare(Arg_0.22547, Arg_1.22548), direction=GT, type=TOTALORDER + ROOT compare.22551 = pred[] compare(Arg_0.22547, Arg_1.22548), direction=GT, order=TOTAL } ENTRY %entry { @@ -17381,7 +17427,7 @@ region_695.22546 { Arg_3.22550 = s32[] parameter(3) Arg_0.22547 = bf16[] parameter(0) Arg_1.22548 = bf16[] parameter(1) - ROOT compare.22551 = pred[] compare(Arg_0.22547, Arg_1.22548), direction=GT, type=TOTALORDER + ROOT compare.22551 = pred[] compare(Arg_0.22547, Arg_1.22548), direction=GT, order=TOTAL } ENTRY %entry { @@ -17415,7 +17461,7 @@ region_695.22546 { Arg_3.22550 = s32[] parameter(3) Arg_0.22547 = bf16[] parameter(0) Arg_1.22548 = bf16[] parameter(1) - ROOT compare.22551 = pred[] compare(Arg_0.22547, Arg_1.22548), direction=GT, type=TOTALORDER + ROOT compare.22551 = pred[] compare(Arg_0.22547, Arg_1.22548), direction=GT, order=TOTAL } ENTRY %entry { @@ -17447,7 +17493,7 @@ region_695.22546 { Arg_3.22550 = s32[] parameter(3) Arg_0.22547 = bf16[] parameter(0) Arg_1.22548 = bf16[] parameter(1) - ROOT compare.22551 = pred[] compare(Arg_0.22547, Arg_1.22548), direction=GT, type=TOTALORDER + ROOT compare.22551 = pred[] compare(Arg_0.22547, Arg_1.22548), direction=GT, order=TOTAL } ENTRY %entry { @@ -17478,7 +17524,7 @@ region { Arg_3.22550 = s32[] parameter(3) Arg_0.22547 = bf16[] parameter(0) Arg_1.22548 = bf16[] parameter(1) - ROOT compare.22551 = pred[] compare(Arg_0.22547, Arg_1.22548), direction=GT, type=TOTALORDER + ROOT compare.22551 = pred[] compare(Arg_0.22547, Arg_1.22548), direction=GT, order=TOTAL } ENTRY %entry { diff --git a/third_party/xla/xla/service/topk_rewriter_test.cc b/third_party/xla/xla/service/topk_rewriter_test.cc index 158507df86c92e..2e689f64d7fd04 100644 --- a/third_party/xla/xla/service/topk_rewriter_test.cc +++ b/third_party/xla/xla/service/topk_rewriter_test.cc @@ -145,7 +145,7 @@ std::string getCompareComparator() { %Arg_1.101 = f32[] parameter(1) %Arg_2.102 = s32[] parameter(2) %Arg_3.103 = s32[] parameter(3) - ROOT %compare.56364 = pred[] compare(f32[] %Arg_0.100, f32[] %Arg_1.101), direction=GT, type=TOTALORDER + ROOT %compare.56364 = pred[] compare(f32[] %Arg_0.100, f32[] %Arg_1.101), direction=GT, order=TOTAL })"; } @@ -158,7 +158,7 @@ std::string getStableComparator() { %broadcast.40631 = pred[] broadcast(pred[] %constant.40630), dimensions={} %p.0.lhs.40626 = f32[] parameter(0) %p.0.rhs.40627 = f32[] parameter(1) - %compare.40632 = pred[] compare(f32[] %p.0.lhs.40626, f32[] %p.0.rhs.40627), direction=GT, type=TOTALORDER + %compare.40632 = pred[] compare(f32[] %p.0.lhs.40626, f32[] %p.0.rhs.40627), direction=GT, order=TOTAL ROOT %select.40633 = pred[] select(pred[] %broadcast.40631, pred[] %compare.40632, pred[] %broadcast.40631) })"; } @@ -766,7 +766,7 @@ HloModule topk c2 { p0 = f32[] parameter(0) p1 = f32[] parameter(1) - ROOT cmp = pred[] compare(p0, p1), direction=GT, type=TOTALORDER + ROOT cmp = pred[] compare(p0, p1), direction=GT, order=TOTAL } c4 { @@ -774,7 +774,7 @@ c4 { p1 = f32[] parameter(1) p2 = s32[] parameter(2) p3 = s32[] parameter(3) - ROOT cmp = pred[] compare(p0, p1), direction=GT, type=TOTALORDER + ROOT cmp = pred[] compare(p0, p1), direction=GT, order=TOTAL } ENTRY TopK { diff --git a/third_party/xla/xla/sh_test_with_runfiles.py b/third_party/xla/xla/sh_test_with_runfiles.py index 9362b3a5ecb11f..8d58713dc721e3 100644 --- a/third_party/xla/xla/sh_test_with_runfiles.py +++ b/third_party/xla/xla/sh_test_with_runfiles.py @@ -25,20 +25,28 @@ class ShTestWithRunfiles(lit.formats.ShTest): def execute(self, test, lit_config): runfiles_env = os.environ.get("RUNFILES_DIR") created_symlinks = [] + if runfiles_env: rf_path = pathlib.Path(runfiles_env) test_exec_dir = pathlib.Path(test.getExecPath()).parent test_exec_dir.mkdir(parents=True, exist_ok=True) - # Symlink all directories from runfiles root to test_exec_dir.parent - # RUNPATH has "../+rocm_configure_ext+local_config_rocm/..." patterns + # Symlink the entire runfiles structure to match what RUNPATH expects + # Binaries have RUNPATH like $ORIGIN/../../../../../_solib_x86_64/... + # They execute from lit_bin which is at runfiles/_main/xla/.../lit_bin/ + # So we need _main accessible from there + # Symlink to test_exec_dir.parent to match RUNPATH patterns for item in rf_path.iterdir(): if item.is_dir(): test_exec_symlink = test_exec_dir.parent / item.name if not test_exec_symlink.exists(): try: - test_exec_symlink.symlink_to(item, target_is_directory=True) + # Use relative symlinks for portability between local and RBE + relative_target = os.path.relpath(item, test_exec_dir.parent) + test_exec_symlink.symlink_to( + relative_target, target_is_directory=True + ) created_symlinks.append(test_exec_symlink) except FileExistsError: pass @@ -74,6 +82,24 @@ def execute(self, test, lit_config): f"{existing_flags} {flag}".strip() ) + # For --dynamic_mode=fully, set LD_LIBRARY_PATH to point to + # library directories. + # Support both bzlmod and workspace mode layouts + lib_dirs = [] + + for candidate_root in [rf_path / "_main", rf_path / "xla", rf_path]: + if candidate_root.is_dir(): + for solib in candidate_root.glob("_solib_*"): + if solib.is_dir(): + lib_dirs.append(str(solib)) + + if lib_dirs: + existing_ld_path = test.config.environment.get("LD_LIBRARY_PATH", "") + new_ld_path = ":".join(lib_dirs) + if existing_ld_path: + new_ld_path = f"{new_ld_path}:{existing_ld_path}" + test.config.environment["LD_LIBRARY_PATH"] = new_ld_path + result = super().execute(test, lit_config) # Clean up created symlinks diff --git a/third_party/xla/xla/shape_util.cc b/third_party/xla/xla/shape_util.cc index 668ff5abe239b8..4db4a138d7b5b9 100644 --- a/third_party/xla/xla/shape_util.cc +++ b/third_party/xla/xla/shape_util.cc @@ -1086,7 +1086,11 @@ Shape ShapeUtil::PrependMajorDimension(int64_t bound, Shape shape) { if (subshape.is_dynamic()) { size += sizeof(DynamicSizeType) * subshape.dimensions().size(); } - if (primitive_util::IsSubByteNonPredType(subshape.element_type())) { + if (subshape.element_type() == PRED) { + // PRED is packed 8 elements per byte. + size += CeilOfRatio(ElementsIn(subshape), 8); + } else if (primitive_util::IsSubByteNonPredType( + subshape.element_type())) { // 4-bit types are packed 2 elements per byte. size += CeilOfRatio( ElementsIn(subshape), diff --git a/third_party/xla/xla/stream_executor/command_buffer.h b/third_party/xla/xla/stream_executor/command_buffer.h index 5e3b34410226ad..e10feee73baadc 100644 --- a/third_party/xla/xla/stream_executor/command_buffer.h +++ b/third_party/xla/xla/stream_executor/command_buffer.h @@ -328,6 +328,13 @@ class CommandBuffer { UpdateCommands update_cond, UpdateCommands update_body) = 0; + // Adds a host node (callback) to the command buffer. + // The `callback` will be executed on the CPU (host) when this node is + // processed during command buffer execution. + virtual absl::StatusOr CreateHost( + absl::AnyInvocable callback, + absl::Span dependencies) = 0; + // Set the priority of all nodes in the command buffer. virtual absl::Status SetPriority(StreamPriority priority) = 0; diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.cc b/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.cc index 29e7dc349269f9..a9f5c52482a0b1 100644 --- a/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.cc +++ b/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.cc @@ -65,6 +65,61 @@ namespace stream_executor::gpu { namespace { constexpr bool kHasCuda12090 = CUDA_VERSION >= 12090; +// Wraps a C++ object of type `T` and binds its lifetime to a CUDA graph using +// `CUuserObject` mechanisms. This ensures that the underlying C++ resource +// is retained as long as the CUDA graph exists, +// and gets correctly deleted when the graph is destroyed. +template +class CudaUserObject { + public: + CudaUserObject(const CudaUserObject&) = delete; + CudaUserObject& operator=(const CudaUserObject&) = delete; + + // Creates and initializes a `CudaUserObject`. + // The object `instance` is wrapped, registered as a `CUuserObject` with CUDA, + // and attached to the supplied `graph` to retain its reference. + static absl::StatusOr Create(T&& instance, CUgraph graph) { + auto obj_ptr = absl::WrapUnique(new CudaUserObject(std::move(instance))); + + ABSL_RETURN_IF_ERROR(cuda::ToStatus(cuUserObjectCreate( + &obj_ptr->handle_, obj_ptr.get(), + [](void* user_data) { + auto* cb = static_cast*>(user_data); + delete cb; + }, + 1, CU_USER_OBJECT_NO_DESTRUCTOR_SYNC))); + + CudaUserObject* res = std::move(obj_ptr).release(); + ABSL_RETURN_IF_ERROR(res->AttachTo(graph)); + + return res; + } + + T* get() { return &instance_; } + CUuserObject handle() { return handle_; } + + private: + explicit CudaUserObject(T instance) : instance_(std::move(instance)) {} + + absl::Status AttachTo(CUgraph graph) { + absl::Status res = cuda::ToStatus( + cuGraphRetainUserObject(graph, handle_, 1, CU_GRAPH_USER_OBJECT_MOVE)); + if (!res.ok()) { + cuUserObjectRelease(handle_, 1); + } + return res; + } + + T instance_; + CUuserObject handle_; +}; + +void HostCallbackTrampoline(void* user_data) { + auto* user_object = + static_cast>*>(user_data); + (*user_object->get())(); +} + template void LogAppend(std::string& out, const absl::FormatSpec& format, const Args&... args) { @@ -137,7 +192,7 @@ std::string FormatGraphConditionalHandles( } // namespace // Converts a list of platform independent GraphNodeHandles into a list of -// CUDA specific CUgraphNode. +// CUDA specific CUgraphNodes. std::vector CudaCommandBuffer::ToCudaGraphHandles( absl::Span opaque_handles) { std::vector handles; @@ -298,7 +353,7 @@ CudaCommandBuffer::CreateConditionalNode( "Conditional nodes require CUDA driver version >= 12.3"); } - // Add conditional node to a graph. + // Add a conditional node to a graph. VLOG(2) << "Add conditional node to a graph " << graph_ << "; type: " << ConditionalTypeToString(type) << "; conditional: " << conditional << "; deps(" @@ -399,7 +454,7 @@ absl::Status CudaCommandBuffer::UpdateMemsetNode(GraphNodeHandle node_handle, absl::StatusOr CudaCommandBuffer::CreateMemcpyD2DNode( absl::Span dependencies, DeviceAddressBase destination, DeviceAddressBase source, uint64_t size) { - VLOG(2) << "Add memcpy d2d node to a graph " << graph_ + VLOG(2) << "Add memcpy D2D node to a graph " << graph_ << "; dst: " << destination.opaque() << "; src: " << source.opaque() << "; size: " << size << "; context: " << cuda_context_->context() << "; deps(" << dependencies.size() @@ -420,14 +475,14 @@ absl::StatusOr CudaCommandBuffer::CreateMemcpyD2DNode( ABSL_RETURN_IF_ERROR(cuda::ToStatus( cuGraphAddMemcpyNode(&node_handle, graph_, deps.data(), deps.size(), ¶ms, cuda_context_->context()), - "Failed to add memcpy d2d node to a CUDA graph")); + "Failed to add memcpy D2D node to a CUDA graph")); return FromCudaGraphHandle(node_handle); } absl::Status CudaCommandBuffer::UpdateMemcpyD2DNode( GraphNodeHandle node_handle, DeviceAddressBase destination, DeviceAddressBase source, uint64_t size) { - VLOG(2) << "Set memcpy d2d node params " << node_handle + VLOG(2) << "Set memcpy D2D node params " << node_handle << " in graph executable " << graph_exec() << "; dst: " << destination.opaque() << "; src: " << source.opaque() << "; size: " << size << "; context: " << cuda_context_->context(); @@ -443,13 +498,13 @@ absl::Status CudaCommandBuffer::UpdateMemcpyD2DNode( return cuda::ToStatus(cuGraphExecMemcpyNodeSetParams( graph_exec(), ToCudaGraphHandle(node_handle), ¶ms, cuda_context_->context()), - "Failed to set memcpy d2d node params"); + "Failed to set memcpy D2D node params"); } absl::StatusOr CudaCommandBuffer::CreateMemcpyD2HNode( absl::Span dependencies, void* destination, DeviceAddressBase source, uint64_t size) { - VLOG(2) << "Add memcpy d2h node to a graph " << graph_ + VLOG(2) << "Add memcpy D2H node to a graph " << graph_ << "; dst: " << destination << "; src: " << source.opaque() << "; size: " << size << "; context: " << cuda_context_->context() << "; deps(" << dependencies.size() @@ -470,7 +525,7 @@ absl::StatusOr CudaCommandBuffer::CreateMemcpyD2HNode( ABSL_RETURN_IF_ERROR(cuda::ToStatus( cuGraphAddMemcpyNode(&node_handle, graph_, deps.data(), deps.size(), ¶ms, cuda_context_->context()), - "Failed to add memcpy d2h node to a CUDA graph")); + "Failed to add memcpy D2H node to a CUDA graph")); return FromCudaGraphHandle(node_handle); } @@ -478,7 +533,7 @@ absl::Status CudaCommandBuffer::UpdateMemcpyD2HNode(GraphNodeHandle node_handle, void* destination, DeviceAddressBase source, uint64_t size) { - VLOG(2) << "Set memcpy d2h node params " << node_handle + VLOG(2) << "Set memcpy D2H node params " << node_handle << " in graph executable " << graph_exec() << "; dst: " << destination << "; src: " << source.opaque() << "; size: " << size << "; context: " << cuda_context_->context(); @@ -494,13 +549,13 @@ absl::Status CudaCommandBuffer::UpdateMemcpyD2HNode(GraphNodeHandle node_handle, return cuda::ToStatus(cuGraphExecMemcpyNodeSetParams( graph_exec(), ToCudaGraphHandle(node_handle), ¶ms, cuda_context_->context()), - "Failed to set memcpy d2h node params"); + "Failed to set memcpy D2H node params"); } absl::StatusOr CudaCommandBuffer::CreateMemcpyH2DNode( absl::Span dependencies, DeviceAddressBase destination, const void* source, uint64_t size) { - VLOG(2) << "Add memcpy h2d node to a graph " << graph_ + VLOG(2) << "Add memcpy H2D node to a graph " << graph_ << "; dst: " << destination.opaque() << "; src: " << source << "; size: " << size << "; context: " << cuda_context_->context() << "; deps(" << dependencies.size() @@ -521,14 +576,14 @@ absl::StatusOr CudaCommandBuffer::CreateMemcpyH2DNode( ABSL_RETURN_IF_ERROR(cuda::ToStatus( cuGraphAddMemcpyNode(&node_handle, graph_, deps.data(), deps.size(), ¶ms, cuda_context_->context()), - "Failed to add memcpy h2d node to a CUDA graph")); + "Failed to add memcpy H2D node to a CUDA graph")); return FromCudaGraphHandle(node_handle); } absl::Status CudaCommandBuffer::UpdateMemcpyH2DNode( GraphNodeHandle node_handle, DeviceAddressBase destination, const void* source, uint64_t size) { - VLOG(2) << "Set memcpy h2d node params " << node_handle + VLOG(2) << "Set memcpy H2D node params " << node_handle << " in graph executable " << graph_exec() << "; dst: " << destination.opaque() << "; src: " << source << "; size: " << size << "; context: " << cuda_context_->context(); @@ -544,7 +599,7 @@ absl::Status CudaCommandBuffer::UpdateMemcpyH2DNode( return cuda::ToStatus(cuGraphExecMemcpyNodeSetParams( graph_exec(), ToCudaGraphHandle(node_handle), ¶ms, cuda_context_->context()), - "Failed to set memcpy h2d node params"); + "Failed to set memcpy H2D node params"); } absl::Status CudaCommandBuffer::PopulateDnnGraphNode( @@ -617,6 +672,51 @@ absl::Status CudaCommandBuffer::UpdateClonedChildNode( "Failed to set CUDA graph child node params"); } +absl::StatusOr CudaCommandBuffer::CreateHostNode( + absl::Span dependencies, + absl::AnyInvocable callback) { + XLA_VLOG_DEVICE(2, stream_exec_->device_ordinal()) + << "Add host node to a graph " << graph_; + + std::string log_msg; + absl::Cleanup cleanup = [&] { + XLA_VLOG_DEVICE(5, stream_exec_->device_ordinal()) << log_msg; + }; + + CUgraphNodeParams cu_params = {}; + cu_params.type = CU_GRAPH_NODE_TYPE_HOST; + + CUDA_HOST_NODE_PARAMS_v2& params = cu_params.host; + params.fn = &HostCallbackTrampoline; + + ABSL_ASSIGN_OR_RETURN(params.userData, + CudaUserObject>::Create( + std::move(callback), graph_)); + + std::vector deps = ToCudaGraphHandles(dependencies); + std::vector edge_data; + edge_data.reserve(deps.size()); + for (size_t i = 0; i < deps.size(); ++i) { + CUgraphEdgeData edge_data_item = {}; + CUgraphNodeType type; + ABSL_RETURN_IF_ERROR(cuda::ToStatus( + cuGraphNodeGetType(deps[i], &type), + absl::StrCat("Failed to get CUDA graph node type for dependency ", i))); + LogAppend(log_msg, " dep %d node: %p, type: %d", i, deps[i], type); + edge_data.push_back(edge_data_item); + } + + CUgraphNode node_handle = nullptr; + ABSL_RETURN_IF_ERROR(cuda::ToStatus( + cuGraphAddNode_v2(&node_handle, graph_, deps.data(), edge_data.data(), + deps.size(), &cu_params), + "Failed to add host node to a CUDA graph")); + + LogAppend(log_msg, "CudaCommandBuffer::CreateHostNode: created node: %p", + node_handle); + return FromCudaGraphHandle(node_handle); +} + absl::StatusOr CudaCommandBuffer::CreateKernelNode( absl::Span dependencies, StreamPriority priority, const ThreadDim& threads, const BlockDim& blocks, @@ -839,9 +939,9 @@ absl::Status CudaCommandBuffer::Trace( graph_, /*dependencies=*/nullptr, /*dependency_data=*/nullptr, /*num_dependencies=*/0, // THREAD_LOCAL implies that capturing is done only on the current - // stream. Cuda calls can be made on other streams without + // stream. CUDA calls can be made on other streams without // interrupting the capture. - // The default mode CU_STREAM_CAPTURE_MODE_GLOBAL, will capture at + // The default mode CU_STREAM_CAPTURE_MODE_GLOBAL will capture at // at a global level. That would stall everything at a driver level. CU_STREAM_CAPTURE_MODE_THREAD_LOCAL)); Stream* capture_stream = capture_handle.capturing_stream(); @@ -920,7 +1020,7 @@ absl::Status CudaCommandBuffer::SetPriority(StreamPriority priority) { absl::Status CudaCommandBuffer::PrepareFinalization() { if (stream_exec_->GetDeviceDescription().driver_version() < SemanticVersion{12, 8, 0}) { - // For CUDA < 12080, cuda graph conditional node does not support + // For CUDA < 12080, a CUDA graph conditional node does not support // empty body graph. ABSL_ASSIGN_OR_RETURN(auto node_count, GetNodeCount()); if (node_count > 0) { diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.h b/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.h index 1e73aff2c0557d..1145053f600854 100644 --- a/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.h +++ b/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.h @@ -163,6 +163,10 @@ class CudaCommandBuffer final : public GpuCommandBuffer { absl::Status UpdateClonedChildNode(GraphNodeHandle node_handle, const CommandBuffer& nested) override; + absl::StatusOr CreateHostNode( + absl::Span dependencies, + absl::AnyInvocable callback) override; + absl::StatusOr CreateKernelNode( absl::Span dependencies, StreamPriority priority, const ThreadDim& threads, const BlockDim& blocks, diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer_test.cc b/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer_test.cc index d577f73d484c1b..6113615cdd1f84 100644 --- a/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer_test.cc +++ b/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer_test.cc @@ -304,6 +304,35 @@ TEST(CudaCommandBufferTest, LaunchClusterKernelWithClusterDimsSucceeds) { ASSERT_OK(stream->BlockHostUntilDone()); } +TEST(CudaCommandBufferTest, LaunchHostCallback) { + Platform* platform = CudaPlatform(); + ASSERT_OK_AND_ASSIGN(StreamExecutor * executor, + platform->ExecutorForDevice(0)); + if (!executor->GetDeviceDescription() + .cuda_compute_capability() + .IsAtLeastVolta()) { + GTEST_SKIP() << "Requires at least a Volta GPU."; + } + ASSERT_OK_AND_ASSIGN(std::unique_ptr stream, + executor->CreateStream()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr cmd_buffer, + executor->CreateCommandBuffer(primary)); + + int counter = 0; + ASSERT_OK_AND_ASSIGN(const CommandBuffer::Command* cmd, + cmd_buffer->CreateHost([&]() { counter++; }, {})); + ASSERT_NE(cmd, nullptr); + + ASSERT_OK(cmd_buffer->Finalize()); + ASSERT_OK(cmd_buffer->Submit(stream.get())); + ASSERT_OK(stream->BlockHostUntilDone()); + EXPECT_EQ(counter, 1); + + ASSERT_OK(cmd_buffer->Submit(stream.get())); + ASSERT_OK(stream->BlockHostUntilDone()); + EXPECT_EQ(counter, 2); +} + TEST(CudaCommandBufferTest, MemcpyH2D2H) { Platform* platform = CudaPlatform(); ASSERT_OK_AND_ASSIGN(StreamExecutor * executor, diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_executor.cc b/third_party/xla/xla/stream_executor/cuda/cuda_executor.cc index ca3ace26d0aee6..8c6ad20a798760 100644 --- a/third_party/xla/xla/stream_executor/cuda/cuda_executor.cc +++ b/third_party/xla/xla/stream_executor/cuda/cuda_executor.cc @@ -218,13 +218,13 @@ absl::StatusOr LoadPtx(Context* context, const char* ptx_contents) { if (!status.ok()) { XLA_LOG_DEVICE(ERROR, context->device_ordinal()) - << "failed to load PTX text as a module: " << status; + << "Failed to load PTX text as a module: " << status; // As a precaution for null termination of the API-provided value, // ensure that at least the last byte is null. error_log_buffer[error_log_buffer_bytes ? error_log_buffer_bytes - 1 : 0] = '\0'; XLA_LOG_DEVICE(ERROR, context->device_ordinal()) - << "error log buffer (" << error_log_buffer_bytes + << "Error log buffer (" << error_log_buffer_bytes << " bytes): " << error_log_buffer.data(); if (absl::StrContains(error_log_buffer.data(), "Register allocation failed")) { @@ -311,7 +311,7 @@ void UnloadCudaModule(Context* context, CUmodule module) { auto status = cuda::ToStatus(cuModuleUnload(module)); if (!status.ok()) { XLA_LOG_DEVICE(ERROR, context->device_ordinal()) - << "failed to unload module " << module << "; leaking: " << status; + << "Failed to unload module " << module << "; leaking: " << status; } } @@ -487,7 +487,7 @@ bool GetDeviceTotalMemory(CUdevice device, uint64_t* result) { auto status = cuda::ToStatus(cuDeviceTotalMem(&value, device)); if (!status.ok()) { XLA_LOG_DEVICE(ERROR, device) - << "failed to query total available memory: " << status; + << "Failed to query total available memory: " << status; return false; } @@ -500,7 +500,7 @@ bool IsEccEnabled(CUdevice device, bool* result) { auto status = cuda::ToStatus( cuDeviceGetAttribute(&value, CU_DEVICE_ATTRIBUTE_ECC_ENABLED, device)); if (!status.ok()) { - XLA_LOG_DEVICE(ERROR, device) << "failed to query ECC status: " << status; + XLA_LOG_DEVICE(ERROR, device) << "Failed to query ECC status: " << status; return false; } @@ -517,7 +517,7 @@ std::string GetPCIBusID(CUdevice device) { cuDeviceGetPCIBusId(raw_pci_bus_id.data(), kBufferSize, device)); if (!status.ok()) { XLA_LOG_DEVICE(ERROR, device) - << "failed to query PCI bus id for device: " << status; + << "Failed to query PCI bus id for device: " << status; return ""; } if (!absl::c_linear_search(raw_pci_bus_id, '\0')) { @@ -534,7 +534,7 @@ bool HostRegister(Context* context, void* location, uint64_t size) { auto status = cuda::ToStatus( cuMemHostRegister(location, size, CU_MEMHOSTREGISTER_PORTABLE)); if (!status.ok()) { - LOG(ERROR) << "error registering host memory at " << location << ": " + LOG(ERROR) << "Error registering host memory at " << location << ": " << status; return false; } @@ -1094,7 +1094,7 @@ absl::StatusOr> CudaExecutor::LoadKernel( CUmodule module = gpu_binary_to_module_.at(module_handle).module; XLA_VLOG_DEVICE(2, device_ordinal()) - << "getting function " << kernel_name << " from module " << module; + << "Getting function " << kernel_name << " from module " << module; ABSL_ASSIGN_OR_RETURN( CUfunction function, GetModuleFunction(cuda_context_, module, kernel_name.c_str())); @@ -1450,7 +1450,7 @@ absl::Status CudaExecutor::SynchronousMemcpy(DeviceAddressBase* gpu_dst, xla::XlaFormatDevice(device_ordinal()), AsCudaDevicePtr(gpu_dst), host_src, size, size))); XLA_VLOG_DEVICE(2, device_ordinal()) - << "successfully enqueued sync memcpy h2d of " << size << " bytes"; + << "Successfully enqueued sync memcpy H2D of " << size << " bytes"; return absl::OkStatus(); } @@ -1464,7 +1464,7 @@ absl::Status CudaExecutor::SynchronousMemcpy(void* host_dst, "host dst: %p; GPU src: %llx; size: %u=0x%x", xla::XlaFormatDevice(device_ordinal()), host_dst, AsCudaDevicePtr(gpu_src), size, size))); - XLA_VLOG_DEVICE(2, device_ordinal()) << "successfully sync memcpy'd d2h of " + XLA_VLOG_DEVICE(2, device_ordinal()) << "Successfully sync memcpy'd D2H of " << size << " bytes to " << host_dst; return absl::OkStatus(); } diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_stream.cc b/third_party/xla/xla/stream_executor/cuda/cuda_stream.cc index 2387399ad187fe..9cb2828e6242a0 100644 --- a/third_party/xla/xla/stream_executor/cuda/cuda_stream.cc +++ b/third_party/xla/xla/stream_executor/cuda/cuda_stream.cc @@ -92,7 +92,7 @@ absl::StatusOr CreateStream(StreamExecutor* executor, int priority) { cuStreamCreateWithPriority(&stream, CU_STREAM_NON_BLOCKING, priority))); } - VLOG(2) << "successfully created stream " << stream << " for executor " + VLOG(2) << "Successfully created stream " << stream << " for executor " << executor << " on thread"; return stream; } @@ -115,7 +115,7 @@ absl::Status AsynchronousMemcpyD2H(StreamExecutor* executor, void* host_dst, ABSL_RETURN_IF_ERROR( cuda::ToStatus(cuMemcpyDtoHAsync(host_dst, gpu_src, size, stream))); - VLOG(2) << "successfully enqueued async memcpy d2h of " << size + VLOG(2) << "Successfully enqueued async memcpy D2H of " << size << " bytes from " << absl::bit_cast(gpu_src) << " to " << host_dst << " on stream " << stream; return absl::OkStatus(); @@ -128,7 +128,7 @@ absl::Status AsynchronousMemcpyH2D(StreamExecutor* executor, ABSL_RETURN_IF_ERROR( cuda::ToStatus(cuMemcpyHtoDAsync(gpu_dst, host_src, size, stream))); - VLOG(2) << "successfully enqueued async memcpy h2d of " << size << " bytes" + VLOG(2) << "Successfully enqueued async memcpy H2D of " << size << " bytes" << " from " << host_src << " to " << absl::bit_cast(gpu_dst) << " on stream " << stream; return absl::OkStatus(); @@ -178,7 +178,7 @@ absl::Status AsynchronousMemcpyD2D(StreamExecutor* executor, } } - VLOG(2) << "successfully enqueued async memcpy d2d of " << size << " bytes" + VLOG(2) << "Successfully enqueued async memcpy D2H of " << size << " bytes" << " from " << absl::bit_cast(gpu_src) << " to " << absl::bit_cast(gpu_dst) << " on stream " << stream; return absl::OkStatus(); @@ -328,15 +328,15 @@ void DestroyStream(StreamExecutor* executor, CUstream stream) { std::unique_ptr activation = executor->Activate(); CUresult res = cuStreamQuery(stream); if (res != CUDA_SUCCESS) { - LOG(ERROR) << "stream not idle on destroy: " << cuda::ToStatus(res); + LOG(ERROR) << "Stream not idle on destroy: " << cuda::ToStatus(res); } auto status = cuda::ToStatus(cuStreamDestroy(stream)); if (!status.ok()) { - LOG(ERROR) << "failed to destroy CUDA stream for executor " << executor + LOG(ERROR) << "Failed to destroy CUDA stream for executor " << executor << ": " << status; } else { - VLOG(2) << "successfully destroyed stream " << stream << " for executor " + VLOG(2) << "Successfully destroyed stream " << stream << " for executor " << executor; } } @@ -533,7 +533,7 @@ absl::Status LaunchCudaKernel( std::unique_ptr activation = executor->Activate(); if (VLOG_IS_ON(2)) { - std::string msg = absl::StrCat("launching kernel: ", kernel_name); + std::string msg = absl::StrCat("Launching kernel: ", kernel_name); if (cluster_dims.has_value()) { absl::StrAppend(&msg, "; cdx: ", cluster_dims->x, " cdy: ", cluster_dims->y, " cdz: ", cluster_dims->z); diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.cc b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.cc index 8719039f1d67f7..0a06de913f30ea 100644 --- a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.cc +++ b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.cc @@ -606,6 +606,18 @@ absl::Status GpuCommandBuffer::UpdateWhile(const Command* command, return absl::OkStatus(); } +absl::StatusOr GpuCommandBuffer::CreateHost( + absl::AnyInvocable callback, + absl::Span dependencies) { + ABSL_RETURN_IF_ERROR(CheckInState(State::kCreate)); + + ABSL_ASSIGN_OR_RETURN(GraphNodeHandle handle, + CreateHostNode(ToGraphNodeDependencies(dependencies), + std::move(callback))); + + return AppendCommand(GpuCommand{handle}); +} + absl::Status GpuCommandBuffer::Finalize() { ABSL_RETURN_IF_ERROR(CheckNotFinalized()); ABSL_RETURN_IF_ERROR(PrepareFinalization()); diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.h b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.h index 3bfaaa70999d58..3053c21b5dcb0f 100644 --- a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.h +++ b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.h @@ -218,6 +218,10 @@ class GpuCommandBuffer : public CommandBuffer { UpdateCommands update_cond, UpdateCommands update_body) override; + absl::StatusOr CreateHost( + absl::AnyInvocable callback, + absl::Span dependencies) override; + absl::Status Finalize() override; absl::Status Update() override; absl::Status Submit(Stream* stream) override; @@ -423,6 +427,10 @@ class GpuCommandBuffer : public CommandBuffer { const BlockDim& blocks, const std::optional& cluster_dims, const Kernel& kernel, const KernelArgsPackedArrayBase& args) = 0; + virtual absl::StatusOr CreateHostNode( + absl::Span dependencies, + absl::AnyInvocable callback) = 0; + //===--------------------------------------------------------------------===// // Launches an instantiated graph. Only supported on primary command diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_cudamallocasync_allocator.cc b/third_party/xla/xla/stream_executor/gpu/gpu_cudamallocasync_allocator.cc index f8779d7dc11812..51badb365e0d83 100644 --- a/third_party/xla/xla/stream_executor/gpu/gpu_cudamallocasync_allocator.cc +++ b/third_party/xla/xla/stream_executor/gpu/gpu_cudamallocasync_allocator.cc @@ -44,8 +44,8 @@ namespace stream_executor { struct GpuCudaMallocAsyncAllocator::CudaState { // cudaMallocAsync is stream aware. But TF StreamExecutor use only 1 - // compute stream and already synchronize with the h2d, d2h and d2d - // stream. So we do not need to ask cudaMallocAsync to add extra + // compute stream and already synchronize with the H2D, D2H and D2D + // stream. So, we do not need to ask cudaMallocAsync to add extra // synchronization. // Not owned. CUstream cuda_stream{}; @@ -136,14 +136,15 @@ GpuCudaMallocAsyncAllocator::GpuCudaMallocAsyncAllocator( << " We detected a version compatible with: " << driverVersion; } - // WAR an CUDA 11.2 driver bug for multiple-GPU. It currently + // Work around a CUDA 11.2 driver bug for multiple-GPU. It currently // request that the context on GPU 0 is initialized. Which isn't the // case for TF+horovod. if (platform_device_id.value() > 0 && driverVersion < 11030) { - CUcontext pctx; // We loose track of it. But this is fine. - if (auto result = cuDevicePrimaryCtxRetain(&pctx, 0)) + CUcontext pctx; // We lose track of it, but this is fine. + if (auto result = cuDevicePrimaryCtxRetain(&pctx, 0)) { LOG(FATAL) // Crash OK. << "Failed to retain context: " << cuda::ToStatus(result); + } } std::unique_ptr scoped_activation = stream_exec_->Activate(); @@ -165,13 +166,14 @@ GpuCudaMallocAsyncAllocator::GpuCudaMallocAsyncAllocator( << " Current driver: " << driverVersion << ". Failed to get device attribute : " << cuda::ToStatus(status); } - if (!cuda_malloc_async_supported) + if (!cuda_malloc_async_supported) { LOG(FATAL) // Crash OK. << "TF_GPU_ALLOCATOR=cuda_malloc_async isn't currently supported on " << "GPU id " << platform_device_id.value() << ":" << " Possible causes: device not supported (request SM60+), driver too " "old, " << " OS not supported, CUDA version too old(request CUDA11.2+)."; + } size_t pool_size; if (create_new_pool_) { @@ -185,15 +187,17 @@ GpuCudaMallocAsyncAllocator::GpuCudaMallocAsyncAllocator( #if CUDA_VERSION >= 12030 pool_props.maxSize = new_pool_size; #endif // CUDA_VERSION >= 12030 - if (auto status = cuMemPoolCreate(&cuda_state_->pool, &pool_props)) + if (auto status = cuMemPoolCreate(&cuda_state_->pool, &pool_props)) { LOG(FATAL) << // Crash OK. "Failed to create CUDA pool: " << cuda::ToStatus(status); + } } else { pool_size = reserve_memory_size; if (auto status = cuDeviceGetDefaultMemPool(&cuda_state_->pool, - platform_device_id.value())) + platform_device_id.value())) { LOG(FATAL) << // Crash OK. "Failed to get default CUDA pool: " << cuda::ToStatus(status); + } VLOG(2) << "using default memory pool " << cuda_state_->pool; } @@ -203,9 +207,10 @@ GpuCudaMallocAsyncAllocator::GpuCudaMallocAsyncAllocator( uint64_t release_threshold_64 = reserve_memory_size; if (auto status = cuMemPoolSetAttribute(cuda_state_->pool, CU_MEMPOOL_ATTR_RELEASE_THRESHOLD, - &release_threshold_64)) + &release_threshold_64)) { LOG(FATAL) << // Crash OK. "Failed to set CUDA pool attribute: " << cuda::ToStatus(status); + } if (compute_stats) { stats_ = std::make_unique(); @@ -213,7 +218,7 @@ GpuCudaMallocAsyncAllocator::GpuCudaMallocAsyncAllocator( } // If not set, it means we do not compute stats. // If in TF_DETERMINISTIC_ALLOCATOR is set, then make the allocator behave - // determistically. + // deterministically. bool deterministic = false; CHECK_OK(tsl::ReadBoolFromEnvVar("TF_DETERMINISTIC_ALLOCATOR", /*default_val=*/false, &deterministic)); @@ -319,8 +324,9 @@ GpuCudaMallocAsyncAllocator::~GpuCudaMallocAsyncAllocator() { if (create_new_pool_) { VLOG(2) << "Delete memory pool " << reinterpret_cast(cuda_state_->pool); - if (auto status = cuMemPoolDestroy(cuda_state_->pool)) + if (auto status = cuMemPoolDestroy(cuda_state_->pool)) { LOG(FATAL) << "Failed to destroy memory pool:" << cuda::ToStatus(status); + } } } @@ -346,7 +352,7 @@ void* GpuCudaMallocAsyncAllocator::AllocateRaw(size_t alignment, cuMemAllocFromPoolAsync(reinterpret_cast(&ptr), num_bytes, cuda_state_->pool, cuda_state_->cuda_stream); if (result == CUDA_ERROR_OUT_OF_MEMORY) { - // Doing a stream synchronization give the driver more flexibility + // Doing a stream synchronization gives the driver more flexibility // for blocks coalescing and doing memory remapping. So it can // solve some OOM cases when memory is tight. cuStreamSynchronize(cuda_state_->cuda_stream); @@ -396,7 +402,9 @@ void* GpuCudaMallocAsyncAllocator::AllocateRaw(size_t alignment, } void GpuCudaMallocAsyncAllocator::DeallocateRaw(void* ptr) { - if (ptr == nullptr) return; + if (ptr == nullptr) { + return; + } // The lock is only needed when stats are enabled, but it must be around // the cuMemFreeAsync call as well to ensure consistency of the stats update. std::optional lock; @@ -406,7 +414,7 @@ void GpuCudaMallocAsyncAllocator::DeallocateRaw(void* ptr) { if (auto result = cuMemFreeAsync(reinterpret_cast(ptr), cuda_state_->cuda_stream)) { if (result == CUDA_ERROR_DEINITIALIZED) { - // It happens with multi-GPU that TF free the GPU allocation after + // It happens with multi-GPU that TF frees the GPU allocations after // the driver is unloaded. It is safe to ignore this error here. // TODO: Find how to fix the shutdown steps in TF. VLOG(1) << "Ignoring CUDA error: " << cuda::ToStatus(result); @@ -444,25 +452,33 @@ bool GpuCudaMallocAsyncAllocator::TracksAllocationSizes() const { } size_t GpuCudaMallocAsyncAllocator::RequestedSize(const void* ptr) const { - if (!stats_ || !ptr) return 0; + if (!stats_ || !ptr) { + return 0; + } absl::MutexLock l(mutex_); return size_map_.at(ptr); } size_t GpuCudaMallocAsyncAllocator::AllocatedSize(const void* ptr) const { - if (!stats_ || !ptr) return 0; + if (!stats_ || !ptr) { + return 0; + } absl::MutexLock l(mutex_); return size_map_.at(ptr); } std::optional GpuCudaMallocAsyncAllocator::GetStats() { - if (!stats_) return std::nullopt; + if (!stats_) { + return std::nullopt; + } absl::MutexLock l(mutex_); return *stats_; } bool GpuCudaMallocAsyncAllocator::ClearStats() { - if (!stats_) return false; + if (!stats_) { + return false; + } absl::MutexLock l(mutex_); stats_->num_allocs = 0; stats_->peak_bytes_in_use = stats_->bytes_in_use; diff --git a/third_party/xla/xla/stream_executor/mock_command_buffer.h b/third_party/xla/xla/stream_executor/mock_command_buffer.h index 3dfe4fe7d35d53..15aae998c0c2c2 100644 --- a/third_party/xla/xla/stream_executor/mock_command_buffer.h +++ b/third_party/xla/xla/stream_executor/mock_command_buffer.h @@ -144,6 +144,10 @@ class MockCommandBuffer : public CommandBuffer { (const Command* command, DeviceAddress pred, UpdateCommands update_cond, UpdateCommands update_body), (override)); + MOCK_METHOD(absl::StatusOr, CreateHost, + (absl::AnyInvocable callback, + absl::Span dependencies), + (override)); MOCK_METHOD(absl::Status, SetPriority, (StreamPriority priority), (override)); MOCK_METHOD(absl::Status, Submit, (Stream * stream), (override)); MOCK_METHOD(absl::Status, Finalize, (), (override)); diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_command_buffer.h b/third_party/xla/xla/stream_executor/rocm/rocm_command_buffer.h index 876be379f02bfd..b8e3ac54e6016a 100644 --- a/third_party/xla/xla/stream_executor/rocm/rocm_command_buffer.h +++ b/third_party/xla/xla/stream_executor/rocm/rocm_command_buffer.h @@ -147,6 +147,12 @@ class RocmCommandBuffer : public GpuCommandBuffer { absl::Status UpdateClonedChildNode(GraphNodeHandle node_handle, const CommandBuffer& nested) override; + absl::StatusOr CreateHostNode( + absl::Span dependencies, + absl::AnyInvocable callback) override { + return absl::UnimplementedError("Not implemented."); + } + absl::StatusOr CreateKernelNode( absl::Span dependencies, StreamPriority priority, const ThreadDim& threads, const BlockDim& blocks, diff --git a/third_party/xla/xla/tests/BUILD b/third_party/xla/xla/tests/BUILD index 5b4d1905423859..1e533f63b889af 100644 --- a/third_party/xla/xla/tests/BUILD +++ b/third_party/xla/xla/tests/BUILD @@ -905,9 +905,12 @@ xla_test( ":hlo_pjrt_interpreter_reference_mixin", ":hlo_pjrt_test_base", ":xla_internal_test_main", # fixdeps: keep + ":xla_test_backend_predicates", "//xla:error_spec", + "//xla:fp_util", "//xla:literal", "//xla:literal_util", + "//xla:types", "//xla:xla_data_proto_cc", "//xla/hlo/builder:xla_builder", "//xla/hlo/builder/lib:math", @@ -2759,12 +2762,15 @@ xla_test( ":client_library_test_runner_mixin", ":hlo_pjrt_interpreter_reference_mixin", ":hlo_pjrt_test_base", + ":hlo_runner_agnostic_reference_mixin", ":xla_internal_test_main", "//xla:error_spec", "//xla:shape_util", "//xla:xla_data_proto_cc", "//xla/hlo/builder:xla_builder", + "//xla/hlo/evaluator:hlo_evaluator", "//xla/pjrt/interpreter:interpreter_client", + "//xla/service:hlo_runner_pjrt", "//xla/tsl/platform:test", "@com_google_absl//absl/strings:string_view", "@tsl//tsl/platform:ml_dtypes", diff --git a/third_party/xla/xla/tests/array_elementwise_ops_test.cc b/third_party/xla/xla/tests/array_elementwise_ops_test.cc index 9f0a17faf3c19e..6166b63b00573c 100644 --- a/third_party/xla/xla/tests/array_elementwise_ops_test.cc +++ b/third_party/xla/xla/tests/array_elementwise_ops_test.cc @@ -112,9 +112,6 @@ class ArrayElementwiseOpTest : public ClientLibraryTestRunnerMixin< static constexpr double kEpsF64 = std::numeric_limits::epsilon(); ErrorSpec error_spec_{60 * kEpsF32, 60 * kEpsF32}; ErrorSpec strict_error_spec_{100 * kEpsF64, 100 * kEpsF64}; - - template - void TestComplexTanhUlps(int64_t max_ulps = 2); }; class ArrayElementwiseOpTestParamCount @@ -2953,91 +2950,6 @@ TEST_F(ArrayElementwiseOpTest, TanhF64sVector) { ComputeAndCompare(&builder, {}, strict_error_spec_); } -template -std::vector> GetComplexTanhTestInputs() { - return { - // Small inputs where (exp(a))^2 - (exp(-a))^2 suffered precision loss - // due to exp2 cancellation on older TPUs: - {T(0.0017180424), T(0.0017180424)}, - {T(-0.0017180424), T(0.0017180424)}, - {T(0.0017180424), T(-0.0017180424)}, - {T(-0.0017180424), T(-0.0017180424)}, - - // Other small magnitudes: - {T(1e-5), T(1e-5)}, - {T(-1e-5), T(1e-5)}, - {T(1e-4), T(1e-4)}, - {T(-1e-4), T(-1e-4)}, - {T(1e-3), T(1e-3)}, - {T(-1e-3), T(-1e-3)}, - {T(1e-2), T(1e-2)}, - {T(0.05), T(0.05)}, - - // Points near or on axes: - {T(0.0), T(0.0)}, - {T(1e-3), T(0.0)}, - {T(0.0), T(1e-3)}, - {T(-1e-3), T(0.0)}, - {T(0.0), T(-1e-3)}, - - // Moderate inputs: - {T(0.5), T(0.5)}, - {T(-0.5), T(0.5)}, - {T(1.0), T(1.0)}, - {T(2.0), T(-1.5)}, - - // Large inputs (overflow handling, Re(z) > 15 region where tanh(z) -> +/- - // 1): - {T(15.0), T(0.5)}, - {T(-15.0), T(0.5)}, - {T(20.0), T(1.0)}, - {T(-20.0), T(1.0)}, - }; -} - -template -void ArrayElementwiseOpTest::TestComplexTanhUlps(int64_t max_ulps) { - using RealT = typename ComplexT::value_type; - std::vector xs = GetComplexTanhTestInputs(); - XlaBuilder builder(TestName()); - auto a = ConstantR1(&builder, xs); - Tanh(a); - ASSERT_OK_AND_ASSIGN(Literal actual, this->ExecuteAndTransfer(&builder, {})); - for (int64_t i = 0; i < xs.size(); ++i) { - ComplexT act = actual.Get({i}); - std::complex zd(static_cast(xs[i].real()), - static_cast(xs[i].imag())); - std::complex ref = std::tanh(zd); - ComplexT exp(static_cast(ref.real()), - static_cast(ref.imag())); - - auto real_ulps = UlpDistance(act.real(), exp.real()); - auto imag_ulps = UlpDistance(act.imag(), exp.imag()); - ASSERT_TRUE(real_ulps.has_value()) - << "NaN/Inf mismatch on real part for input " << xs[i]; - ASSERT_TRUE(imag_ulps.has_value()) - << "NaN/Inf mismatch on imag part for input " << xs[i]; - EXPECT_LE(*real_ulps, max_ulps) - << "Real part ULP error exceeded for input " << xs[i] - << ": actual=" << act.real() << ", expected=" << exp.real() - << ", ulp_distance=" << *real_ulps; - EXPECT_LE(*imag_ulps, max_ulps) - << "Imag part ULP error exceeded for input " << xs[i] - << ": actual=" << act.imag() << ", expected=" << exp.imag() - << ", ulp_distance=" << *imag_ulps; - } -} - -TEST_F(ArrayElementwiseOpTest, TanhC64s) { TestComplexTanhUlps(); } - -TEST_F(ArrayElementwiseOpTest, TanhC128s) { - // Float64 transcendentals on TPU use software emulation, which has an error - // bound of up to 100 ULPs (matching strict_error_spec_). On platforms with - // native float64 units (CPU, GPU), enforce <= 2 ULPs. - const int64_t max_ulps = test::DeviceTypeIs(test::kTpu) ? 100 : 2; - TestComplexTanhUlps(max_ulps); -} - TEST_F(ArrayElementwiseOpTest, ExpF32sVector) { // The input tensor is large enough to exercise the vectorized exp // implementation on XLA CPU. diff --git a/third_party/xla/xla/tests/bitcast_convert_test.cc b/third_party/xla/xla/tests/bitcast_convert_test.cc index 1899c086204d48..49ca31d1fc1328 100644 --- a/third_party/xla/xla/tests/bitcast_convert_test.cc +++ b/third_party/xla/xla/tests/bitcast_convert_test.cc @@ -23,11 +23,14 @@ limitations under the License. #include "absl/strings/string_view.h" #include "xla/error_spec.h" #include "xla/hlo/builder/xla_builder.h" +#include "xla/hlo/evaluator/hlo_evaluator.h" #include "xla/pjrt/interpreter/interpreter_client.h" +#include "xla/service/hlo_runner_pjrt.h" #include "xla/shape_util.h" #include "xla/tests/client_library_test_runner_mixin.h" #include "xla/tests/hlo_pjrt_interpreter_reference_mixin.h" #include "xla/tests/hlo_pjrt_test_base.h" +#include "xla/tests/hlo_runner_agnostic_reference_mixin.h" #include "xla/tsl/platform/test.h" #include "xla/xla_data.pb.h" #include "tsl/platform/ml_dtypes.h" @@ -213,7 +216,23 @@ ENTRY main { EXPECT_TRUE(RunAndCompare(hlo_string, ErrorSpec{1e-5, 1e-5})); } -TEST_F(BitcastConvertTest, S8ToPred) { +template +class HloPjRtInterpreterReferenceMixinNoAot + : public HloRunnerAgnosticReferenceMixin { + protected: + template + explicit HloPjRtInterpreterReferenceMixinNoAot(BaseArgs&&... base_args) + : HloRunnerAgnosticReferenceMixin( + std::make_unique(std::make_unique( + []() { return std::make_unique(); })), + std::forward(base_args)...) {} + ~HloPjRtInterpreterReferenceMixinNoAot() override = default; +}; + +class BitcastConvertNoAotTest + : public HloPjRtInterpreterReferenceMixinNoAot {}; + +TEST_F(BitcastConvertNoAotTest, S8ToPred) { absl::string_view hlo_string = R"( HloModule bitcast_to_smaller diff --git a/third_party/xla/xla/tests/complex_unary_op_test.cc b/third_party/xla/xla/tests/complex_unary_op_test.cc index 5939664727dd4e..cbe279a99dd5a1 100644 --- a/third_party/xla/xla/tests/complex_unary_op_test.cc +++ b/third_party/xla/xla/tests/complex_unary_op_test.cc @@ -13,12 +13,17 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include #include +#include #include +#include #include +#include #include #include "xla/error_spec.h" +#include "xla/fp_util.h" #include "xla/hlo/builder/lib/math.h" #include "xla/hlo/builder/xla_builder.h" #include "xla/literal.h" @@ -27,7 +32,9 @@ limitations under the License. #include "xla/tests/complex_unary_op_samples.h" #include "xla/tests/hlo_pjrt_interpreter_reference_mixin.h" #include "xla/tests/hlo_pjrt_test_base.h" +#include "xla/tests/xla_test_backend_predicates.h" #include "xla/tsl/platform/test.h" +#include "xla/types.h" #include "xla/xla_data.pb.h" namespace xla { @@ -106,6 +113,9 @@ class ComplexUnaryOpTest : public ClientLibraryTestRunnerMixin< ComputeAndCompareLiteral(&builder, expected, {}, ErrorSpec(atol)); } } + + template + void TestComplexTanhUlps(); }; TEST_F(ComplexUnaryOpTest, Log1pTest) { @@ -143,5 +153,165 @@ TEST_F(ComplexUnaryOpTest, ExpTest) { [](XlaOp x) { return Exp(x); }); } +template +struct ComplexTanhTestCase { + std::complex input; + // Maximum allowed ULP error on CPU and GPU. + // TODO(phawkins): Tighten near-pole cases to 2 ULPs once MLIR's + // ComplexToStandard pass is fixed. + // tanh(a + bi) = sinh(a + bi) / cosh(a + bi). Multiplying numerator and + // denominator by 2*cosh(a - bi) gives denominator + // 2*cosh(a + bi)*cosh(a - bi) = cosh(2a) + cos(2b), using the identity + // 2*cosh(u)*cosh(v) = cosh(u + v) + cosh(u - v) and cosh(2bi) = cos(2b). + // Using cos(2b) = 2*cos^2(b) - 1, + // cosh(2a) + cos(2b) = (cosh(2a) - 1) + 2*cos^2(b). + // Doubling numerator and denominator (so the real numerator is + // 2*sinh(2a) = expm1(2a) - expm1(-2a)) gives denominator: + // 2*(cosh(2a) - 1) + 4*cos^2(b) + // which avoids cancellation near poles (where cos(2b) ≈ -1). + // However, evaluating 2*(cosh(2a) - 1) as expm1(2a) + expm1(-2a) cancels + // as a -> 0. Using (e^(2a) - 1)*(e^(-2a) - 1) = 2 - 2*cosh(2a), that term + // can be computed via multiplication without cancellation: + // denom = -expm1(2a) * expm1(-2a) + 4 * cos(b)^2. + int64_t cpu_gpu_max_ulps = 2; + // Maximum allowed ULP error across all TPU generations (including older TPUs + // v2-v5p). + // TODO(phawkins): Lower the default bound to 4 (or 2) on newer TPUs (TPU v6e, + // TPU 7x, and later). + int64_t tpu_max_ulps = 30; +}; + +template +std::vector> GetComplexTanhTestInputs() { + return { + // Small inputs where (exp(a))^2 - (exp(-a))^2 suffered precision loss + // due to exp2 cancellation on older TPUs: + {{T(0.0017180424), T(0.0017180424)}}, + {{T(-0.0017180424), T(0.0017180424)}}, + {{T(0.0017180424), T(-0.0017180424)}}, + {{T(-0.0017180424), T(-0.0017180424)}}, + + // Other small magnitudes: + {{T(1e-5), T(1e-5)}}, + {{T(-1e-5), T(1e-5)}}, + {{T(1e-4), T(1e-4)}}, + {{T(-1e-4), T(-1e-4)}}, + {{T(1e-3), T(1e-3)}}, + {{T(-1e-3), T(-1e-3)}}, + // (0.01, 0.01): Older TPUs (v2-v5p) experience up to 602 ULPs due to + // expm1 approximation error. TPU v6+ achieves <= 2 ULPs. + {{T(1e-2), T(1e-2)}, /*cpu_gpu=*/2, /*tpu=*/650}, + // (0.05, 0.05): Older TPUs achieve <= 80 ULPs (72 observed). TPU v6+ + // achieves <= 2 ULPs. + {{T(0.05), T(0.05)}, /*cpu_gpu=*/2, /*tpu=*/80}, + + // Points near or on axes: + {{T(0.0), T(0.0)}}, + {{T(1e-3), T(0.0)}}, + {{T(0.0), T(1e-3)}}, + {{T(-1e-3), T(0.0)}}, + {{T(0.0), T(-1e-3)}}, + + // Moderate inputs: + {{T(0.5), T(0.5)}}, + {{T(-0.5), T(0.5)}}, + {{T(1.0), T(1.0)}}, + {{T(2.0), T(-1.5)}}, + + // Inputs near poles (b ≈ ±pi/2) where cos(b) ≈ 0 and + // expm1(2a) + expm1(-2a) cancellation causes up to 131 ULPs on CPU/GPU. + // TODO(phawkins): Tighten cpu_gpu bounds to 2 ULPs once MLIR's + // ComplexToStandard pass is fixed. The fix is to compute the + // denominator's + // exponential term as (expm1(a) - expm1(-a))^2 instead of + // expm1(2a) + expm1(-2a), avoiding catastrophic cancellation when a -> 0. + // On older TPUs (v2-v5p), 0.002 achieves <= 22 ULPs, 0.005 achieves <= + // 125 ULPs. + // On TPU v6+, all achieve <= 4 ULPs. + {{T(0.002), T(1.5707963267948966)}, /*cpu_gpu=*/150}, + {{T(-0.002), T(1.5707963267948966)}, /*cpu_gpu=*/150}, + {{T(0.002), T(-1.5707963267948966)}, /*cpu_gpu=*/150}, + {{T(0.005), T(1.5707963267948966)}, /*cpu_gpu=*/150, /*tpu=*/150}, + // Input near pole (b ≈ pi/2): + // Older TPUs achieve <= 130 ULPs (111 observed); TPU v6+ achieves <= 2 + // ULPs; H100 and B200 achieve <= 3 ULPs. + {{T(0.11656774), T(1.7330385)}, /*cpu_gpu=*/4, /*tpu=*/130}, + + // Large inputs (overflow handling, Re(z) > 15 region where tanh(z) -> +/- + // 1): + // Older TPUs achieve <= 80 ULPs (74 observed); TPU v6+ achieves <= 16 + // ULPs. + {{T(15.0), T(0.5)}, /*cpu_gpu=*/2, /*tpu=*/80}, + {{T(-15.0), T(0.5)}, /*cpu_gpu=*/2, /*tpu=*/80}, + {{T(20.0), T(1.0)}}, + {{T(-20.0), T(1.0)}}, + }; +} + +template +void ComplexUnaryOpTest::TestComplexTanhUlps() { + using RealT = typename ComplexT::value_type; + std::vector> cases = + GetComplexTanhTestInputs(); + std::vector xs; + xs.reserve(cases.size()); + for (const auto& c : cases) { + xs.push_back(c.input); + } + XlaBuilder builder(TestName()); + Literal input_literal = LiteralUtil::CreateR1(xs); + auto a = Parameter(&builder, 0, input_literal.shape(), "a"); + Tanh(a); + std::vector args = {&input_literal}; + ASSERT_OK_AND_ASSIGN(Literal actual, + this->ExecuteAndTransfer(&builder, args)); + const bool is_tpu = test::DeviceTypeIs(test::kTpu); + for (int64_t i = 0; i < cases.size(); ++i) { + ComplexT act = actual.Get({i}); + std::complex zd(static_cast(xs[i].real()), + static_cast(xs[i].imag())); + std::complex ref = std::tanh(zd); + ComplexT exp(static_cast(ref.real()), + static_cast(ref.imag())); + + auto real_ulps = UlpDistance(act.real(), exp.real()); + auto imag_ulps = UlpDistance(act.imag(), exp.imag()); + ASSERT_TRUE(real_ulps.has_value()) + << "NaN/Inf mismatch on real part for input " << xs[i]; + ASSERT_TRUE(imag_ulps.has_value()) + << "NaN/Inf mismatch on imag part for input " << xs[i]; + + int64_t max_ulps; + if constexpr (std::is_same_v) { + // Host libc++ std::complex::tanh suffers from catastrophic cancellation + // near poles (b ≈ ±pi/2). For complex64, computing the reference in + // double avoids this, but for complex128 there is no higher precision + // type to use for the reference, leading to large observed ULP errors. + max_ulps = 100000; + } else { + max_ulps = is_tpu ? cases[i].tpu_max_ulps : cases[i].cpu_gpu_max_ulps; + } + + if (test::DeviceTypeIs(test::kInterpreter)) { + // Similarly, the interpreter evaluates tanh via host libc++ + // std::complex, suffering from the same cancellation near poles. + max_ulps = std::max(max_ulps, int64_t{30000}); + } + + EXPECT_LE(*real_ulps, max_ulps) + << "Real part ULP error exceeded for input " << xs[i] + << ": actual=" << act.real() << ", expected=" << exp.real() + << ", ulp_distance=" << *real_ulps; + EXPECT_LE(*imag_ulps, max_ulps) + << "Imag part ULP error exceeded for input " << xs[i] + << ": actual=" << act.imag() << ", expected=" << exp.imag() + << ", ulp_distance=" << *imag_ulps; + } +} + +TEST_F(ComplexUnaryOpTest, TanhC64s) { TestComplexTanhUlps(); } + +TEST_F(ComplexUnaryOpTest, TanhC128s) { TestComplexTanhUlps(); } + } // namespace } // namespace xla diff --git a/third_party/xla/xla/tests/constraint_propagator.cc b/third_party/xla/xla/tests/constraint_propagator.cc index 37cf1f89250ad2..ac980eb7a61ec1 100644 --- a/third_party/xla/xla/tests/constraint_propagator.cc +++ b/third_party/xla/xla/tests/constraint_propagator.cc @@ -1291,8 +1291,19 @@ void ConstraintPropagator::PropagateReduceApprox( int64_t num_inputs = instruction->operand_count() / 2; int64_t num_elements = 1; const Shape& operand_shape = instruction->operand(0)->shape(); - for (int64_t dim : instruction->dimensions()) { - num_elements *= operand_shape.dimensions(dim); + if (instruction->opcode() == HloOpcode::kReduce) { + for (int64_t dim : instruction->dimensions()) { + num_elements *= operand_shape.dimensions(dim); + } + } else if (instruction->opcode() == HloOpcode::kReduceWindow) { + const Window& window = instruction->window(); + for (int64_t d = 0; + d < window.dimensions_size() && d < operand_shape.dimensions_size(); + ++d) { + int64_t win_size = window.dimensions(d).size(); + int64_t op_dim = operand_shape.dimensions(d); + num_elements *= std::max(1, std::min(win_size, op_dim)); + } } if (num_elements > 1) { std::optional root_op = @@ -1633,8 +1644,11 @@ absl::Status ConstraintPropagator::PropagateConstraintsApprox( const HloInstruction* instruction) { ConstraintInterval output_interval = states_[instruction].GetConstraintInterval(); + // Exempt reductions: tuple-shaped outputs have no top-level interval and + // resolve bounds from downstream GTE users. if ((output_interval.IsEmpty() || output_interval.IsUnconstrained()) && - instruction->opcode() != HloOpcode::kReduce) { + instruction->opcode() != HloOpcode::kReduce && + instruction->opcode() != HloOpcode::kReduceWindow) { return absl::OkStatus(); } switch (instruction->opcode()) { @@ -1648,6 +1662,7 @@ absl::Status ConstraintPropagator::PropagateConstraintsApprox( PropagateMultiplyApprox(instruction, output_interval); break; case HloOpcode::kReduce: + case HloOpcode::kReduceWindow: PropagateReduceApprox(instruction, output_interval); break; case HloOpcode::kConvolution: diff --git a/third_party/xla/xla/tests/constraint_propagator_test.cc b/third_party/xla/xla/tests/constraint_propagator_test.cc index 651fdc2ecb8a0e..47d0b515bc739e 100644 --- a/third_party/xla/xla/tests/constraint_propagator_test.cc +++ b/third_party/xla/xla/tests/constraint_propagator_test.cc @@ -567,6 +567,38 @@ ENTRY main { EXPECT_NEAR(p0_int.min, -expected_max_in, 1e-3); } +TEST_F(ConstraintPropagatorTest, ReduceWindowSumReverseConstraintShapes) { + const char* hlo = R"( +HloModule TestModule +add_computation { + x = bf16[] parameter(0) + y = bf16[] parameter(1) + ROOT add = bf16[] add(x, y) +} +ENTRY main { + param_0 = bf16[16,16,256,256] parameter(0) + init = bf16[] constant(0) + ROOT reduce_window = bf16[16,16,256,256] reduce-window(param_0, init), + window={size=1x1x1x511 pad=0_0x0_0x0_0x255_255}, + to_apply=add_computation +} +)"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo)); + ASSERT_OK_AND_ASSIGN(auto states, ConstraintPropagator::Run(*module)); + + auto p0_int = states[module->entry_computation()->parameter_instruction(0)] + .GetConstraintInterval(); + + // Window size is 511 along dim 3, but operand dim 3 size is 256. + // Effective reduction size N = min(511, 256) = 256 elements. + // Root max_out = 65504.0. + // max_in = 65504.0 / 256 = 255.875. + double expected_max_in = 65504.0 / 256.0; + EXPECT_FALSE(p0_int.IsEmpty()); + EXPECT_NEAR(p0_int.max, expected_max_in, 1e-3); + EXPECT_NEAR(p0_int.min, -expected_max_in, 1e-3); +} + TEST_F(ConstraintPropagatorTest, DotReverseConstraintShapes) { const char* hlo = R"( HloModule TestModule diff --git a/third_party/xla/xla/tests/test_utils.cc b/third_party/xla/xla/tests/test_utils.cc index be58d91c2de931..a34c66b735cb38 100644 --- a/third_party/xla/xla/tests/test_utils.cc +++ b/third_party/xla/xla/tests/test_utils.cc @@ -111,11 +111,14 @@ void FindConstrainedUsesHelper( op_num == 1) { constrained_uses.push_back(use); } else if (opcode == HloOpcode::kFusion) { - const HloInstruction* const to_analyze = - user->fused_parameter(op_num); - FindConstrainedUsesHelper(dataflow, *to_analyze, - treat_gte_as_data_formatting, visited, - constrained_uses); + // Fusions can have unused operands with no corresponding parameter. + if (op_num < user->fused_parameters().size()) { + const HloInstruction* const to_analyze = + user->fused_parameter(op_num); + FindConstrainedUsesHelper(dataflow, *to_analyze, + treat_gte_as_data_formatting, visited, + constrained_uses); + } } else if (NeedsInitValue(use)) { constrained_uses.push_back(use); } else if (IsDataFormattingOp(user, treat_gte_as_data_formatting)) { diff --git a/third_party/xla/xla/tests/test_utils_test.cc b/third_party/xla/xla/tests/test_utils_test.cc index 1308d2a4db91eb..90192dcb520d96 100644 --- a/third_party/xla/xla/tests/test_utils_test.cc +++ b/third_party/xla/xla/tests/test_utils_test.cc @@ -683,5 +683,29 @@ ENTRY entry { EXPECT_EQ(args[1].Get({}), LiteralUtil::MaxValue(F32).Get({})); } +// Tests that MakeFakeArguments succeeds when a fusion instruction has unused +// operands (operands beyond the number of parameters in the fused computation). +TEST_F(TestUtilsTest, FusionWithUnusedOperand) { + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"( +HloModule FusionWithUnusedOperandModule + +fused_computation (param_0: f32[4,5,128,256]) -> f32[4,5,128,256] { + param_0 = f32[4,5,128,256] parameter(0) + ROOT copy = f32[4,5,128,256] copy(param_0) +} + +ENTRY entry { + param_0 = f32[4,5,128,256] parameter(0) + param_1 = f32[] parameter(1) + copy = f32[] copy(param_1) + ROOT fusion = f32[4,5,128,256] fusion(param_0, copy), kind=kOutput, calls=fused_computation +} +)")); + + ASSERT_OK_AND_ASSIGN(std::vector args, + MakeFakeArguments(module.get())); + EXPECT_EQ(args.size(), 2); +} + } // namespace } // namespace xla diff --git a/third_party/xla/xla/tsl/profiler/utils/xplane_schema.cc b/third_party/xla/xla/tsl/profiler/utils/xplane_schema.cc index a36ff25dc0d6e9..df0fc9f2e50f89 100644 --- a/third_party/xla/xla/tsl/profiler/utils/xplane_schema.cc +++ b/third_party/xla/xla/tsl/profiler/utils/xplane_schema.cc @@ -353,6 +353,12 @@ const StatTypeMap& GetStatTypeMap() { kDevCapPeakSramRdBwGigabytesPerSecond}, {"peak_sram_wr_bw_gigabytes_per_second", kDevCapPeakSramWrBwGigabytesPerSecond}, + {"peak_spmem_rd_bw_gigabytes_per_second", + kDevCapPeakSpmemRdBwGigabytesPerSecond}, + {"peak_spmem_wr_bw_gigabytes_per_second", + kDevCapPeakSpmemWrBwGigabytesPerSecond}, + {"peak_sc_teraflops_per_second", kDevCapPeakScTeraflopsPerSecond}, + {"num_sparse_core_tiles", kDevCapNumSparseCoreTiles}, {"device_vendor", kDevVendor}, {"has_megacore", kDevHasMegacore}, {"has_merged_vmem", kDevHasMergedVmem}, diff --git a/third_party/xla/xla/tsl/profiler/utils/xplane_schema.h b/third_party/xla/xla/tsl/profiler/utils/xplane_schema.h index 3f776c23721ae0..a532605a80049d 100644 --- a/third_party/xla/xla/tsl/profiler/utils/xplane_schema.h +++ b/third_party/xla/xla/tsl/profiler/utils/xplane_schema.h @@ -410,11 +410,15 @@ enum StatType { kUsesIci, kDimensions, kType, + kDevCapPeakSpmemRdBwGigabytesPerSecond, + kDevCapPeakSpmemWrBwGigabytesPerSecond, + kDevCapPeakScTeraflopsPerSecond, + kDevCapNumSparseCoreTiles, // LINT.ThenChange(:last_stat_type) // LINT.IfChange(last_stat_type) // Change this to point to the last stat type when adding a new one. - kLastStatType = kType, + kLastStatType = kDevCapNumSparseCoreTiles, // LINT.ThenChange(:stat_type_enum) }; diff --git a/third_party/xla/xla/xla.proto b/third_party/xla/xla/xla.proto index e8e40de0502825..7a8b0f552f3896 100644 --- a/third_party/xla/xla/xla.proto +++ b/third_party/xla/xla/xla.proto @@ -1837,7 +1837,15 @@ message DebugOptions { // kept inline. Default is MAX_INT (deduplication disabled). optional int64 xla_deduplicate_backend_configs_min_size = 542; - // Next id: 543 + // Single serialized config to override config of all instructions, bypassing + // cache and autotuning. Accepts textproto of autotuner.Config. + optional string xla_force_config = 543; + + // File containing a list of serialized configs to override candidate configs + // for all instructions. Accepts textproto of autotuner.CandidateConfigs. + optional string xla_candidate_configs_file = 544; + + // Next id: 545 // Extra options to pass to the compilation backend (e.g. LLVM); specific // interpretation of these values is left to the backend.